ニューラルネットワークの量子化まとめ①

論文読みが主です

　現在勉強中のニューラルネットワークの量子化について、論文を調べて学んだことや、実際にPythonで（気が向いたらFPGAにも）実装したことをまとめていこうと思います。

　今回は、こちらの論文を参考に、ニューラルネットワークの量子化についてざっくりと雰囲気だけ話していこうと思います。Pythonで、ちょっとだけ実装もしてみました。

ニューラルネットワークの量子化とは

　量子化とは、ニューラルネットワークにおける重みなどのパラメータや、入力値・活性値を、float 型から低ビットの int 型に置き換えようというものです。これによって、ハードウェア上に実装するにあたって、メモリ量や計算資源、消費電力の削減や推論の高速化が期待されます。

　まずは、具体的にどうやって浮動小数点を整数に直すのかについてまとめていきます。

量子化の方法

　8bit 整数への量子化を前提に話していきます。量子化前の浮動小数点の値を $x_f$ 、量子化後の整数の値を $x_i$ とすると、次のように変換がなされます。

$$　　x_f=S(x_i−Z) $$

　ここで出てきた $S$ は、スケーリングパラメータというもので、$x_i$ が 1 増加すると $x_f$ がどれだけ増加するかを表す float 値です。また、$Z$ は、ゼロポイントといって、$x_f$ が 0 となるような $x_i$ の値です。

　逆変換はこうなります。$Round$ は、丸め関数です。

$$ x_i=Round(x_f/S)+Z $$

　$x_i$ のとりうる値は $ -128 \leq x_i \leq 127 $ と制限されているので、量子化すべき $x_f$ の値の範囲が広くなると、Sを大きくしないといけません。一方で、$S$ が大きくなれば、量子化による丸め誤差は大きくなってしまいます。値が大きすぎる、あるいは小さすぎる $x_f$ については、 量子化の範囲の外に出し、$ -128 \leq x_i \leq 127 $ の範囲にクリップすることも時には必要です。

　　　　　　　　　　
　　　　　　　　　

　　　　　　　　　

　ところで、画像の入力値や ReLU の出力のような場合、$x_f$ は、正の値しかとらないことになります。この場合は、$x_i$ を unsigned int として、$ 0 \leq x_i \leq 255 $ の範囲に量子化します。

値の伝播

　上のやり方で、重みや活性値を量子化したとして、値の伝播はどのように行われるのでしょうか。浮動小数点の入力を $x_f$ , 重みを $w_f$ として、その行列積 $y$ を

$$ y = w_f ・ x_f $$

と書くことにします。これに先ほどの量子化の式を当てはめると、
　

$$ y = (S_ww_i + Z_w)・(S_xx_i + Z_x) $$

となります。簡単のため、$Z_w, Z_x$ は 0 として話を進めていきます（後で考えます）

　活性化関数としては、シグモイド関数などの曲線的な関数の量子化を考えるのはすこし難しそうなので、ReLUを考えます。活性値を、

$$ a_f = ReLU(y) $$

として、$a_f$ のスケーリングパラメータを$S_a$とします。これを、ゼロポイント $Z_a$ を 0 とし、unsigned で量子化したものを $a_i$ とします。

