好久不见,优化系列得继续更啊。。。
前言
今天来介绍一个工程上的常用Trick,即折叠Batch Normalization,也叫作折叠BN。我们知道一般BN是跟在卷积层后面,一般还会接上激活函数,也就是conv+BN+relu
这种基本组件,但在部署的时候前向推理框架一般都会自动的将BN和它前面的卷积层折叠在一起,实现高效的前向推理网络,今天我就从原理和实现来讲一讲这个操作。
原理
我们知道卷积层的计算可以表示为:
然后BN层的计算可以表示为:
我们把二者组合一下,公式如下:
然后令
那么,合并BN层后的卷积层的权重和偏置可以表示为:
值得一提的是,一般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
(完)