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

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

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

Dynamic Quantizationの概要
各種ライブラリでの量子化方法紹介
量子化実験
まとめ

Dynamic Quantizationの概要

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

この手法の利点として、一度学習したモデルがあれば、簡単に適用することができることが挙げられます。今回は実際に、自然言語処理でよく使われるBERT（Bidirectional Encoder Representations from Transformers）¹と呼ばれるモデルに対していくつかのライブラリでDynamic Quantizationを試してみたいと思います。具体的には以下のライブラリでDynamic Quantizationを試してみました。

PyTorch
ONNX Runtime
Intel Neural Compressor

各種ライブラリでの量子化方法紹介

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

1. PyTorchでDynamic Quantization

PyTorchには次の2種類の量子化実装が存在しています。

Eager Mode Quantization
FX Graph Mode Quantization

今回の実験設定でお手軽にDynamic Quantizationをやりたい、ということであれば1番のEager Mode Quantizationを使うことになります。詳細はPyTorchの公式ドキュメントを確認してください。

BERTは線型結合層（torch.nn.Linear）やEmbedding層（torch.nn.Embedding）などによって構成されています。本来であればそれらすべてを量子化したいところですが、公式ドキュメントによると残念ながら2022年2月現在、torch.nn.EmbeddingはDynamic Quantizationに対応していないようです。そこで、今回は線型結合層のみを量子化してみようと思います。

この量子化は以下のようにすることで実現できます。

# huggingface/transformersのBERTモデルの読み込み
model_fp32 = AutoModelForSequenceClassification.from_pretrained(/path/to/model_fp32)

# 線型結合層を量子化
model_int8 = torch.quantization.quantize_dynamic(model_fp32, {torch.nn.Linear}, dtype=torch.qint8)

このmodel_int8はすでに量子化済みのモデルなので、このモデルを使って推論を行うことで通常のfp32モデルより速く推論ができ、量子化によってモデルのデータサイズも小さくなっています。

2. ONNX RuntimeでDynamic Quantization

ONNX Runtimeで量子化を行う際は、以下のようにONNXモデルを量子化する処理を行う必要があります。

from onnxruntime.quantization import quantize_dynamic

quantize_dynamic(
     model_input=/path/to/onnx_model_fp32,
     model_output=/path/to/onnx_model_int8
)

このようにquantize_dynamicを使って量子化したONNXモデルを使用した推論を行うことで、Dynamic Quantizationを実現できます。

しかし、この場合だと、Transformersライブラリでfine-tuningしたモデルをONNXモデルに変換し、それをquantize_dynamicに通して量子化するというように少し段階を踏まないといけません。便利なもので、実はPyTorch BERTをONNXモデルに変換して、それを量子化するライブラリ、スクリプトはいくつか存在します。今回は2つ、お手軽にBERTを量子化したONNXモデルに変換する方法を紹介します。

Transformers主体で量子化

Transformersの中にはBERTをこのONNXモデルに変換し、そのまま量子化まで行うスクリプトconvert_graph_to_onnx.pyがあります。具体的には以下のスクリプトを実行すると、自動で量子化済みのONNXモデルが作成されます。²

python convert_graph_to_onnx.py \
  --model /path/to/model_fp32 \
  --pipeline sentiment-analysis \ # たとえば分類タスクでfine-tuningした場合
  --tokenizer tokenizer_name \
  --framework pt \
  --quantize \
  output_dir/onnx

このスクリプトを実行することで、output_dir/onnx-quantizedに量子化済みのモデルが作成されるので、これを使って推論を行うことでよりお手軽にDynamic Quantizationを実現できます。実際のスクリプトを確認していただけるとわかる通り、このスクリプトは内部でBERTをONNXモデルに変換した後、onnxruntime.quantization.quantizeを用いてモデルを量子化しています。

