神經網絡量化入門--基本原理

最近打算寫一個關於神經網絡量化的入門教程，包括網絡量化的基本原理、離線量化、量化訓練，以及全量化模型的推理過程，最后我會用 pytorch 從零構建一個量化模型，幫助讀者形成更深刻的理解。

之所以要寫這系列教程，主要是想幫助初次接觸量化的同學快速入門。筆者在剛開始接觸模型量化時走了很多彎路，並且發現網上的資料和論文對初學者來說太不友好。目前學術界的量化方法都過於花俏，能落地的極少，工業界廣泛使用的還是 Google TFLite 那一套量化方法，而 TFLite 對應的大部分資料都只告訴你如何使用，能講清楚原理的也非常少。這系列教程不會涉及學術上那些花俏的量化方法，主要是想介紹工業界用得最多的量化方案 (即 TFLite 的量化原理，對應 Google 的論文 Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference )

話不多說，我們開始。這一章中，主要介紹網絡量化的基本原理，以及推理的時候如何跑量化模型。

背景知識

量化並不是什么新知識，我們在對圖像做預處理時就用到了量化。回想一下，我們通常會將一張 uint8 類型、數值范圍在 0~255 的圖片歸一成 float32 類型、數值范圍在 0.0~1.0 的張量，這個過程就是反量化。類似地，我們經常將網絡輸出的范圍在 0.0~1.0 之間的張量調整成數值為 0~255、uint8 類型的圖片數據，這個過程就是量化。所以量化本質上只是對數值范圍的重新調整，可以「粗略」理解為是一種線性映射。(之所以加「粗略」二字，是因為有些論文會用非線性量化，但目前在工業界落地的還都是線性量化，所以本文只討論線性量化的方案)。

不過，可以明顯看出，反量化一般沒有信息損失，而量化一般都會有精度損失。這也非常好理解，float32 能保存的數值范圍本身就比 uint8 多，因此必定有大量數值無法用 uint8 表示，只能四舍五入成 uint8 型的數值。量化模型和全精度模型的誤差也來自四舍五入的 clip 操作。

這篇文章中會用到一些公式，這里我們用 \(r\) 表示浮點實數，\(q\) 表示量化后的定點整數。浮點和整型之間的換算公式為：

\[r = S(q-Z) \tag{1} \]

\[q = round(\frac{r}{S}+Z) \tag{2} \]

其中，\(S\) 是 scale，表示實數和整數之間的比例關系，\(Z\) 是 zero point，表示實數中的 0 經過量化后對應的整數，它們的計算方法為：

\[S = \frac{r_{max}-r_{min}}{q_{max}-q_{min}} \tag{3} \]

\[Z = round(q_{max} - \frac{r_{max}}{S}) \tag{4} \]

\(r_{max}\)、\(r_{min}\)分別是 \(r\) 的最大值和最小值，\(q_{min}\)、\(q_{max}\)同理。這個公式的推導比較簡單，很多資料也有詳細的介紹，這里不過多介紹。需要強調的一點是，定點整數的 zero point 就代表浮點實數的 0，二者之間的換算不存在精度損失，這一點可以從公式 (2) 中看出來，把 \(r=0\) 代入后就可以得到 \(q=Z\)。這么做的目的是為了在 padding 時保證浮點數值的 0 和定點整數的 zero point 完全等價，保證定點和浮點之間的表征能夠一致。

矩陣運算的量化

由於卷積網絡中的卷積層和全連接層本質上都是一堆矩陣乘法，因此我們先看如何將浮點運算上的矩陣轉換為定點運算。

假設 \(r_1\)、\(r_2\) 是浮點實數上的兩個 \(N \times N\) 的矩陣，\(r_3\) 是 \(r_1\)、\(r_2\) 相乘后的矩陣：

\[r_3^{i,k}=\sum_{j=1}^N r_1^{i,j}r_2^{j,k} \tag{5} \]

假設 \(S_1\)、\(Z_1\) 是 \(r_1\) 矩陣對應的 scale 和 zero point，\(S_2\)、\(Z_2\)、\(S_3\)、\(Z_3\)同理，那么由 (5) 式可以推出：

