Conv、BN、Relu融合

神经网络量化入门–Folding BN ReLU

上一篇文章介绍了量化训练的基本流程，本文介绍量化中如何把BatchNorm和ReLU合并到Conv中。

Folding BatchNorm

BatchNorm是Google提出的一种加速神经网络训练的技术，在很多网络中基本是标配。

回忆一下，BatchNorm其实就是在每一层输出的时候做了一遍归一化操作：
Input: Values of x over a mini-batch: B = { x 1 . . x m } B = \{x_1 ..x_m\} B={x1​..xm​} ; Parameters to be learned: $\gamma$, $\beta$
Output: { y i = B N B ( x i ) } \{y_i = BN_B(x_i)\} {yi​=BNB​(xi​)}

${\mu }_{\mathcal{B}}\leftarrow \frac{1}{m}{\sum }_{i=1}^m{x}_i$ // mini-batch mean
${\sigma }_{\mathcal{B}}^2\leftarrow \frac{1}{m}{\sum }_{i=1}^m(x_i-{\mu }_{\mathcal{B}}{)}^2$ // mini-batch variance
${\hat{x}}_i\leftarrow \frac{x_i-{\mu }_{\mathcal{B}}}{\sqrt{{\sigma }_{\mathcal{B}}^2+ϵ}}$ // normalize
$y_i\leftarrow \gamma {\hat{x}}_i+\beta \equiv B{N}_{\gamma ,\beta }(x_i)$ // scale and shift

Algorithm 1: Batch Normalizing Transform, applied to activation x over a mini-batch.

其中$x_i$是网络中间某一层的激活值，${\mu }_{\mathcal{B}}$、${\sigma }_{\mathcal{B}}^2$分别是其均值和方差，$y_i$则是过了BN后的输出。

一般卷积层与BN合并

Folding BatchNorm不是量化才有的操作，在一般的网络中，为了加速网络推理，我们也可以把BN合并到Conv中。

合并的过程是这样的。假设有一个已经训练好的Conv和BN：

假设Conv的weight和bias分别是$W$和$b$（实际上卷积的kernal就是weight，当weight和图片fold后的做完einsum之后再加上bias就是输出）。那么卷积层的输出为：
$$y=\sum _{i=1}^N{w}_i{x}_i+b\text{\hspace{1em}(1)}$$
图中 BN 层的均值和标准差可以表示为 ${\mu }_y$，${\sigma }_y$，那么根据论文的表述，BN 层的输出为：

$$\begin{gathered} y_{bn}=\gamma \hat{y}+\beta \\ =\gamma \frac{y-{\mu }_y}{\sqrt{{\sigma }_y^2+ϵ}}+\beta \text{\hspace{1em}(2)} \end{gathered}$$

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

$$y_{bn}=\frac{\gamma }{\sqrt{{\sigma }_y^2+ϵ}}\left(\sum _{i=1}^N{w}_i{x}_i+b-{\mu }_y\right)+\beta \text{\hspace{1em}(3)}$$

我们用 ${\gamma }^{\prime }$ 来表示 $\frac{\gamma }{\sqrt{{\sigma }_y^2+ϵ}}$，那么 (3) 可以简化为：

$$\begin{gathered} y_{bn}={\gamma }^{\prime }\left(\sum _{i=1}^N{w}_i{x}_i+b-{\mu }_y\right)+\beta \\ =\sum _{i=1}^N{\gamma }^{\prime }{w}_i{x}_i+{\gamma }^{\prime }(b-{\mu }_y)+\beta \text{\hspace{1em}(4)} \end{gathered}$$

发现没有，(4) 式形式上跟 (1) 式一样一样，因此它本质上也是一个 Conv 运算，我们只需要用 $w_i^{\prime }={\gamma }^{\prime }{w}_i$ 和 $b^{\prime }={\gamma }^{\prime }(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{P}_{in}=128$、$Z{P}_{out}=0$。再假设ReLU前的浮点数是$r_{in}=0.5$，那么经过ReLU后的值依然是$0.5$。换算成整型的话，ReLU前的整数是$q_{in}=192$，由于$q_{in}\in [0,255]$，因此过完ReLU后的数值依然是$192$。但是，$S_{in}$和$Z{P}_{in}$已经发生了变化，因此反量化后的$r_{in}$不再是$0.5$，而这不是我们想要的。所以，如果想要保证量化的ReLU和浮点型的ReLU之间的一致性，就必须保证$S_{in}$、$S_{out}$以及$Z{P}_{in}$、$Z{P}_{out}$是一致的。

但是保证前后的scale和zp一致，没规定一定得用$S_{in}$和$Z{P}_{in}$，我们一样可以用ReLU之后的scale和zp。不过，使用哪一个scale和zp，意义完全不一样。如果使用ReLU之后的scale和zp，那我们就可以用量化本身的截断功能来实现ReLU的作用。

想要理解这一点，需要回顾一下量化的基本公式：
注意，这里的round除了把float型四舍五入转成int型外，还需要保证$q$的数值在特定范围内「例如$0～255$」，相当于要做一遍clip操作。因此，这个公式更准确的写法应该是「假设量化到uint8数值」：
而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变成了$S_{out}=\frac{1}{255}$和$Z{P}_{out}=0$，因此，量化的整型数值为：
注意，上面的量化过程中，我们执行了截断操作，把$q_{in}$从$-128$截断成$0$，而这一步本来应该是在ReLU里面计算的！然后，我们如果根据$S_{out}$和$Z{P}_{out}$反量化回去，就会得到$r_{out}=0$，而它正是原先ReLU计算后得到的数值。

因此，通过在Conv后直接使用ReLU后的scale和zp，我们实现了将ReLU合并到Conv里面的过程。

那对于ReLU外的其他激活函数，是否可以同样合并到Conv里面呢？这取决于其他函数是否也只是在做clip操作，例如ReLU6也有同样的性质。但对于其他绝大部分函数来说，由于它们本身包含其他数学运算，因此就不具备类似性质。

总结

这篇文章主要介绍了如何把BatchNorm和ReLU合并成一个Conv，从而加速量化推理。按照计划，应该和之前的文章一样，给出代码实现。但我在测试代码的时候发现有一些bug需要解决，正好也控制一下篇幅，下篇文章会给出相关的代码实现。

转载

• https://mp.weixin.qq.com/s?__biz=Mzg4ODA3MDkyMA==&mid=2247483746&idx=1&sn=18eebe5df9d43c487adbdf7da25d11ca&chksm=ceceadf868049d48123510f37841fac69b7355d2e606bef9c6363e62fdf3004212aeee49533a#rd
• https://mp.weixin.qq.com/s?__biz=Mzg4ODA3MDkyMA==&mid=2247483692&idx=1&sn=3e28db4881d591f4e6a66c83d4213823&chksm=cf81f74bf8f67e5d0f2a98fd7bf7a91864d14010d88a5ed89120b7b4fcd94fc34789f0d0db9a&scene=21#wechat_redirect
• https://mp.weixin.qq.com/s?__biz=Mzg4ODA3MDkyMA==&mid=2247483723&idx=1&sn=ad844b16f6d7a7eb006a01fa6201c9c7&chksm=cf81f72cf8f67e3a1349d53518187652d84328117f19ec9a3fa19b693ffe4e8f43142f1c8740&scene=21#wechat_redirect