ちなみに、現在ONNX Runtimeライブラリには量子化関数としてquantize、quantize_dynamic、quantize_staticが存在しています。その中のquantizeは引数次第でDynamic QuantizationとStatic Quantizationのどちらも可能な関数なのですが、実際に使ってみると

onnxruntime.quantization.quantize is deprecated. Please use quantize_static for static quantization, quantize_dynamic for dynamic quantization.

と出力されるので、基本的にはquantize関数はもう使わない方が良さそうです。

huggingface/optimumを使って量子化

上述の通り、convert_to_graph_to_onnx.pyは非推奨の関数を使っているため、そこまで率先して使うようなスクリプトではないと思います。 Hugging Face社はTransformersの拡張として、optimumというライブラリを用意していて、これを使うことで簡単にPyTorch BERTをONNXモデルに変換、量子化までしてくれます。

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

from optimum.onnxruntime import ORTConfig, ORTQuantizer

# https://github.com/huggingface/optimum/blob/main/optimum/onnxruntime/configuration.pyを確認してconfigを設定
ort_config = ORTConfig(quantization_approach="dynamic")
quantizer = ORTQuantizer(ort_config)
quantizer.fit(
       /path/to/model_fp32,
       output_dir,
       feature="sequence-classification" # たとえば分類タスクでfine-tuningした場合
)

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

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

3. Intel Neural CompressorでDynamic Quantization

Intel Neural CompressorはIntel社が公開しているPythonライブラリで、こちらはIntel CPUまたはGPU上でモデルの圧縮を行い動作させるライブラリです。

実はこのライブラリを用いた量子化もoptimumがサポートしていて、簡単に使うことができます。具体的には以下の通りです。

import torch
from optimum.intel.neural_compressor import IncOptimizer, IncQuantizationConfig, IncQuantizationMode, IncQuantizer, IncTrainer
from optimum.intel.neural_compressor.quantization import IncQuantizedModelForSequenceClassification
from transformers import AutoModelForSequenceClassification

model_fp32 = AutoModelForSequenceClassification.from_pretrained(/path/to/model_fp32)

q8_config = IncQuantizationConfig.from_pretrained(/path/to/config_dir, config_file_name="quantization.yml")
quant_approach = getattr(IncQuantizationMode, "DYNAMIC").value
q8_config.set_config("quantization.approach", quant_approach)

# 本来、eval_funcには、modelを与えたらそれを使って推論を行い、その推論の評価スコアが返ってくる関数が来る想定です
# ですが、今回は量子化だけなので、適当な関数を入れています
inc_quantizer = IncQuantizer(model_fp32, q8_config, eval_func=lambda x: 1.0)
# q_model.modelが量子化後のモデル
q_model = inc_quantizer.fit_dynamic()
# torch標準の方法で保存する場合は以下
torch.save(q_model.model.state_dict(), /path/to/model_int8)
# optimumのTrainerを利用する場合は以下
trainer = IncTrainer(args) # Transformersと同じように引数を埋める
trainer.model = q_model.model
trainer.save_model(/path/to/model_int8)

なお、この手順を実行する前に、このquantization.ymlを手元に用意し、/path/to/config_dirで配置したディレクトリを指定してください。また、今回の検証でIntel Neural Compressorを動かす際に、LANG=en_us.utf-8の設定が必要でした。

この手順を行うことで/path/to/model_int8以下に量子化されたモデルが作成されます。 ONNXの場合と同様に、Intel Neural Compressorでも、huggingface/optimumを使うことで、お手軽にPyTorchのBERTを量子化できました。 optimumは本当に便利ですね。

量子化実験

それでは、実際にここまで紹介してきたライブラリを使い、量子化（Dynamic Quantization）をしてみようと思います。

実験設定

今回は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で行いました。また推論時間³はテストデータそれぞれの平均をとしています。

Dynamic Quantizationの実装はそれぞれ以下の通りとなっています。

