神经网络量化入门--Folding BN ReLU代码实现

AI小男孩

已于 2022-03-15 15:32:17 修改

阅读量1.9k

点赞数 3

分类专栏：深度学习文章标签：神经网络 pytorch 深度学习

于 2022-03-13 21:50:38 首次发布

本文链接：https://blog.csdn.net/u012954083/article/details/123445553

版权

本文介绍了如何在PyTorch中实现神经网络的BN与ReLU折叠，详细解析了Folding BN的公式与代码实现，以及合并ReLU的统计minmax方法。通过实验表明，量化感知训练在BN存在时仍能提升低比特精度，但后训练量化精度下降明显。

摘要由CSDN通过智能技术生成

(本文首发于公众号)

上一篇文章介绍了如何把 BatchNorm 和 ReLU 合并到 Conv 中，这篇文章会介绍具体的代码实现。本文相关代码都可以在 github 上找到。

Folding BN

回顾一下前文把 BN 合并到 Conv 中的公式：
$KaTeX parse error: No such environment: align at position 8: \begin{̲a̲l̲i̲g̲n̲}̲ y_{bn}&=\frac{…$
其中， $x$ 是卷积层的输入， $w$ 、 $b$ 分别是 Conv 的参数 weight 和 bias， $\gamma$ 、 $\beta$ 是 BN 层的参数。

对于 BN 的合并，首先，我们需要熟悉 pytorch 中的 BatchNorm2d 模块。

pytorch 中的 BatchNorm2d 针对 feature map 的每一个 channel 都会计算一个均值和方差，所以公式 (1) 需要对 weight 和 bias 进行 channel wise 的计算。另外，BatchNorm2d 中有一个布尔变量 affine，当该变量为 true 的时候，(1) 式中的 $\gamma$ 和 $\beta$ 就是可学习的， BatchNorm2d 会中有两个变量：weight 和 bias，来分别存放这两个参数。而当 affine 为 false 的时候，就直接默认 $\gamma=1$ ， $\beta=0$ ，相当于 BN 中没有可学习的参数。默认情况下，我们都设置 affine=True。

我们沿用之前的代码，先定义一个 QConvBNReLU 模块：

class QConvBNReLU(QModule):

    def __init__(self, conv_module, bn_module, qi=True, qo=True, num_bits=8):
        super(QConvBNReLU, self).__init__(qi=qi, qo=qo, num_bits=num_bits)
        self.num_bits = num_bits
        self.conv_module = conv_module
        self