深層学習の量子化に入門してみた〜BERTをIntel Neural CompressorでStatic Quantization〜

こんにちは。リサーチャーの勝又です。私はレトリバで自然言語処理、とくに要約や文法誤り訂正に関する研究の最新動向の調査・キャッチアップなどを行っております。

前々回、Dynamic Quantizationを様々なライブラリで試した記事を公開しました。前回はStatic QuantizationをONNX Runtimeで試した記事を公開しました。今回はStatic QuantizationをIntel Neural Compressorで試してみようと思います。

Static Quantizationの概要
ONNX Runtimeを利用した量子化方法紹介
量子化実験
まとめ

Static Quantizationの概要

Static Quantization（Post Training Quantization）は、量子化手法の1つで、入力 $x$ に対して量子化のscale factor $s_{x}$ を事前に決定する手法です。詳細については以前の記事を確認していただければと思います。

前回はONNX Runtimeライブラリを利用したStatic Quantizationを試しました。その結果、Static QuantizationよりもDynamic Quantizationの方が推論速度が速いことがわかりました。

今回はバックエンドをONNX RuntimeからIntel Neural Compressor（INC）に変えてStatic Quantizationを行い、同様の傾向が得られるのかを簡単に調べてみました。具体的には、自然言語処理でよく使われるBERT（Bidirectional Encoder Representations from Transformers）¹と呼ばれるモデルに対して、Static Quantizationを試しました。

ちなみに、INCはDynamic Quantizationを試す際にも紹介しています。詳しくはこちらの記事をご確認ください。

ONNX Runtimeを利用した量子化方法紹介

今回もTransformersライブラリでfine-tuningを行ったPyTorchモデルを、バッチサイズ1でのCPU推論を行う想定で量子化します。

Dynamic Quantizationの記事やONNX RuntimeでStatic Quantizationを試した記事でご紹介した通り、Transformersの拡張ライブラリoptimumを使用することで、Transformersでfine-tuningしたモデルを簡単に量子化できます。

INCでStatic Quantizationを行うプログラムの流れは以下の通りです。

データの準備および単語分割を行う
Trainerおよび精度の評価を行う処理を作成
量子化のConfigを作成
量子化したいモデルをGraph Moduleに変換
モデルを引数に受け取り、そのモデルの精度の評価結果を返す処理を作成
量子化の実行

データの準備および単語分割を行う

Static Quantizationはscale factor $s_{x}$ を事前に計算するために、実際にfine-tuningで使用した学習データを用意する必要があります。

from datasets import load_dataset
from transformers import AutoTokenizer

train_file_path = "/path/to/train.json"
valid_file_path = "/path/to/valid.json"  # INCはvalidationも必要になります
data_files = {"train": train_file_path, "valid": valid_file_path}
raw_dataset = load_dataset("json", data_files=data_files)

model_name_or_path = "/path/to/fine-tuned-model"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)

# 今回は評価にTrainerを実装を使用するので、Static Quantizationではありますがラベル周りの処理も記述します
label_list = raw_dataset["train"].unique("label")
label_list.sort()
num_labels = len(label_list)
label_to_id = {v: i for i, v in enumerate(label_list)}

def preprocess_function(examples):
  result = tokenizer(examples["sentence1"], padding=True, truncation=True, max_length=512)
  result = [(label_to_id[l] if l != -1 else -1) for l in examples["label"]]
  return result
# 単語がid表現となった学習データを作成
raw_dataset = raw_dataset.map(preprocess_function, batched=True)

calibration_dataset_size = 100  # 学習データの内、100サンプルをcalibrationに使用する
train_dataset = raw_dataset["train"].select(range(calibration_dataset_size))
eval_dataset = raw_dataset["valid"]

ここで作成したtrain_datasetとeval_datasetを用いてモデルを量子化します。

Trainerおよび精度の評価を行う処理を作成

INCは量子化後の精度が量子化前とどの程度離れているかを計算する実装となっているため、精度評価を行う処理を記述する必要があります。せっかくTransformersやOptimumを利用しているので、今回は簡単にTrainerを用いてこの部分を記述しようと思います。

