前言¶

今天来介绍一个工程上的常用Trick，即折叠Batch Normalization，也叫作折叠BN。我们知道一般BN是跟在卷积层后面，一般还会接上激活函数，也就是conv+BN+relu这种基本组件，但在部署的时候前向推理框架一般都会自动的将BN和它前面的卷积层折叠在一起，实现高效的前向推理网络，今天我就从原理和实现来讲一讲这个操作。

原理¶

我们知道卷积层的计算可以表示为：

$Y = W * X + B$

然后BN层的计算可以表示为：

$\mu = \frac{1}{m}\sum_{i=1}^mx_i$

$\sigma^2=\frac{1}{m}\sum_{i=1}^m(x_i-\mu)^2$

$x_i=\frac{x_i-\mu}{\sqrt{\sigma^2+ \epsilon}}$

$y_i=\gamma * x_i + \beta$

我们把二者组合一下，公式如下：

$Y=\gamma*(\frac{(W*X+B)-\mu}{\sqrt{\sigma^2+\epsilon}})+\beta$

$Y=\frac{\gamma*W}{\sqrt{\sigma^2+\epsilon}}*X+\frac{\gamma*(B-\mu)}{\sqrt{\sigma^2+\epsilon}}+\beta$

然后令 $a = \frac{\gamma}{\sqrt{\delta^2+\epsilon}}$

那么，合并BN层后的卷积层的权重和偏置可以表示为：

$W_{merged}=W*a$

$B_{merged}=(B-\mu)*a+\beta$

值得一提的是，一般Conv后面接BN的时候很多情况下是不带Bias的，这个时候上面的公式就会少第二项。

Pytorch代码实现¶

原理介绍完了，接下来我们来看一下代码实现，就以Pytorch的代码为例好了。

def fuse(conv, bn):
    global i
    i = i + 1
    # ********************BN参数*********************
    mean = bn.running_mean
    var_sqrt = torch.sqrt(bn.running_var + bn.eps)
    gamma = bn.weight
    beta = bn.bias
    # *******************conv参数********************
    w = conv.weight
    if conv.bias is not None:
        b = conv.bias
    else:
        b = mean.new_zeros(mean.shape)

    if(i >= 2 and i <= 7):
        b = b - mean + beta * var_sqrt
    else:
        w = w * (gamma / var_sqrt).reshape([conv.out_channels, 1, 1, 1])
        b = (b - mean)/var_sqrt * gamma + beta
    fused_conv = nn.Conv2d(conv.in_channels,
                         conv.out_channels,
                         conv.kernel_size,
                         conv.stride,
                         conv.padding,
                         groups=conv.groups,
                         bias=True)
    fused_conv.weight = nn.Parameter(w)
    fused_conv.bias = nn.Parameter(b)
    return fused_conv

可以看到这个代码实现和上面介绍的公式完全对应，实现了折叠BN的效果。

Caffe代码实现¶

下面这个代码实现了将Mobilenet Caffe模型的卷积和BN进行合并的代码，亲测无误，放心使用。

由于完整代码太长，我放这个工程里了：https://github.com/BBuf/cv_tools/blob/master/merge_bn.py

思考¶

从公式推导可知，只有Conv后接BN才可以折叠，Preact结构的ResNet不可以折叠BN。
如果是反卷积接的BN，如何合并？

效果¶

在参考2的博客中看到看到了折叠BN的速度测试结果，分别在1080Ti和i7 CPU上对比了Resnet50 模型合并BN层前后的性能，可以看到分类精度保持不变，速度显著提升，具体测试结果如下：

模型	CPU前向时间	GPU前向时间
Resnet50（合并前）	176.17ms	11.03ms
Resnet50（合并后）	161.69ms	7.3ms
提升	10%	51%

附录¶

参考1：https://github.com/666DZY666/model-compression
参考2：https://blog.csdn.net/kangdi7547/article/details/81348254

前言¶

原理¶

Pytorch代码实现¶

Caffe代码实现¶

思考¶

效果¶

附录¶

推荐阅读¶