TadaoYamaokaの開発日記

個人開発しているスマホアプリや将棋AIの開発ネタを中心に書いていきます。

PyTorchのFXグラフモードで量子化認識トレーニングを試す

ディープラーニングのモデルを推論する際、通常GPUが必要である。しかし、GPUがない場合でも、モデルのパラメータを浮動小数点から整数へと変換する手法を使うことで、CPUだけでも推論処理を高速に行うことが可能である。この手法を量子化と呼ぶ。

PyTorchの量子化

PyTorchには、量子化の方法がいくつか用意されている。

実装は「Eager Mode Quantization」と「FX Graph Mode Quantization」に分かれており、「FX Graph Mode Quantization」の方が新しい実装で、「Eager Mode Quantization」ではモデルのレイヤー融合を手動で行う必要があったところが自動化されている。

また、一般的に量子化の手法として、

がある。

ダイナミック量子化は、キャリブレーション(浮動小数点から整数へ変換する際のスケールの調整)なしで実行できてお手軽だが、畳み込み層は量子化されないなど制限がある。

レーニング後の量子化量子化認識トレーニングは、どちらも事前にキャリブレーションを行う方法だが、量子化認識トレーニングは追加で行う訓練時に調整を行うため、より精度が高くなる。

手法の比較は、PyTorchのドキュメントにまとまっている。
Quantization — PyTorch 2.0 documentation

FXグラフモードの量子化認識トレーニングを試す

ここでは、CNNで最も精度が高くなる量子化認識トレーニングを試す。
また、レイヤー融合を自動で行えるFXグラフモードで実装する。

実装は、Eager Modeでの量子化認識トレーニングのサンプルコードがあった以下の記事を参考に、FXグラフモードに変更した。
PyTorch Quantization Aware Training - Lei Mao's Log Book

処理の流れ
  1. モデルを通常に訓練する(データセットCIFAR10を使用してResNet18のモデルを訓練)
  2. get_default_qat_qconfig_mappingとprepare_qat_fxで、量子化認識トレーニング用モデルを準備する
  3. 量子化のために追加で訓練する
  4. convert_fxで量子化モデルに変換する
  5. ONNXに出力する


結果

精度
FP32 evaluation accuracy: 0.812
INT8 evaluation accuracy: 0.859

通常の訓練を行った時点のFP32のモデルと、量子化認識トレーニングを行ったINT8のモデルを比較すると、後者の方が正解率が高くなっている。
これは、量子化により精度が高くなったというより、量子化のために追加の訓練を行った効果も含まれている。
量子化により大きく精度が落ちないことは確認できた。

推論速度
FP32 CPU Inference Latency: 5.52 ms / sample
FP32 CUDA Inference Latency: 2.17 ms / sample
INT8 CPU Inference Latency: 1.29 ms / sample
INT8 JIT CPU Inference Latency: 0.75 ms / sample

FP32モデルをCPUとGPU(CUDA)で実行した場合は、後者が速い。
FP32モデルとINT8モデルをCPUで実行した場合は、後者が速く、量子化に効果があることが確かめられた。
INT8モデルをPyTorchでモデルをコンパイルするJITモデルに変換すると、さらに速くなっている。

なお、FP32モデルをGPUで推論するより、INT8モデルをCPUで推論する方が速くなっているのは、バッチサイズが1であるためである。バッチサイズが大きい場合は通常GPUの方が速くなる。

まとめ

PyTorchのFXグラフモードで量子化認識トレーニングを試した。
結果、精度が大きく落ちることなく、推論速度が速くなることが確認できた。