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

機械学習・深層学習 自然言語処理

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

前々回、深層学習の量子化について簡単な解説記事を公開しました。 前回はDynamic Quantizationを試した記事を公開しました。 今回はStatic Quantizationを実際に試してみようと思います。

Static Quantizationの概要

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

以前紹介したDynamic Quantizationは動的にscale factor s_{x}を決定するため、その分だけStatic Quantizationの高速化が期待できます。 今回は実際に、自然言語処理でよく使われるBERT（Bidirectional Encoder Representations from Transformers）1と呼ばれるモデルに対して、ONNX Runtimeライブラリを利用してStatic Quantizationを試してみたいと思います。

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

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

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

流れとしては以下の通りです。

  1. データの準備および単語分割を行う
  2. 量子化のConfigを作成
  3. Calibrationを実行してscale factorを計算する
  4. 量子化適用前に除外するノードを定義
  5. 量子化の実行

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

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

from datasets import load_dataset
from transformers import AutoTokenizer

train_file_path = "/path/to/train.json"
data_files = {"train": train_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)

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

量子化のConfigを作成

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

from onnxruntime.quantization import QuantFormat
from optimum.onnxruntime.configuration import QuantizationConfig

qconfig = QuantizationConfig(
  is_static=True,
  format=QuantFormat.QDQ,
  operators_to_quantize=["MatMul", "Add"],
)

Calibrationを実行してscale factorを計算する

fine-tuneしたモデルを用意し、calibrationを実行してscale factorを計算します。

from optimum.onnxruntime import ORTQuantizer
from optimum.onnxruntime.configuration import AutoCalibrationConfig

# fine-tuneしたモデルを用意します
quantizer = ORTQuantizer.from_pretrained(model_name_or_path, feature="sequence-classification")

# 1で用意したデータセットを準備します
calibration_dataset_size = 100  # 学習データの内、100サンプルをcalibrationに使用するとします
calibration_dataset = raw_dataset["train"]
calibration_dataset = calibration_dataset.select(range(calibration_dataset_size))
calibration_dataset = quantizer.clean_calibration_dataset(calibration_dataset)
calibration_config = AutoCalibrationConfig.minmax(calibration_dataset)

# calibrationを実行します
onnx_model_path = "/path/to/model.onnx"  # （量子化してない）onnxモデルの保存先
ranges = quantizer.fit(
  dataset=calibration_dataset,
  calibration_config=calibration_config,
  onnx_model_path=onnx_model_path,
  operators_to_quantize=qconfig.operators_to_quantize,
  batch_size=8,  # calibrationの際のbatch_size
)

最後の出力rangesがcalibrationで得られたscale factorです。

ここで、calibration用のconfigとしてcalibration_configを作成しています。 今回はcalibration方法としてminmaxを選択しましたが、他にもentropyやpercentilesなども存在します。 気になる方はこちらの公式実装をご確認ください。

量子化適用前に除外するノードを定義

calibrationもできたので量子化するぞ、と行きたいのですが、もう1つやらないといけないことがあります。 量子化を行う際に、いくつかの処理（ここではノードと呼称します）は除外しないと精度が下がるといった話があったりする2のでその対策を行います。 今回は公式のexampleの除外処理をそのまま使用しています。

from optimum.onnxruntime.preprocessors import QuantizationPreprocessor
from optimum.onnxruntime.preprocessors.passes import (
  ExcludeGeLUNodes,
  ExcludeLayerNormNodes,
  ExcludeNodeAfter,
  ExcludeNodeFollowedBy,
)

quantization_preprocessor = QuantizationPreprocessor(onnx_model_path)
quantization_preprocessor.register_pass(ExcludeLayerNormNodes())
quantization_preprocessor.register_pass(ExcludeGeLUNodes())
quantization_preprocessor.register_pass(ExcludeNodeAfter("Add", "Add"))  # residual
quantization_preprocessor.register_pass(ExcludeNodeAfter("Gather", "Add"))
quantization_preprocessor.register_pass(ExcludeNodeFollowedBy("Add", "Softmax"))

