工程Trick | 合并BN层加速前向推理

1.BatchNorm

在训练深度网络模型时，BN（Batch Normalization）层能够加速网络收敛，并且能够控制过拟合，一般放在卷积层之后。如下示意图所示，BatchNorm是以通道为单位，对当前通道中所有的N、H、W做归一化。

BN 层将特征归一化后，能够有效解决由于特征分布不均匀而导致的梯度消失与梯度爆炸问题。并通过可学习参数保证特征的有效性。虽然 BN 层在训练时起到了积极作用，然而，在网络前向推断时多了一些层的运算，影响了模型的性能，且占用了更多的内存或者显存空间。基于此将卷积后的BN操作直接融合到卷积中，这是现在工程上普遍会使用的方法。

2.将BN操作合并到卷积中

3.pytorch代码实现

很明显，下述代码就是将上部分中的公式进行了代码化。在这里不得不说一句pytorch大法NB。

conv_param = model_state_dict[merge_bn_model_dict_key[index]]
conv_bias = model_state_dict[merge_bn_model_dict_key[index+1]]
outchannel = conv_param.size()[1]
bn_weignt = model_state_dict[merge_bn_model_dict_key[index+2]]
bn_bias = model_state_dict[merge_bn_model_dict_key[index+3]]
bn_mean = model_state_dict[merge_bn_model_dict_key[index+4]]
bn_var = model_state_dict[merge_bn_model_dict_key[index+5]]

#merge
var_sqrt = torch.sqrt(bn_var+bn_eps)
merge_bn_conv_param = conv_param*((bn_weight/var_sqrt)
merge_bn_conv_bias = bn_bias + bn_weight*(conv_bias-bn_mean/var_sqrt)

new_conv_key = merge_bn_model_dict_key[index]
new_bias_key = new_conv_key[:,-6]+'bias'


merge_bn_model_dict[new_conv_key] = merge_bn_conv_param
merge_bn_model_dict[new_bias_key] = merge_bn_conv_bias