神經網絡量化入門--Folding BatchNorm ReLU

上一篇文章介紹了量化訓練的基本流程，本文介紹量化中如何把 BatchNorm 和 ReLU 合並到 Conv 中。

Folding BatchNorm

BatchNorm 是 Google 提出的一種加速神經網絡訓練的技術，在很多網絡中基本是標配。

回憶一下，BatchNorm 其實就是在每一層輸出的時候做了一遍歸一化操作：

其中 \(x_i\) 是網絡中間某一層的激活值，\(\mu_{\beta}\)、\(\sigma_{\beta}\) 分別是其均值和方差，\(y_i\) 則是過了 BN 后的輸出。

一般卷積層與BN合並

Folding BatchNorm 不是量化才有的操作，在一般的網絡中，為了加速網絡推理，我們也可以把 BN 合並到 Conv 中。

合並的過程是這樣的，假設有一個已經訓練好的 Conv 和 BN：

假設 Conv 的 weight 和 bias 分別是 \(w\) 和 \(b\)。那么卷積層的輸出為：

\[y=\sum_{i}^N w_i x_i + b \tag{1} \]

圖中 BN 層的均值和標准差可以表示為 \(\mu_{y}\)、\(\sigma_{y}\)，那么根據論文的表述，BN 層的輸出為：

\[\begin{align} y_{bn}&=\gamma \hat{y}+\beta \notag \\ &=\gamma \frac{y-\mu_y}{\sqrt{\sigma_y^2+\epsilon}}+\beta \tag{2} \end{align} \]

然后我們把 (1) 代入 (2) 中可以得到：

\[y_{bn}=\frac{\gamma}{\sqrt{\sigma_y^2+\epsilon}}(\sum_{i}^N w_i x_i + b-\mu_y)+\beta \tag{3} \]