import numpy as np
from optimum.intel.neural_compressor import IncTrainer
from transformers import default_data_collator

def compute_metrics(pred_output):
  preds = pred_output.predictions[0] if isinstance(pred_output.predictions, tuple) else pred_output.predictions
  preds = np.argmax(preds, axis=1)
  return {"accuracy": (preds == pred_output.label_ids).astype(np.float32).mean().item()}

trainer = IncTrainer(
  model=model,
  args=training_args,  # ここでのtraining_argsはhuggingface/transformersのTrainingArgumentsをそのまま使っています
  train_dataset=train_dataset,
  eval_dataset=eval_dataset,
  tokenizer=tokenizer,
  data_collator=default_data_collator,
  compute_metrics=compute_metrics,
)

量子化のConfigを作成

次に量子化に関するConfigを作成します。 Configの内容については公式実装を確認していただければと思います。

from optimum.intel.neural_compressor import IncQuantizationConfig, IncQuantizationMode

# /path/to/config_dir/quantization.ymlの場合、引数は以下の通り
q_config = IncQuantizationConfig.from_pretrained(
  "/path/to/config_dir",
  config_file_name="quantization.yml",
)
quantization_approach = "static"
quant_approach = getattr(IncQuantizationMode, quantization_approach.upper()).value
q_config.set_config("quantization.approach", quant_approach)

上記のq_configが量子化に使用するConfigをまとめたものになります。

量子化したいモデルをGraph Moduleに変換

今回量子化する対象であるBERTをPyTorchのGraph Moduleに変換します。便利なことに、この変換スクリプトもTransformersに実装されているため、今回はそちらを使用します。

from transformers.utils.fx import symbolic_trace

eval_dataloader = trainer.get_eval_dataloader()
it = iter(eval_dataloader)
input_names = next(it).keys()

q_config.set_config("model.framework", "pytorch_fx")
model = symbolic_trace(
  model,
  input_names=input_names,
  batch_size=training_args.per_device_eval_batch_size,
  sequence_length=512,  # 推論に使用したい系列長を入力
)

モデルを引数に受け取り、そのモデルの精度の評価結果を返す処理を作成

前述の通り、INCは量子化時に精度評価を行う処理が必要で、今回はTrainerを用いて対応します。具体的には以下のように、eval_func(model)を作成し、その中で精度評価を行います。

metric_name = "eval_accuracy"

def take_eval_steps(model, trainer, metric_name):
  trainer.model = model
  # modelがeval_datasetに対して推論を行い、その推論結果をcompute_metricsで評価し、スコアを取得します
  metrics = trainer.evaluate()
  return metrics.get(metric_name)

def eval_func(model):
  return take_eval_steps(model, trainer, metric_name)

量子化の実行

最後に、以下のようにして量子化を実行します。

import yaml
from optimum.intel.neural_compressor import IncQuantizer
from optimum.intel.neural_compressor.utils import CONFIG_NAME

# scale factorを求めるデータはtrain_datasetを使用します
calib_dataloader = trainer.get_train_dataloader()
inc_quantizer = IncQuantizer(
  model, q_config, calib_dataloader=calib_dataloader, eval_func=eval_func
)
q_model = inc_quantizer.fit_static()

trainer.save_model(training_args.output_dir)
with open(os.path.join(training_args.output_dir, CONFIG_NAME), "w") as o_f:
  yaml.dump(q_model.tune_cfg, o_f, default_flow_style=False)

上記スクリプトを実行すると、training_args.output_dir/pytorch_model.binに量子化されたモデルが作成され、training_args.output_dir/best_configure.yamlに量子化時の設定が保存されます。

このように、optimumを使うことで、お手軽にPyTorchのBERTを量子化できました。今回、簡単に量子化実験も行っているのですが、量子化にはこのoptimumを使った方法を使用しています²。

量子化実験

それでは、実際に量子化を用いた実験を行おうと思います。

実験設定

