機械学習モデルの量子化
量子化はモデル圧縮の手法の一つであり、モデルのパラメータなどを低bitで近似表現するものです。決して新しい手法ではありませんが、LLM（大規模言語モデル）が一般的になってきた中で耳にすることも多くなってきたかと思います。多くのパラメータを有し、モデルサイズも大きくなりがちなLLMモデルにおいてその効果は大きく、さまざまなモデルや場面において活用されています。
量子化については大規模なモデルに限らず、
- 演算あたりのビット数が削減できる
- モデルサイズが削減できる
という観点において汎用的に活用していけるものです。
しかしながら、上記効果は事実ではあるものの「ありとあらゆる場面において効果があるか」という意味では必ずしもそうとは限りません。
この記事では、機械学習モデルの量子化が効果的でない場面・条件について確認していきたいと思います。
量子化の整理と本記事の前提条件
簡単に言うと、FP32は32 bitの浮動小数点数で有効桁数はおよそ7桁、FP16は16 bitで有効桁数はおよそ3桁強になります。一方、INT8量子化では8 bitの整数とスケール係数で元の FP32値を近似表現します。
これによって、近年肥大化してきているモデルサイズを圧縮し、メモリに乗せることを手助けしています。
特にこのようなメモリ帯域に対する貢献はわかりやすい点かと思います。
ここで本稿の背景について少し触れておきます。
個人としては実務において、大規模なモデルよりも小型で高速なモデルを取り扱うことの方が多いです。
CPU上で、単バッチ、30fps以上、数msの低遅延でリアルタイム処理が可能なモデルが一つの参考です。また基本的にはPython/PyTorchで学習させたのち、ONNXに変換してデプロイするシチュエーションを想定します。今回の検証はONNX Runtime CPU EPを利用し、Apple M2で実施します。
結論としてこのような条件下では、量子化処理は相性が悪く、むしろパフォーマンスに悪影響を及ぼすことがあります。
モデルサイズは圧縮されますが、推論効率については注意が必要であるために具体例とともに観測していきたいと思います。
なお、この記事で扱った条件は、量子化モデルにとってはかなり不利な設定になっている点には注意が必要です。
特に、推論にGPUを用いる場合や、他の推論エンジン（OpenVINOなど）を利用する場合は、本記事と異なる傾向になることが多く、ここでの結果をそのまま当てはめることはできません。
検証
それではモデルを定義して、モデルサイズの圧縮を観測しながら推論速度を計測して量子化の効果をみていきます。
モデル定義
まずは軽量の適当なモデルを定義します。
全結合層を4層重ねただけの非常に単純なネットワークで検証します。
class TestNetwork(nn.Module):
def __init__(self, input_size=100, hidden_size=256, output_size=10):
super(TestNetwork, self).__init__()
self.fc1 = nn.Linear(input_size, hidden_size)
self.relu1 = nn.ReLU()
self.fc2 = nn.Linear(hidden_size, hidden_size)
self.relu2 = nn.ReLU()
self.fc3 = nn.Linear(hidden_size, hidden_size)
self.relu3 = nn.ReLU()
self.fc4 = nn.Linear(hidden_size, output_size)
def forward(self, x):
x = self.fc1(x)
x = self.relu1(x)
x = self.fc2(x)
x = self.relu2(x)
x = self.fc3(x)
x = self.relu3(x)
x = self.fc4(x)
return x
モデルサイズ
torchinfoのsummaryを利用してモデルのパラメータ数などを確認してみます。
Model Summary:
----------------------------------------------------------------------
==========================================================================================
Layer (type:depth-idx) Output Shape Param #
==========================================================================================
SimpleNetwork [1, 10] --
├─Linear: 1-1 [1, 256] 25,856
├─ReLU: 1-2 [1, 256] --
├─Linear: 1-3 [1, 256] 65,792
├─ReLU: 1-4 [1, 256] --
├─Linear: 1-5 [1, 256] 65,792
├─ReLU: 1-6 [1, 256] --
├─Linear: 1-7 [1, 10] 2,570
==========================================================================================
Total params: 160,010
Trainable params: 160,010
Non-trainable params: 0
Total mult-adds (M): 0.16
==========================================================================================
また実際のONNXファイルのサイズ比較は下記の通りです。
ONNX化したモデルを量子化し、確認していきます。
手順としてはONNX公式で記される方法に基づきます。
FP32: 0.61 MB
FP16: 0.31 MB (50.1%)
INT8: 0.16 MB (26.0%)
想定通り、モデルサイズを圧縮することができています。
大きなモデルであればその削減効果も大きいということが容易に想像がつく結果かと思います。
推論速度
それでは、それぞれのモデルについて推論速度を計測していきます。
ランタイムはONNX Runtime CPU EPです。
バッチサイズは1で、100回の推論でWarmupさせたのち、1000回の推論時間を計測して比較します。
Metric FP32 FP16 INT8
----------------------------------------------------------------------
Mean (ms) 0.0131 0.0345 0.0189
Std Dev (ms) 0.0010 0.0032 0.0050
Min (ms) 0.0128 0.0338 0.0178
Max (ms) 0.0309 0.0960 0.0920
Total (ms) 13.11 34.49 18.91
----------------------------------------------------------------------
結果として、量子化を施していないFP32が最も早く、次点でINT8が早く、FP16はわかりやすく所要時間が長いことがわかりました。
参考までにResNet50の結果を記載しておきます。
こちらでもFP32が最も高速であり、FP16とINT8はいずれもFP32より遅い結果となりました。
Metric FP32 FP16 INT8
----------------------------------------------------------------------
Mean (ms) 29.4577 32.1128 42.7559
Std Dev (ms) 4.0022 8.9304 3.1175
Min (ms) 28.0056 30.2958 41.9033
Max (ms) 128.1853 241.7661 123.7010
Total (ms) 29457.74 32112.75 42755.90
----------------------------------------------------------------------
考察
今回の結果に至った要因についてまとめていきます。
まずFP16だけ明らかに遅くなったのは、一般的な汎用CPUでは、FP16専用の演算ユニットが十分でない場合が多く、ライブラリ側で内部的にFP16 → FP32に変換して計算し、結果だけFP16に戻す実装が想定されます。
その場合、変換コストがかかるため、FP32より速くならないケースがありえます。
INT8がFP32よりも劣ったのは、INT8量子化されたモデルであっても、すべての処理を真にINT8だけで完結させることが難しいためだと考えられます。一部の演算はINT8では扱いにくく、部分的にFP32にキャストして処理している可能性があります。そのため、行列演算そのものの高速化よりも、その前後で発生するquantize / dequantizeやFP32へのupcastのコストが相対的に大きくなり、結果として所要時間が増加していると考察できます。実際こうした内容についてはONNX RuntimeのFAQにも記載があります。
なお、今回のモデルはFP32でも約0.61MBと小さく、近年のCPUのL2/L3キャッシュに十分収まるサイズです。
そのため、DRAMからのメモリアクセスが支配的になっているとは考えにくく、量子化によるメモリ帯域削減のメリットが表面化しにくい条件でもあります。一方で、ResNet50のようなより大きなモデルではキャッシュに収まりきらない部分も増えるため、理屈上では量子化によるメモリ帯域削減効果が現れやすくなるはずです。それでも今回の条件では、FP32が最速であったのはメモリ帯域よりも各種変換処理に伴うオーバーヘッドのほうが支配的であったからであると考えることができます。
おわりに
今回の検証によって、量子化はモデルサイズの圧縮には非常に有効ですが、今回のような条件（CPU / 小型モデル / バッチサイズ 1 / 低遅延重視 / ONNX Runtime CPU EP）では、推論速度が必ずしも改善せず、むしろ悪化する場合があることが分かりました。実運用では、対象モデル・ハードウェアに合わせて、量子化を採用するかどうか判断をしつつ、適切なランタイムを選択することが推奨されます。
テストで用いたモデルは非常にシンプルなのですが、実際にはもう少し複雑でも同様な傾向が観察できます。
また一般に、最近のGPUにはFP16やINT8専用の演算ユニットを備えたものも多く、量子化との相性が良い傾向があります。さらにバッチ処理のように一度に多くのデータを処理する場合は、quantize / dequantizeといった固定オーバーヘッドが相対的に小さくなるため、スループットの向上を実感しやすくなります。
その意味で、今回の条件は量子化にとっては不利なものだったと言えます。
量子化に伴う速度劣化について悩まれることがありましたら今回の内容を思い出していただければ幸いです。