PyTorch: Transformersのrun_glue.py内部に、本記事で紹介した量子化を行う処理を追記
ONNX: optimumのONNX用run_glue.pyを使用
Intel Neural Compressor: optimumのIntel Neural Compressor用run_glue.pyを使用

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

neural-compressor: 1.9.1
onnx: 1.11.0
onnxruntime: 1.10.0
optimum: 48b2a41commitをpip install
transformers: 4.16.2

モデルサイズの比較

まずは、量子化によってどのくらいモデルのデータサイズが減少するか調査しました。結果が以下のグラフとなります。

f:id:ssskkk420:20220225103058p:plain — 図1: 量子化した際のモデルのデータサイズ

PyTorch fp32がベースライン、つまり量子化前のモデルのデータサイズとなっています。このグラフからわかる通り、各量子化手法を使うことでモデルのデータサイズを削減できました。とくにONNXモデルにする⁴ことで、モデルのデータサイズを約1/4程度まで減少させることができました。

推論速度の比較

次に、量子化（今回はDynamic Quantization）によってどのくらい推論速度が減少するか調査しました。結果が以下のグラフとなります。

f:id:ssskkk420:20220225103139p:plain — 図2: 量子化した際の推論速度

こちらのグラフでも、PyTorch fp32が量子化前のモデルの推論速度となっています。このグラフからわかる通り、各量子化手法を使うことで推論をより高速にできました。とくにONNX Runtimeでは、量子化前から推論速度を約41.1%ほど減少できています。

この例を見ると明らかにONNX Runtimeが他の手法と比べて優れているように見えます。これは量子化だけでなく、そもそもONNX Runtime自体が推論高速化されているためにそう見えると思われます。 ONNX Runtimeに関しても、弊社技術ブログにて簡単に紹介していますので、気になった方はそちらをご確認ください。

まとめ

今回の記事では、BERTに対してDynamic Quantizationを簡単に動かしてみました。ライブラリとしてはPyTorch内部実装、ONNX Runtime、Intel Neural Compressorの3種を試しました。結果として、Dynamic Quantizationを行うことでモデルのデータサイズを削減でき、推論速度も高速化できました。とくにONNXおよびONNX Runtimeはデータサイズの削減、推論速度の高速化の2点でもっとも良い結果となりました。

ここまで読んでいただくと、Dynamic QuantizationをしたいときはONNX Runtimeがおすすめ、といった内容なのですが、私としては必ずしもそうではないと思います。 1行追加するだけで量子化できるPyTorch標準の方法は本当にお手軽だと思いますし、Intel Neural Compressorはモデルの枝刈り（pruning）なども簡単に実現できるのでそういった高速化手法と組み合わせる際には選択肢に挙がってくると思います。単純に推論の高速化を頑張りたいのであればONNX Runtimeを使うなど、その時の用途にあったライブラリ選択をおすすめします。

また、ONNX RuntimeとIntel Neural Compressorに関してはoptimumを使用することで、本当に簡単に試すことができました。 optimumはまだまだ開発中なので、これからも触っていきたいと思います。

今後の記事では、Dynamic Quantizationではなく、Static Quantizationを利用した場合、推論速度が向上するのかなどについて検証したいと思います。

BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova. [paper]↩
ちなみに、Transformersは現在ONNXモデルへの変換はconvert_graph_to_onnx.pyではなく、python -m transformers.onnxを利用した方法を推奨（参考ページ）しているようです。残念ながら、推奨された変換方法では一気に量子化はできないので、今回はconvert_graph_to_onnx.pyを利用した方法を使用しています。↩
ここでの推論時間は、モデルロードなどは含んでいません。テストデータをdata_loaderでモデルに入力し、全てのテストデータが推論されるまでの時間を計測、1サンプルあたりの平均時間を計測しました。↩
ちなみに、PyTorch fp32モデルをONNXモデルに変換した際のモデルデータサイズは423MBでした。つまり、ONNXは単純な変換ではモデルサイズは減少せず、量子化を行うことでモデルサイズを減少させることができるようです。↩