\[S_3(q_3^{i,k}-Z_3)=\sum_{j=1}^{N}S_1(q_{1}^{i,j}-Z_1)S_2(q_2^{j,k}-Z_2) \tag{6} \]

整理一下可以得到：

\[q_3^{i,k}=\frac{S_1 S_2}{S_3}\sum_{j=1}^N(q_1^{i,j}-Z_1)(q_2^{j,k}-Z_2)+Z_3 \tag{7} \]

仔細觀察 (7) 式可以發現，除了\(\frac{S_1 S_2}{S_3}\)，其他都是定點整數運算。那如何把 \(\frac{S_1 S_2}{S_3}\) 也變成定點運算呢？這里要用到一個 trick。假設 \(M=\frac{S_1 S_2}{S_3}\)，由於 \(M\) 通常都是 (0, 1) 之間的實數 (這是通過大量實驗統計出來的)，因此可以表示成 \(M=2^{-n}M_0\)，其中 \(M_0\) 是一個定點實數。注意，定點數並不一定是整數，所謂定點，指的是小數點的位置是固定的，即小數位數是固定的。因此，如果存在 \(M=2^{-n}M_0\)，那我們就可以通過\(M_0\)的 bit 位移操作實現 \(2^{-n}M_0\)，這樣整個過程就都在定點上計算了。

很多剛接觸量化的同學對這一點比較疑惑，下面我就用一個簡單的示例說明這一點。我們把 \(M=\frac{S_1 S_2}{S_3}\) 代入 (7) 式可以得到：

\[q_3^{i,k}=M\sum_{j=1}^N(q_1^{i,j}-Z_1)(q_2^{j,k}-Z_2)+Z_3=MP+Z_3 \tag{8} \]

這里面 \(P\) 是一個在定點域上計算好的整數。

假設 \(P=7091\)，\(M=0.0072474273418460\) (\(M\) 可以通過 \(S\) 事先計算得到)，那下面我們就是要找到一個 \(M_0\) 和 \(n\)，使得 \(MP=2^{-n}M_0 P\) 成立。我們可以用一段代碼來找到這兩個數：

M = 0.0072474273418460
P = 7091

def multiply(n, M, P):
    result = M * P
    Mo = int(round(2 ** n * M)) # 這里不一定要四舍五入截斷，因為python定點數不好表示才這樣處理

    approx_result = (Mo * P) >> n
    print("n=%d, Mo=%d, approx=%f, error=%f"%\
          (n, Mo, approx_result, result-approx_result))

for n in range(1, 16):
    multiply(n, M, P)

輸出：

n=1, Mo=0, approx=0.000000, error=51.391507
n=2, Mo=0, approx=0.000000, error=51.391507
n=3, Mo=0, approx=0.000000, error=51.391507
n=4, Mo=0, approx=0.000000, error=51.391507
n=5, Mo=0, approx=0.000000, error=51.391507
n=6, Mo=0, approx=0.000000, error=51.391507
n=7, Mo=1, approx=55.000000, error=-3.608493
n=8, Mo=2, approx=55.000000, error=-3.608493
n=9, Mo=4, approx=55.000000, error=-3.608493
n=10, Mo=7, approx=48.000000, error=3.391507
n=11, Mo=15, approx=51.000000, error=0.391507
n=12, Mo=30, approx=51.000000, error=0.391507
n=13, Mo=59, approx=51.000000, error=0.391507
n=14, Mo=119, approx=51.000000, error=0.391507
n=15, Mo=237, approx=51.000000, error=0.391507

可以看到，在 n=11、\(M_0=15\) 的時候，誤差就已經在 1 以內了。因此，只要 \(M_0P\) 的數值范圍在 21(32-11) 個 bit 內，就可以通過對 \(M_0P\) 右移 \(n\) 個 bit 來近似 \(MP\) 了，而這個誤差本身在可以接受的范圍內。這樣一來，(8) 式就可以完全通過定點運算來計算，即我們實現了浮點矩陣乘法的量化。

卷積網絡的量化

有了上面矩陣乘法的量化，我們就可以進一步嘗試對卷積網絡的量化。

假設一個這樣的網絡：