量子化の実行

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

quantized_model_path = "/path/to/model-quantized.onnx"
quantizer.export(
  onnx_model_path=onnx_model_path,
  onnx_quantized_model_output_path=quantized_model_path,
  calibration_tensors=ranges,  # calibrationの出力
  quantization_config=qconfig,  # 量子化のconfig
  preprocessor=quantization_preprocessor,  # 除外するノード
)

上記スクリプトを実行すると、/path/to/model-quantized.onnxに量子化されたONNXモデルが作成されます。

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

量子化実験

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

実験設定

今回はLivedoorニュースコーパスを利用して、記事が与えられた時にその記事が属するカテゴリを分類するタスク（9クラス分類）で検証しました。 具体的な流れとしては、一度Transformersのrun_glue.pyで東北大BERT（cl-tohoku/bert-base-japanese-whole-word-masking）をfine-tuningします。 その後、各種ライブラリでfine-tuningしたモデルを量子化、その量子化モデルを用いてテストデータ（1,473件）に対して推論を行いました。

推論はCPU上で行い、batch sizeは1、max lengthは512で行いました。 また推論時間3はテストデータそれぞれの平均をとしています。

Static Quantizationの実装はoptimumのONNX用run_glue.pyを元にしました。

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

  • onnx: 1.11.0
  • onnxruntime: 1.10.0
  • optimum: 1.2.1
  • transformers: 4.16.2

推論速度の比較

量子化（今回はStatic Quantization）によってどのくらい推論速度が減少するか調査しました。 比較対象としては、量子化を行わずに実行した場合（PyTorch fp32）とDynamic Quantizationを利用した場合（Dynamic Quant）4を実験しています。 結果が以下のグラフとなります。

図1: 量子化した際の推論速度

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

1つ気になる点として、Static QuantizationよりDynamic Quantizationの方が推論速度が速い点が挙げられます。 実はONNX Runtimeの量子化に関するドキュメントに以下のような記載が存在します。

In general, it is recommended to use dynamic quantization for RNNs and transformer-based models, and static quantization for CNN models.

なぜRNNやTransformer系のモデルはDynamic Quantizationの方が推奨されているか明記されていないので断言できませんが、おそらくTransformerのモデル構造が原因となってDynamic Quantizationの方が推論速度が速い結果になったと思われます。

まとめ

今回の記事では、BERTに対してStatic Quantizationを簡単に動かしてみました。 前回のDynamic Quantizationに引き続き、optimumを利用してONNX Runtimeでの量子化を行いました。 結果としてStatic Quantizationを行うことで、量子化を行わない場合より推論速度を高速化できました。 一方でStatic QuantizationよりもDynamic Quantizationの方が推論速度が速いこともわかりました。

今回の結果で、事前にscale factorを計算する手法をBERTで試すモチベーションが多少失われましましたが、せっかくなので次回はQuantization Aware Trainingを検証しようと思います。

  1. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova. [paper]↩

  2. Efficient 8-Bit Quantization of Transformer Neural Machine Language Translation Model. Aishwarya Bhandare, Vamsi Sripathi, Deepthi Karkada, Vivek Menon, Sun Choi, Kushal Datta, Vikram Saletore. [paper]
    I-BERT: Integer-only BERT Quantization. Sehoon Kim, Amir Gholami, Zhewei Yao, Michael W. Mahoney, Kurt Keutzer. [paper]

    前者のBhandare et al.ではLayer NormalizationやSoftmaxに対する量子化の難しさが議論されています。
    後者のKim et al.はこれら量子化の難しい処理に対して工夫を行い、量子化を行う手法を提案しています。

    ↩

  3. ここでの推論時間は、モデルロードなどは含んでいません。テストデータをdata_loaderでモデルに入力し、すべてのテストデータが推論されるまでの時間を計測、1サンプルあたりの平均時間を計測しました。↩

  4. Dynamic Quantizationは前回も実験していますが、今回ライブラリのバージョンが変わったこともあり改めて実験しています。↩