神经网络量化入门--Folding BN ReLU
上一篇文章介绍了量化训练的基本流程，本文介绍量化中如何把 BatchNorm 和 ReLU 合并到 Conv 中。
Folding BatchNorm
BatchNorm 是 Google 提出的一种加速神经网络训练的技术，在很多网络中基本是标配。
回忆一下，BatchNorm 其实就是在每一层输出的时候做了一遍归一化操作：
其中 x i x_i xi 是网络中间某一层的激活值， μ β \mu_{\beta} μβ、 σ β \sigma_{\beta} σβ 分别是其均值和方差， y i y_i yi 则是过了 BN 后的输出。
一般卷积层与BN合并
Folding BatchNorm 不是量化才有的操作，在一般的网络中，为了加速网络推理，我们也可以把 BN 合并到 Conv 中。
合并的过程是这样的，假设有一个已经训练好的 Conv 和 BN：
假设 Conv 的 weight 和 bias 分别是 w w w 和 b b b。那么卷积层的输出为：
y = ∑ i N w i x i + b (1) y=\sum_{i}^N w_i x_i + b \tag{1} y=i∑Nwixi+b(1)
图中 BN 层的均值和标准差可以表示为 μ y \mu_{y} μy、 σ y \sigma_{y} σy，那么根据论文的表述，BN 层的输出为：
KaTeX parse error: No such environment: align at position 8: \begin{̲a̲l̲i̲g̲n̲}̲ y_{bn}&=\gamma…
然后我们把 (1) 代入 (2) 中可以得到：
y b n = γ σ y 2 + ϵ ( ∑ i N w i x i + b − μ y ) + β (3) y_{bn}=\frac{\gamma}{\sqrt{\sigma_y^2+\epsilon}}(\sum_{i}^N w_i x_i + b-\mu_y)+\beta \tag{3} ybn=σy2+ϵγ(i∑Nwixi+b−μy)+β(3)
我们用 γ ′ \gamma' γ′ 来表示 γ σ y 2 + ϵ \frac{\gamma}{\sqrt{\sigma_y^2+\epsilon}} σy2+ϵγ，那么 (3) 可以简化为：
y b n = γ ′ ( ∑ i N w i x i + b − μ y ) + β = ∑ i N γ ′ w i x i + γ ′ ( b − μ y ) + β (4) \begin{aligned} y_{bn}&=\gamma'(\sum_{i}^N w_ix_i+b-\mu_y)+\beta \\ &=\sum_{i}^N \gamma'w_ix_i+\gamma'(b-\mu_y)+\beta \tag{4} \end{aligned} ybn=γ′(i∑Nwixi+b−μy)+β=i∑Nγ′wixi+γ′(b−μy)+β(4)
发现没有，(4) 式形式上跟 (1) 式一模一样，因此它本质上也是一个 Conv 运算，我们只需要用 w i ′ = γ ′ w i w_i'=\gamma'w_i wi′=γ′wi 和 b ′ = γ ′ ( b − μ y ) + β b'=\gamma'(b-\mu_y)+\beta b′=γ′(b−μy)+β 来作为原来卷积的 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 i n ∈ [ − 1 , 1 ] r_{in} \in [-1, 1] rin∈[−1,1]，那么经过 ReLU 后的数值范围是 r o u t ∈ [ 0 , 1 ] r_{out} \in [0,1] rout∈[0,1]。假设量化到 uint8 类型，即 [0, 255]，那么 ReLU 前后的 scale 分别为 S i n = 2 255 S_{in}=\frac{2}{255} Sin=2552、 S o u t = 1 255 S_{out}=\frac{1}{255} Sout=2551，zp 分别为 Z i n = 128 Z_{in}=128 Zin=128、 Z o u t = 0 Z_{out}=0 Zout=0。 再假设 ReLU 前的浮点数是 r i n = 0.5 r_{in}=0.5 rin=0.5，那么经过 ReLU 后的值依然是 0.5。换算成整型的话，ReLU 前的整数是 q i n = 192 q_{in}=192 qin=192，由于 Z i n = 128 Z_{in}=128 Zin=128，因此过完 ReLU 后的数值依然是 192。但是， S o u t S_{out} Sout 和 Z o u t Z_{out} Zout 已经发生了变化，因此反量化后的 r o u t r_{out} rout 不再是 0.5，而这不是我们想要的。所以，如果想要保证量化的 ReLU 和浮点型的 ReLU 之间的一致性，就必须保证 S i n S_{in} Sin、 S o u t S_{out} Sout 以及 Z i n Z_{in} Zin、 Z o u t Z_{out} Zout 是一致的。
但是保证前后的 scale 和 zp 一致，没规定一定得用 S i n S_{in} Sin 和 Z i n Z_{in} Zin，我们一样可以用 ReLU 之后的 scale 和 zp。不过，使用哪一个 scale 和 zp，意义完全不一样。如果使用 ReLU 之后的 scale 和 zp，那我们就可以用量化本身的截断功能来实现 ReLU 的作用。
想要理解这一点，需要回顾一下量化的基本公式：
q = r o u n d ( r S + Z ) (5) q=round(\frac{r}{S}+Z) \tag{5} q=round(Sr+Z)(5)
注意，这里的 round 除了把 float 型四舍五入转成 int 型外，还需要保证 q q q 的数值在特定范围内「例如 0～255」，相当于要做一遍 clip 操作。因此，这个公式更准确的写法应该是「假设量化到 uint8 数值」：
q = r o u n d ( c l i p ( r S + Z , 0 , 255 ) ) (6) q=round(clip(\frac{r}{S}+Z, 0, 255)) \tag{6} q=round(clip(Sr+Z,0,255))(6)
记住，ReLU 本身就是在做 clip。所以，我们才能用量化的截断功能来模拟 ReLU 的功能。
再举个例子。
假设有一个上图所示的 Conv+ReLU 的结构，其中，Conv 后的数值范围是 r i n ∈ [ − 1 , 1 ] r_{in} \in [-1,1] rin∈[−1,1]。在前面的文章中，我们都是用 ReLU 前的数值来统计 minmax 并计算 scale 和 zp，并把该 scale 和 zp 沿用到 ReLU 之后。这部分的计算可以参照图中上半部分。
但现在，我们想在 ReLU 之后统计 minmax，并用 ReLU 后的 scale 和 zp 作为 ReLU 前的 scale 和 zp「即 Conv 后面的 scale 和 zp」，结果会怎样呢？
看图中下半部分，假设 Conv 后的数值是 r i n = − 0.5 r_{in}=-0.5 rin=−0.5，此时，由于 Conv 之后的 scale 和 zp 变成了 1 255 \frac{1}{255} 2551 和 0 0 0，因此，量化的整型数值为：
q = r o u n d ( − 0.5 1 255 + 0 ) = r o u n d ( − 128 ) = 0 (7) \begin{aligned} q&=round(\frac{-0.5}{\frac{1}{255}}+0) \\ &=round(-128) \\ &=0 \tag{7} \end{aligned} q=round(2551−0.5+0)=round(−128)=0(7)
注意，上面的量化过程中，我们执行了截断操作，把 q q q 从 -128 截断成 0，而这一步本来应该是在 ReLU 里面计算的！然后，我们如果根据 S o u t S_{out} Sout 和 Z o u t Z_{out} Zout 反量化回去，就会得到 r o u t = 0 r_{out}=0 rout=0，而它正是原先 ReLU 计算后得到的数值。
因此，通过在 Conv 后直接使用 ReLU 后的 scale 和 zp，我们实现了将 ReLU 合并到 Conv 里面的过程。
那对于 ReLU 外的其他激活函数，是否可以同样合并到 Conv 里面呢？这取决于其他函数是否也只是在做 clip 操作，例如 ReLU6 也有同样的性质。但对于其他绝大部分函数来说，由于它们本身包含其他数学运算，因此就不具备类似性质。
总结
这篇文章主要介绍了如何把 BatchNorm 和 ReLU 合并成一个 Conv，从而加速量化推理。按照计划，应该和之前的文章一样，给出代码实现。但我在测试代码的时候发现有一些 bug 需要解决，正好也控制一下篇幅，下篇文章会给出相关的代码实现。
欢迎关注我的公众号：大白话AI，立志用大白话讲懂AI。