好久不见,优化系列得继续更啊。。。
今天来介绍一个工程上的常用Trick,即折叠Batch Normalization,也叫作折叠BN。我们知道一般BN是跟在卷积层后面,一般还会接上激活函数,也就是conv+BN+relu
这种基本组件,但在部署的时候前向推理框架一般都会自动的将BN和它前面的卷积层折叠在一起,实现高效的前向推理网络,今天我就从原理和实现来讲一讲这个操作。
我们知道卷积层的计算可以表示为:
然后BN层的计算可以表示为:
我们把二者组合一下,公式如下:
然后令
那么,合并BN层后的卷积层的权重和偏置可以表示为:
值得一提的是,一般Conv后面接BN的时候很多情况下是不带Bias的,这个时候上面的公式就会少第二项。
原理介绍完了,接下来我们来看一下代码实现,就以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的效果。
下面这个代码实现了将Mobilenet Caffe模型的卷积和BN进行合并的代码,亲测无误,放心使用。
由于完整代码太长,我放这个工程里了:https://github.com/BBuf/cv_tools/blob/master/merge_bn.py
在参考2的博客中看到看到了折叠BN的速度测试结果,分别在1080Ti和i7 CPU上对比了Resnet50 模型合并BN层前后的性能,可以看到分类精度保持不变,速度显著提升,具体测试结果如下:
模型 | CPU前向时间 | GPU前向时间 |
---|---|---|
Resnet50(合并前) | 176.17ms | 11.03ms |
Resnet50(合并后) | 161.69ms | 7.3ms |
提升 | 10% | 51% |
本文分享自 GiantPandaCV 微信公众号,前往查看
如有侵权,请联系 cloudcommunity@tencent.com 删除。
本文参与 腾讯云自媒体同步曝光计划 ,欢迎热爱写作的你一起参与!