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

使用 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 torchimport torchvisiondef 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 weightsw_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 biasif conv.bias is not None:b_conv = conv.biaselse: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 ittorch.set_grad_enabled(False)x = torch.randn(16, 3, 256, 256)resnet18 = torchvision.models.resnet18(pretrained=True)# removing all learning variables, etcresnet18.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)

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