　以上をまとめると、量子化された値 $x_i, w_i, a_i $ の間には次の式が成り立ちます。
　($w_i・x_i$ < 0 の場合は $a_i$ は 0 です）

$$ S_aa_i = S_wS_xw_i・x_i $$

ここで、$a_i$ と $w_i・x_i$ は整数で、$S_a, S_w, S_x$ は浮動小数点です。このままでは、浮動小数点演算が残ってしまいます。すべての演算を整数演算で置き換えるためには、 $S_wS_x/S_a$ をシフト演算で近似すると良さそうです。つまり、

$$ 2^{-n} \leq S_wS_x/S_a < 2^{-n+1} $$

を満たす $n$ を求めて、

$ \qquad\qquad\qquad\qquad\qquad\qquad\quad a_i = w_i・x_i $ >> $n$

とすればOKです。(>> は右シフトです)

　 これで、量子化された活性値 $a_i$ を、量子化された重みと入力から、整数演算のみにより計算することができました。


　ところで、しれっとなきものにしていたバイアス項ですが、$S_wS_x$をスケールパラメータとして量子化し、$w_i・x_i$ に足せば良いかなと思いました。（ちょっと自信ないです）

ゼロポイントの考慮

　上の計算では、ゼロポイントを無視していましたが、これを考慮する場合、行列積は次のようになります。

$$ y 　=　 S_wS_xw_i・x_i 　+　 S_ww_i・Z_x　+　 S_xZ_w・x_i 　+　 Z_w・Z_x $$

　$ Z_w, Z_x$ は、それぞれ $w_i, x_i$ と同じサイズに拡張されています。

　ゼロポイントの追加で、第２項〜第４項が追加されました。第２項と第４項は、推論するより前にあらかじめ計算しておける値です。一方で、第３項は、推論時に入力値 $x_i$ をみて計算する値です。第３項の計算時間等のロスを考えると、重みの量子化のゼロポイント $Z_w$ は 0 として、第３項をなくすのが良いかもしれません。

学習後の量子化 & 学習中の量子化

　ここまで話してきたパラメータの量子化についてですが、量子化を行うタイミングとしては、以下の２種類があるようです。

1. 学習後に量子化する。
2. 量子化しながら学習する。

　1. は非常に分かりやすいやり方で、モデルの学習は普通のニューラルネットワークと同様に行い、学習後にそのパラメータを量子化したり、活性値を量子化する部分を追加したりするというものです。これについて、この後に、PyTorchでの実装の一例を載せます。

　2. は、ニューラルネットワークのパラメータとして浮動小数点の値をもっておきつつ、値の伝播は、それらを量子化したパラメータを用いて行うというものです。これについては、私はまだよく理解できていないので、また勉強して詳しいことが分かったら、記事を投稿できればと思っています。

PyTorchで実装

　学習後の量子化について、Githubの方にコードを載せました。PyTorchには量子化のためのライブラリはもちろんありましたが、理解を深めるという意図で、量子化部分は一から実装しました。

　PyTorchのライブラリを用いた量子化方法については、こちら。

モデルの定義

　モデルは、とりあえず簡単なものでやってみました。別のモデルを実装してみたものも、Githubの方にあります。

　入力は MNISTデータセットで、1024ノードの全結合層 + ReLU と出力層が続きます。量子化に際して、追加したものが、ReLU の後のシフト演算です。これは、全結合層の出力を 8bit の整数に直すためのものです。 shiftM については、quantize.py というプログラムで計算されます。

model.py

 
  4 # 入力はMNIST
  5 class Model(nn.Module):
  6     def __init__(self):
  7         super().__init__()
  8         self.fc1 = nn.Linear(784, 1024, bias = False)
  9         self.fc2 = nn.Linear(1024, 10, bias = False)
 10     def forward(self, x):
 11         x = x.reshape(-1, 784)
 12         x = self.fc1(x)
 13         x = torch.relu(x)
 14         x = self.fc2(x)
 15         return x
 16 
 17 class QuantizedModel(Model):
 18     def __init__(self, shiftM):
 19         super().__init__()
 20         self.shiftM = shiftM
 21     def forward(self, x):
 22         x = x.reshape(-1, 784)
 23         x = self.fc1(x)
 24         x = torch.relu(x)
 25         # 再量子化
 26         x = (x.int() >> self.shiftM).float()
 27         x = x.clip(0, 255)
 28         x = self.fc2(x)
 29         return x

量子化

入力の量子化

　入力値は、MNISTの各ピクセルの値（ 0 ~ 1 ) なので、スケーリングパラメータ s0 は
1 / 255、ゼロポイント z0 は 0 、unsigned で量子化します。

 20 # 入力の量子化パラメータ
 21 s0 = 1 / 255
 22 z0 = 0

全結合層の重みの量子化

　ゼロポイントは 0 、singed で量子化することにしました。浮動小数点の重みで絶対値が最大のものを 127 でわって、スケーリングパラメータとしたのち、量子化します。
　出力層 fc2 の重みの量子化も同様です。

 24 # fc1 の重みの量子化パラメータ
 25 w1 = model1.fc1.weight.data
 26 s1 = max(w1.max(), w1.min() * -1) / 127
 27 z1 = 0
 28 quantized_w1 = (w1 / s1).round()

ReLU の出力のシフト

　活性値の量子化は、重みや入力と違い、量子化の幅がわかりません。今回は、訓練データを１回回して出てきた活性値の最大値 amax を、量子化の最大値としました。このやり方は、活性値の外れ値の影響を受けてしまうという理由から、考え直すべきかもしれません。

これをもとに、活性値のスケーリングパラメータ sa を決めたら、値の伝播 のところで述べたシフト量の計算を行います。

30 # fc1 + relu の活性値の量子化パラメータ
31 # reluの出力を見て、スケールパラメータを決める
32 amax = 0
33 amin = 0
34 for (x, t) in train_dataloader:
35     x = x.reshape(-1, 784)
36     x = model1.fc1(x)
37     x = torch.relu(x)
38     if amax < x.max():
39         amax = x.max()
40 sa = amax / 255
41 
42 # スケールパラメータsaが決まったら、活性値の出力値を
43 # sa / s0 / s1 で割るようにする。そして、これをシフト演算で置き換える
44 M = sa / s1 / s0
45 shiftM = 0
46 while M > 1:
47     M /= 2
48     shiftM += 1
49 

精度の比較

　量子化前と量子化後で両方 98.3 % と、全く変わりませんでした。もっと複雑なモデル、データセットでもやってみたいですね。

まだまだ続きます（多分）

　今回はこのあたりで終わろうと思います。

　実装したのは、「学習後に量子化する」というものでしたが、「学習中に量子化する」という方法もありますし、他にもまだまだ量子化の手法はあります。また何かしら記事にまとめられたらなと思います。