[TOC]

概述

文章参考：https://zhuanlan.zhihu.com/p/176982058

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

Folding BatchNorm

康奈尔大学的论文：https://arxiv.org/abs/1502.03167

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 x_iw_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}
$$
均值计算：mean，数列的算术平均值，反应了数列的集中趋势,等于有效数值的合除以有效数值的个数。

标准差计算：https://zhuanlan.zhihu.com/p/35435231

然后我们把 (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}$ 和，我们一样可以用 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&=clip(round(\frac{-0.5}{\frac{1}{255}}+0), 0, 255) \notag \ &=clip(round(-128), 0, 255) \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 需要解决，正好也控制一下篇幅，下篇文章会给出相关的代码实现。