模型部署——融合BN和Conv层

今天我们将尝试了解如何使我们的模型在推理上更快一点。

使用 Batch Normalization 作为提高泛化能力的一种方式浪费了大量的网络。 但是在推理过程中，批量归一化被关闭，而是使用近似的每个通道均值和方差。 很酷的是，我们可以通过 1x1 卷积实现相同的行为。 更好的是将Batch Normalization 与前面的卷积合并。

Batch Normalization

假设$x$ 是要归一化的激活信号。 给定一组来自一个batch中不同样本的此类信号$x_1,x_2,...,x_n$，归一化如下：

$${\hat{x}}_i=\gamma \frac{x_i-\mu }{\sqrt{{\sigma }^2+ϵ}}+\beta =\frac{\gamma x_i}{\sqrt{{\sigma }^2+ϵ}}+\beta -\frac{\gamma \mu }{\sqrt{{\sigma }^2+ϵ}}$$

这里 $\mu$ 和 ${\sigma }^2$ 为这个batch上计算得到的均值和方差（在B,H,W维度上计算，每个channel单独计算），而 $ϵ$ 是防止除零所设置的一个极小值， $\gamma$ 是比例参数，而 $\beta$ 是平移系数。在训练过程中， $\mu$ 和 $\sigma$ 在当前batch上计算：

$$\begin{gathered} \mu =\frac{1}{n}\sum x_i \\ {\sigma }^2=\frac{1}{n}\sum (x_i-\mu {)}^2 \end{gathered}$$

参数 $\gamma$ 和 $\beta$ 与网络的其他参数一起通过梯度下降缓慢学习。 在测试期间，通常不会在一个batch图像上运行网络。 因此，不能使用前面提到的 $\mu$ 和 $\sigma$ 公式。 相反，我们使用他们在训练期间通过exponential moving average计算的估计值 $\hat{\mu }$ 和 ${\hat{\sigma }}^2$

如今，批量归一化主要用于卷积神经网络中。 在此设置中，输入特征图的每个通道 $c$ 都有均值 ${\mu }_c$ 和方差估计 ${\sigma }_c^2$ 、平移 ${\beta }_c$ 和比例参数 ${\gamma }_c$

融合方案

对于一个形状为$C\times H\times W$的特征图$F$，记归一化结果$\hat{F}$，计算如下：

上式为$f(x)=Wx+b$的形式，可以看成$1\times 1$卷积，由于BN层常常在Conv层之后，可以将两操作合并。

融合BN和卷积

• ${\mathbf{w}}_{BN}\in {\mathbb{R}}^{C\times C}$ 和 ${\mathbf{b}}_{BN}\in {\mathbb{R}}^C$ 是BN的参数
• ${\mathbf{w}}_{conv}\in {\mathbb{R}}^{C\times C_{pre}.k^2}$和${\mathbf{b}}_{conv}\in {\mathbb{R}}^C$是Conv层的参数
• $F_{prev}$是卷积的输入
• $C_{prev}$：输入层的通道数量
• $k$：卷积核大小

将 $F_{prev}$的每个$k\times k$部分reshape为一个维度为$k^2.C_{prev}$ 的向量$f_{i,j}$，因此Conv层加BN层的操作为：

$${\hat{\mathbf{f}}}_{i,j}={\mathbf{W}}_{BN}.({\mathbf{W}}_{conv}.{\mathbf{f}}_{i,j}+{\mathbf{b}}_{conv})+{\mathbf{b}}_{BN}$$

因此，我们可以用具有以下参数的单个卷积层替换这两层：

• 滤波器权重$W$：$\mathbf{W}={\mathbf{W}}_{BN}.{\mathbf{W}}_{conv}$
• 偏置bias：$\mathbf{b}={\mathbf{W}}_{BN}.{\mathbf{b}}_{conv}+{\mathbf{b}}_{BN}$

pytorch实现：

nn.Conv2d参数：

• 滤波器权重，$\mathbf{W}$：conv.weight
• bias，$\mathbf{b}$：conv.bias

nn.BatchNorm2d参数：

• scaling，$\gamma$：bn.weight
  shift，$\beta$：bn.bias
• mean estimate，$\hat{\mu }$： bn.running_mean
• variance estimate，${\sigma }^2$：bn.running_var
• $ϵ$(for numerical stability):：bn.eps

代码实现：

    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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

参考：https://learnml.today/speeding-up-model-with-fusing-batch-normalization-and-convolution-3