BERTを量子化して高速かつ軽量にする
こんにちは、@vimmodeです。自然言語界隈ではBERTを始めとしたTransformerベースの手法の進化が目覚ましいですが、実運用されている話はあまり聞きません。
その理由としてモデルのサイズの大きさと推論速度の遅さに一定起因すると感じており、この記事はその解消になり得る量子化と呼ばれる手法の紹介とPyTorchで実装されたBERTモデルに量子化を適応する方法を紹介します。
量子化とは
量子化という単語は数学や物理など様々な領域で使われています。ここで述べる量子化は情報理論における量子化であり、主に連続値を離散値で表現することを考えます。
機械学習の枠組みで考えるとモデルのパラメータや学習時の勾配（場合によっては入力と出力データも含める）の数値表現を浮動小数点から整数に変更することを目的にします。
ディープラーニングではパラメータ等をfloat32で表現することが多いですが、もしこれをint8に置き換えることができればそれだけでサイズを1/4に減らすことが出来ます。計算速度はモデルの構造に強く依存しますが、シンプルな線型結合を多用するBERTは恩恵を強く受けます。
量子化により数値の小数点精度が落ちるため、モデルの精度が低下することは想像出来ます。如何に精度を落とすことなく、軽量化することが出来るかがこの分野のチャレンジと言えます。
モデルの軽量化と同じ文脈で蒸留と呼ばれる、モデル自体を圧縮する手法もありますが量子化とは異なる考え方を持ちますので、量子化とは異なるものとして考えて良いと思います。
また、float32の精度にほぼ近似しつつもfloat16で学習と推論を可能にする混合精度演算と呼ばれる技術もありますが、これも量子化と仕組みが異なりますので、これも別物と考えて良いと思います。
量子化の要素整理
一概に量子化と言っても、どの要素を何で量子化するかによって必要な処理が変わるので、ニューラルネットワークで考慮する要素を以下のように整理しました。
量子化の対象
- 重み（パラメータ）
- 勾配
- 活性化情報
量子化を実行するタイミング
量子化のレベル
- もとの値を任意のビットに変更
- 1ビットの場合はバイナリーニューラルネットワークと呼ばれる
- PyTorchではint8(tensor.uint8)がサポートされています
これらの要素を考慮しつつ精度を落とさない工夫はグーグルの研究者によるこちらの論文で解説されています。どれも専門性の高い話なのでここでの説明は割愛いたします。
BERTを量子化する
以降PyTorchで実装されたBERTモデルを量子化します。今回は学習済みのBERTモデルに対して量子化を適用し、適用前後のモデルのサイズと推論速度、精度を比較します。
PyTorch公式チュートリアルに詳細な手順が記載されていますので、ここでは簡略化したコードと補助的な情報をまとめます。
pytorch.org
実験は以下の内容で行います。
- タスク: テキスト分類
- データセット: Microsoft Research Paraphrase Corpusが公開したコーパス
- BERT実装: transformers 実装
前準備として学習済みBERTにファインチューニングを行います。学習コードはこちらであり、以下のように引数を与えて学習を行います。
export GLUE_DIR=./glue_data
export TASK_NAME=MRPC
export OUT_DIR=./$TASK_NAME/
python ./run_glue.py \
--model_type bert \
--model_name_or_path bert-base-uncased \
--task_name $TASK_NAME \
--do_train \
--do_eval \
--do_lower_case \
--data_dir $GLUE_DIR/$TASK_NAME \
--max_seq_length 128 \
--per_gpu_eval_batch_size=8 \
--per_gpu_train_batch_size=8 \
--learning_rate 2e-5 \
--num_train_epochs 3.0 \
--save_steps 100000 \
--output_dir $OUT_DIR
学習後にPythonから学習済みモデルをロードします
import torch from transformers import BertConfig, BertForSequenceClassification, BertTokenizer tokenizer = BertTokenizer.from_pretrained(MODEL_DIR) model = BertForSequenceClassification.from_pretrained(MODEL_DIR)
次にBERTを量子化します。 PyTorchの量子化変換としてtorch.quantation.quantize_dynamic()関数が提供されており、これに定義済みのモデルを第一引数に、第二引数としてモデル内で量子化するレイヤーを指定します（省略する場合は全レイヤーに対して量子化が適用されます）。
BERTの内部では線形結合層(torch.nn.Linear)が多く使用します。今回はここを指定して量子化を実行します。
quantized_model = torch.quantization.quantize_dynamic(
model, {torch.nn.Linear}, dtype=torch.qint8
)
この時点で量子化前と後のモデルのサイズを比較すると以下のようになります。
def print_size_of_model(model): torch.save(model.state_dict(), "temp.p") print('Size (MB):', os.path.getsize("temp.p")/1e6) os.remove('temp.p') print_size_of_model(model) print_size_of_model(quantized_model) >>> Size (MB): 437.980564 Size (MB): 181.440471
モデルサイズが約40%に圧縮できました。
次に推論速度(1スレッド及び4スレッド）と精度を確認します。
def time_model_evaluation(model, configs, tokenizer): eval_start_time = time.time() result = evaluate(configs, model, tokenizer, prefix="") eval_end_time = time.time() eval_duration_time = eval_end_time - eval_start_time print(result) print("Evaluate total time (seconds): {0:.1f}".format(eval_duration_time)) # Evaluate the original FP32 BERT model time_model_evaluation(model, configs, tokenizer) # Evaluate the INT8 BERT model after the dynamic quantization time_model_evaluation(quantized_model, configs, tokenizer) >>> | Prec | F1 score | Model Size | 1 thread | 4 threads | | FP32 | 0.9019 | 438 MB | 160 sec | 85 sec | | INT8 | 0.8953 | 181 MB | 90 sec | 46 sec |
この段階でF1スコアで0.005程度の悪化にとどまる一方で、サイズは40%に圧縮され、推論速度は4スレッドで約50%に減少させることが出来ました。 Kaggleのようなコンペションではこの精度の悪化は課題となりますが、実務で速度パフォーマンスを重視する場合は十分検討価値があると思います。
今回は学習済みモデルに対して量子化を適用させましたが、学習時から量子化を適用することで1段階float32のモデルの精度に近づけることが出来ます。それはまた他の機会に紹介できればと思います。
