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

Folding BatchNorm

BatchNorm 是 Google 提出的一種加速神經(jīng)網(wǎng)絡(luò)訓(xùn)練的技術(shù)，在很多網(wǎng)絡(luò)中基本是標(biāo)配。

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

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

一般卷積層與BN合并

Folding BatchNorm 不是量化才有的操作，在一般的網(wǎng)絡(luò)中，為了加速網(wǎng)絡(luò)推理，我們也可以把 BN 合并到 Conv 中。

合并的過程是這樣的，假設(shè)有一個已經(jīng)訓(xùn)練好的 Conv 和 BN：

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

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

圖中 BN 層的均值和標(biāo)準(zhǔn)差可以表示為 (mu_{y})、(sigma_{y})，那么根據(jù)論文的表述，BN 層的輸出為：

[egin{align}
y_{bn}&=gamma hat{y}+eta
otag \
&=gamma frac{y-mu_y}{sqrt{sigma_y^2+epsilon}}+eta ag{2}
end{align}
]

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

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

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

[egin{align}
y_{bn}&=gamma'(sum_{i}^Nw_ix_i+b-mu_y)+eta
otag \
&=sum_{i}^N gamma'w_ix_i+gamma'(b-mu_y)+eta ag{4}
end{align}
]

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

量化 BatchNorm Folding

量化網(wǎng)絡(luò)時可以用同樣的方法把 BN 合并到 Conv 中。

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

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

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

Conv與ReLU合并

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

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

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

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

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

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

注意，這里的 round 除了把 float 型四舍五入轉(zhuǎn)成 int 型外，還需要保證 (q) 的數(shù)值在特定范圍內(nèi)「例如 0～255」，相當(dāng)于要做一遍 clip 操作。因此，這個公式更準(zhǔn)確的寫法應(yīng)該是「假設(shè)量化到 uint8 數(shù)值」：

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

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

再舉個例子。

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

但現(xiàn)在，我們想在 ReLU 之后統(tǒng)計 minmax，并用 ReLU 后的 scale 和 zp 作為 ReLU 前的 scale 和 zp「即 Conv 后面的 scale 和 zp」，結(jié)果會怎樣呢？

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

[egin{align}
q&=round(frac{-0.5}{frac{1}{255}}+0)
otag \
&=round(-128)
otag \
&=0 ag{7}
end{align}
]

注意，上面的量化過程中，我們執(zhí)行了截斷操作，把 (q) 從 -128 截斷成 0，而這一步本來應(yīng)該是在 ReLU 里面計算的！然后，我們?nèi)绻鶕?jù) (S_{out}) 和 (Z_{out}) 反量化回去，就會得到 (r_{out}=0)，而它正是原先 ReLU 計算后得到的數(shù)值。

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

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

總結(jié)

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

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

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

總結(jié)

以上是生活随笔為你收集整理的神经网络量化入门--Folding BatchNorm ReLU的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。