今回はJGLUEのMARC-jaとJSTS、JNLI、JCommonsenseQAの合計4つのデータセットに対して量子化実験を行いました。具体的な流れとしては、各データセットについて一度Transformersのrun_glue.pyまたはrun_swag.pyで東北大BERT（cl-tohoku/bert-base-japanese-whole-word-masking）をfine-tuningします³。その後、fine-tuningしたモデルを量子化、その量子化モデルを用いて各データセットの開発データに対して推論を行いました。今回は簡単な実験ということでINCのeval_datasetにも開発データをそのまま使用しています。

推論はCPU上で行い、batch sizeは1で行っています。またMARC-jaとJSTS、JNLIはmax lengthを512に設定し、JCommonsenseQAは64に設定しました。また推論時間⁴は開発データそれぞれの平均としています。

Static Quantizationの実装はoptimumのINC用run_glue.pyとINC用run_swag.pyを元にしました。ただし、以前のDynamic Quantizationの際と同様に、LANG=en_us.utf-8の設定とCUDA_VISIBLE_DEVICES=-1によるCPU環境での実行が必要でした。また、dataloader_drop_lastの設定もTrueにしておく必要があります。

なお、今回のライブラリのバージョンは以下の通りです。

neural-compressor: 1.11
optimum: 1.2.1
pytorch: 1.9.1
transformers: 4.16.2

ちなみに最新バージョンのoptimum（v1.2.3）ではINCに関する実装がhuggigface/optimum-intelに移行しておりますのでご注意ください。

推論速度の比較

前回のONNX Runtimeと同様に、Static QuantizationとDynamic Quantizationでどの程度推論速度が変化するか調査しました。比較対象としては、量子化を行わずに実行した場合（PyTorch fp32）とDynamic Quantizationを利用した場合（Dynamic Quant）を実験しています。結果が以下のグラフとなります。

このグラフからわかる通り、量子化手法（Static、Dynamic）を使うことで推論をより高速にできました。 Static Quantizationは量子化前から推論速度を平均約37.6%ほど減少、Dynamic Quantizationは平均約25.5%ほど減少できました。

また、前回のONNX Runtimeの結果とは異なり、Static Quantizationの方がDynamic Quantizationと比べて推論速度が速い結果となりました。（ONNX RuntimeのDynamic Quantizationが速すぎるだけのような気がしますが...）

推論精度の比較

推論速度と同様に、推論精度に関しても比較してみたいと思います。結果としては以下のグラフとなりました。

どの実験結果についても、大きく精度の差はつきませんでした。一貫した傾向として、どのデータセットについても量子化を行わないモデル（Baseline）がもっともよく、量子化を行うとほんの少しではありますが精度が劣化する結果となりました。データセットによってはStatic Quantizationの方がDynamic Quantizationより精度が高く、その逆の結果もいくつか確認できました。

まとめ

今回の記事では、BERTに対してIntel Neural Compressorを基にしたStatic Quantizationを簡単に動かしてみました。前回のONNX RuntimeでのStatic Quantizationに引き続き、optimumを利用して量子化を行いました。結果としてStatic Quantizationを行うことで、量子化を行わない場合より推論速度を高速化できました。とくにIntel Neural Compressorを用いた場合はDynamic QuantizationよりもStatic Quantizationの方が推論速度が速いこともわかりました。このように、実際に量子化を行うライブラリごとにStatic Quantization、Dynamic Quantizationの向き不向きがあるようです。

次回はQuantization Aware Trainingを検証し、その振る舞いを確認しようと思います。

BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova. [paper]↩
今回、regressionタスクであるJSTSでも実験を行っています。JSTSについてはcompute_metricsの評価をaccuracyからmseにしています。↩
正確にはJGLUEで公開していただいているパッチを当てたスクリプトを使用しました。↩
ここでの推論時間は、モデルロードなどは含んでいません。テストデータをdata_loaderでモデルに入力し、すべてのテストデータが推論されるまでの時間を計測、1サンプルあたりの平均時間を計測しました。↩