我們用 \(\gamma'\) 來表示 \(\frac{\gamma}{\sqrt{\sigma_y^2+\epsilon}}\)，那么 (3) 可以簡化為：

\[\begin{align} y_{bn}&=\gamma'(\sum_{i}^Nw_ix_i+b-\mu_y)+\beta \notag \\ &=\sum_{i}^N \gamma'w_ix_i+\gamma'(b-\mu_y)+\beta \tag{4} \end{align} \]

發現沒有，(4) 式形式上跟 (1) 式一模一樣，因此它本質上也是一個 Conv 運算，我們只需要用 \(w_i'=\gamma'w_i\) 和 \(b'=\gamma'(b-\mu_y)+\beta\) 來作為原來卷積的 weight 和 bias，就相當於把 BN 的操作合並到了 Conv 里面。實際 inference 的時候，由於 BN 層的參數已經固定了，因此可以把 BN 層 folding 到 Conv 里面，省去 BN 層的計算開銷。

量化 BatchNorm Folding

量化網絡時可以用同樣的方法把 BN 合並到 Conv 中。

如果量化時不想更新 BN 的參數 (比如后訓練量化)，那我們就先把 BN 合並到 Conv 中，直接量化新的 Conv 即可。

如果量化時需要更新 BN 的參數 (比如量化感知訓練)，那也很好處理。Google 把這個流程的心法寫在一張圖上了：

由於實際 inference 的時候，BN 是 folding 到 Conv 中的，因此在量化訓練的時候也需要模擬這個操作，得到新的 weight 和 bias，並用新的 Conv 估計量化誤差來回傳梯度。

Conv與ReLU合並

在量化中，Conv + ReLU 這樣的結構一般也是合並成一個 Conv 進行運算的，而這一點在全精度模型中則辦不到。

在之前的文章中說過，ReLU 前后應該使用同一個 scale 和 zeropoint。這是因為 ReLU 本身沒有做任何的數學運算，只是一個截斷函數，如果使用不同的 scale 和 zeropoint，會導致無法量化回 float 域。

看下圖這個例子。假設 ReLU 前的數值范圍是 \(r_{in} \in [-1, 1]\)，那么經過 ReLU 后的數值范圍是 \(r_{out} \in [0,1]\)。假設量化到 uint8 類型，即 [0, 255]，那么 ReLU 前后的 scale 分別為 \(S_{in}=\frac{2}{255}\)、\(S_{out}=\frac{1}{255}\)，zp 分別為 \(Z_{in}=128\)、\(Z_{out}=0\)。 再假設 ReLU 前的浮點數是 \(r_{in}=0.5\)，那么經過 ReLU 后的值依然是 0.5。換算成整型的話，ReLU 前的整數是 \(q_{in}=192\)，由於 \(Z_{in}=128\)，因此過完 ReLU 后的數值依然是 192。但是，\(S_{out}\) 和 \(Z_{out}\) 已經發生了變化，因此反量化后的 \(r_{out}\) 不再是 0.5，而這不是我們想要的。所以，如果想要保證量化的 ReLU 和浮點型的 ReLU 之間的一致性，就必須保證 \(S_{in}\)、\(S_{out}\) 以及 \(Z_{in}\)、\(Z_{out}\) 是一致的。

但是保證前后的 scale 和 zp 一致，沒規定一定得用 \(S_{in}\) 和 \(Z_{in}\)，我們一樣可以用 ReLU 之后的 scale 和 zp。不過，使用哪一個 scale 和 zp，意義完全不一樣。如果使用 ReLU 之后的 scale 和 zp，那我們就可以用量化本身的截斷功能來實現 ReLU 的作用。

想要理解這一點，需要回顧一下量化的基本公式：

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

注意，這里的 round 除了把 float 型四舍五入轉成 int 型外，還需要保證 \(q\) 的數值在特定范圍內「例如 0～255」，相當於要做一遍 clip 操作。因此，這個公式更准確的寫法應該是「假設量化到 uint8 數值」：

\[q=round(clip(\frac{r}{S}+Z, 0, 255)) \tag{6} \]

記住，ReLU 本身就是在做 clip。所以，我們才能用量化的截斷功能來模擬 ReLU 的功能。

再舉個例子。

假設有一個上圖所示的 Conv+ReLU 的結構，其中，Conv 后的數值范圍是 \(r_{in} \in [-1,1]\)。在前面的文章中，我們都是用 ReLU 前的數值來統計 minmax 並計算 scale 和 zp，並把該 scale 和 zp 沿用到 ReLU 之后。這部分的計算可以參照圖中上半部分。

但現在，我們想在 ReLU 之后統計 minmax，並用 ReLU 后的 scale 和 zp 作為 ReLU 前的 scale 和 zp「即 Conv 后面的 scale 和 zp」，結果會怎樣呢？

看圖中下半部分，假設 Conv 后的數值是 \(r_{in}=-0.5\)，此時，由於 Conv 之后的 scale 和 zp 變成了 \(\frac{1}{255}\) 和 \(0\)，因此，量化的整型數值為：

\[\begin{align} q&=round(\frac{-0.5}{\frac{1}{255}}+0) \notag \\ &=round(-128) \notag \\ &=0 \tag{7} \end{align} \]

注意，上面的量化過程中，我們執行了截斷操作，把 \(q\) 從 -128 截斷成 0，而這一步本來應該是在 ReLU 里面計算的！然后，我們如果根據 \(S_{out}\) 和 \(Z_{out}\) 反量化回去，就會得到 \(r_{out}=0\)，而它正是原先 ReLU 計算后得到的數值。

因此，通過在 Conv 后直接使用 ReLU 后的 scale 和 zp，我們實現了將 ReLU 合並到 Conv 里面的過程。

那對於 ReLU 外的其他激活函數，是否可以同樣合並到 Conv 里面呢？這取決於其他函數是否也只是在做 clip 操作，例如 ReLU6 也有同樣的性質。但對於其他絕大部分函數來說，由於它們本身包含其他數學運算，因此就不具備類似性質。

總結

這篇文章主要介紹了如何把 BatchNorm 和 ReLU 合並成一個 Conv，從而加速量化推理。按照計划，應該和之前的文章一樣，給出代碼實現。但我在測試代碼的時候發現有一些 bug 需要解決，正好也控制一下篇幅，下篇文章會給出相關的代碼實現。

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

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