這個網絡只有三個模塊，現在需要把 conv、fc、relu 量化。

假設輸入為 \(x\)，我們可以事先統計樣本的最大值和最小值，然后計算出 \(S_x\)(scale) 和 \(Z_x\)(zero point)。

同樣地，假設 conv、fc 的參數為 \(w_1\)、\(w_2\)，以及 scale 和 zero point 為 \(S_{w1}\)、\(Z_{w1}\)、\(S_{w2}\)、\(Z_{w2}\)。中間層的 feature map 為 \(a_1\)，\(a_2\)，並且事先統計出它們的 scale 和 zero point 為 \(S_{a1}\)、\(Z_{a1}\)、\(S_{a2}\)、\(Z_{a2}\)。

卷積運算和全連接層的本質都是矩陣運算，因此我們可以把卷積運算表示成 (這里先忽略加 bias 的操作，這一步同樣可以量化，不過中間有一些 trick，我們在之后的文章再仔細研究)：

\[a_1^{i,k}=\sum_{j=1}^N x^{i,j}w_1^{j,k} \tag{9} \]

根據之前的轉換，我們可以得到：

\[q_{a1}^{i,k}=M\sum_{j=1}^N(q_x^{i,j}-Z_x)(q_{w1}^{j,k}-Z_{w1})+Z_{a1} \tag{10} \]

其中 \(M=\frac{S_{w1}S_{x}}{S_{a1}}\)。

得到 conv 的輸出后，我們不用反量化回 \(a_1\)，直接用 \(q_{a1}\) 繼續后面的計算即可。

對於量化的 relu 來說，計算公式不再是 \(q_{a2}=max(q_{a1}, 0)\)，而是 \(q_{a2}=max(q_{a1},Z_{a1})\)，並且 \(S_{a1}=S_{a2}\)，\(Z_{a1}=Z_{a2}\) (為什么是這樣，這一點留作思考題)。另外，在實際部署的時候，relu 或者 bn 通常是會整合到 conv 中一起計算的，這一點在之后的文章再講解。

得到 \(q_{a2}\) 后，我們可以繼續用 (8) 式來計算 fc 層。假設網絡輸出為 \(y\)，對應的 scale 和 zero point 為 \(S_y\)、\(Z_y\)，則量化后的 fc 層可以用如下公式計算：

\[q_{y}^{i,k}=M\sum_{j=1}^N(q_{a2}^{i,j}-Z_{a2})(q_{w2}^{j,k}-Z_{w2})+Z_{y}\tag{11} \]

然后通過公式 \(y=S_y(q_y-Z_y)\) 把結果反量化回去，就可以得到近似原來全精度模型的輸出了。

可以看到，上面整個流程都是用定點運算實現的。我們在得到全精度的模型后，可以事先統計出 weight 以及中間各個 feature map 的 min、max，並以此計算出 scale 和 zero point，然后把 weight 量化成 int8/int16 型的整數后，整個網絡便完成了量化，然后就可以依據上面的流程做量化推理了。

總結

這篇文章主要介紹了矩陣量化的原理，以及如何把矩陣量化運用到卷積網絡中，實現全量化網絡的計算。這中間忽略了很多細節，比如 relu 和 conv 的合並、激活函數的量化、量化訓練的流程等。后面的文章會繼續補充一些細節，並通過從零搭建一個 pytorch 的量化模型來幫助讀者更好地理解中間的過程。

參考

• 神經網絡量化簡介
• Building a quantization paradigm from first principles
• Post Training Quantization General Questions
• 量化訓練：Quantization Aware Training in Tensorflow（一）
• How to Quantize an MNIST network to 8 bits in Pytorch from scratch (No retraining required)
• Aggressive Quantization: How to run MNIST on a 4 bit Neural Net using Pytorch
• TensorFlow Lite 8-bit quantization specification
• Post-training quantization

PS: 之后的文章更多的會發布在公眾號上，歡迎有興趣的讀者關注我的個人公眾號：AI小男孩，掃描下方的二維碼即可關注

作為曾經在 AI 路上苦苦掙扎的過來人，想幫助小白們更好地思考各種 AI 技術的來龍去脈。