必看必会的模型推理加速技巧：融合BN和Conv层

编辑：我是小将  

很多深度模型采用BN层（Batch Normalization）被很多深度模型来提升泛化能力。在模型推理时，BN层要从训练状态切换到测试状态，此时采用模型训练中近似的均值和方差。BN层最酷的地方是它可以用一个1x1卷积等效替换，更进一步地，我们可以将BN层合并到前面的卷积层中。

Batch Normalization

这里假定是网络的某一激活层特征，我们要对其进行归一化。若模型训练batch中共有个样例，其特征分别是 ，我们采用下列公式进行归一化：

而参数和是和其它模型参数一起通过梯度下降方法训练得到。在测试阶段，我们不太会对一个batch图像进行预测，一般是对单张图像测试。因此，通过前面公式计算 and 就不可能。BN的解决方案是采用训练过程中指数移动平均值和。

目前，BN层大部分用于CNN网络中，此时上面的4个参数就是针对特征图各个通道的参数，这里我们记, , 以及为第个通道的参数。

融合方案

首先我们将测试阶段的BN层（一般称为frozen BN）等效替换为一个1x1卷积层：

对于一个形状为的特征图，归一化后的结果，可以如下计算：

这里的是特征图的各个空间位置，我们可以看到上述计算就是形式，其可以看成是卷积，这里W的行向量就是每个对应输出通道的卷积核（）。由于BN层常常在Conv层之后，我们可以进行两个操作的合并。

然后我们将BN层和Conv层融合：这里我们记：

• 和为BN的参数,
• 和是BN前面Conv层的参数,
• 为Conv层的输入，
• 为输入层的channel数，
• 是Conv层的卷积核大小

我们将的每个卷积部分reshape为一个维度为 的向量，因此Conv层加BN层的操作为：

显然，我们可以将Conv层和BN层合并成一个新的卷积层，其参数为：

• filter weights: 
• bias: 

最后我们在PyTorch中实现这个融合操作：nn.Conv2d参数:

• filter weights,  : conv.weight;
• bias, : conv.bias;

nn.BatchNorm2d参数:

• scaling,  : bn.weight;
• shift, : bn.bias;
• mean estimate,  : bn.running_mean;
• variance estimate,  : bn.running_var;
• (为了数值稳定性): bn.eps.

具体的实现代码如下（Google Colab, https://colab.research.google.com/drive/1mRyq_LlJW4u_rArzzhEe_T6tmEWoNN1K）:

    import torch
    import torchvision
    
    def fuse(conv, bn):
    
        fused = torch.nn.Conv2d(
            conv.in_channels,
            conv.out_channels,
            kernel_size=conv.kernel_size,
            stride=conv.stride,
            padding=conv.padding,
            bias=True
        )
    
        # setting weights
        w_conv = conv.weight.clone().view(conv.out_channels, -1)
        w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps+bn.running_var)))
        fused.weight.copy_( torch.mm(w_bn, w_conv).view(fused.weight.size()) )
        
        # setting bias
        if conv.bias is not None:
            b_conv = conv.bias
        else:
            b_conv = torch.zeros( conv.weight.size(0) )
        b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(
                              torch.sqrt(bn.running_var + bn.eps)
                            )
        fused.bias.copy_( b_conv + b_bn )
    
        return fused
    
    # Testing
    # we need to turn off gradient calculation because we didn't write it
    torch.set_grad_enabled(False)
    x = torch.randn(16, 3, 256, 256)
    resnet18 = torchvision.models.resnet18(pretrained=True)
    # removing all learning variables, etc
    resnet18.eval()
    model = torch.nn.Sequential(
        resnet18.conv1,
        resnet18.bn1
    )
    f1 = model.forward(x)
    fused = fuse(model[0], model[1])
    f2 = fused.forward(x)
    d = (f1 - f2).mean().item()
    print("error:",d)

运行代码，会发现融合BN和Conv层之后推理结果是一样，所以是等效替换。另外也可以对比前后推理时间的差异，会发现融合后推理时间会减少。