ディープラーニングのモデルを推論する際、通常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
結果
精度
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の方が速くなる。