conv与bn的融合

简介

1. 意义：conv层和bn层的融合能够给网络提速。
2. 背景：先学习了这个博客的内容。学习时发现代码的实现和理论的介绍不太一样，遂深入研究了一番，修正了部分内容。
3. 文章结构：本文先介绍基本理论知识，然后解释代码。

理论知识

卷积操作

  在介绍conv与bn的融合之前，需要基本了解卷积核的卷积操作。

卷积层的卷积操作

  对于卷积层而言，若卷积层输入为$F_i$，输出为$F_o$，输入通道数为$i$，输出通道数为$j$，一个卷积核在特征图上卷积的次数为$n$（下文会解释n的含义），则有：
$$F_o=W{F}_i+b$$

$$\left[\begin{array}{c}{F}_{o1}\\ F_{o2}\\ ...\\ F_{oj}\end{array}\right]=\left[\begin{array}{c}{w}_1\\ w_2\\ ...\\ w_j\end{array}\right]\left[\begin{array}{cccc}{F}_{i1}& F_{i2}& ...& F_{in}\end{array}\right]+\left[\begin{array}{c}{b}_1\\ b_2\\ ...\\ b_j\end{array}\right]$$

需要注意的是，$F_i$的大小并非是$(n,c,w,h)$，而是对其进行了特别处理。先回忆每个$w_k(k=1,2,...,j)$在特征图上卷积时的过程：从特征图的第一个区域开始卷积，每次卷积完毕，需要“滑动”到下一个区域再次进行卷积，直到最后一个区域，一共需要卷积$n$次。因此，$F_k(k=1,2,...,n)$应该为特征图中被卷积核第$k$次卷积的区域。

  举个简单的例子，假设特征图为一个二维矩阵：
$$\left[\begin{array}{ccc}1& 2& 3\\ 2& 3& 4\\ 3& 4& 5\end{array}\right]$$

若卷积核大小为$2\ast 2$，则$n=4$，且有：
$$F_{i1}=\left[\begin{array}{cc}1& 2\\ 2& 3\end{array}\right]F_{i2}=\left[\begin{array}{cc}2& 3\\ 3& 4\end{array}\right]F_{i3}=\left[\begin{array}{cc}2& 3\\ 3& 4\end{array}\right]F_{i4}=\left[\begin{array}{cc}3& 4\\ 4& 5\end{array}\right]$$

但是，这里的$F_{i1},F_{i2},...,F_{in}$与$w_1,w_2,...,w_j$均为二维矩阵；在实际应用中，它们均为三维矩阵。又注意到它们之间的运算是对应位置相乘，因此可以将其通通reshape成1维矩阵进行点乘运算，运算后再reshape回去，且不影响计算结果。此时，矩阵$W$和$F_i$变成了二维矩阵，卷积层操作便可以由一个二维矩阵乘法的表达式来表示！

  这一节实际上只为说明一件事：虽然卷积层的卷积操作是基于“滑动窗口”的机制、在每一个窗口区域的运算是对应位置相乘的较为繁琐的运算，但该操作是能够通过二维矩阵乘法来表示的。因此，后续我们可以利用矩阵乘法的性质来对公式进行变形。

将归一化层看成卷积核为1*1的卷积操作

  对于归一化层而言，若某通道均值为$\mu$，方差为$\sigma$，归一化层的权重为$\gamma$，偏移为$\beta$，该通道输入为$X$，输出为$Y$，则有：
$$Y=\gamma \frac{X-\mu }{\sqrt[2]{{\sigma }^2+ϵ}}+\beta =\frac{\gamma }{\sqrt[2]{{\sigma }^2+ϵ}}X+(\beta -\frac{\gamma \mu }{\sqrt[2]{{\sigma }^2+ϵ}})$$
因此，归一化层可以看成输入和输出通道数相同的、卷积核大小为1*1的卷积层。其中，归一化层输出的第$i$个通道对应输入的第$i$个通道，与其它通道没有关系。因此，若将归一化层的操作用卷积层的矩阵相乘来表示，则$W$应该是一个对角矩阵。事实上，令归一化层的权重为$W_{bn}$，偏移为$b_{bn}，$则有：
$$W_{bn}=\left[\begin{array}{cccc}\frac{{\gamma }_1}{\sqrt[2]{{\sigma }_1^2+ϵ}}& & & \\ & \frac{{\gamma }_2}{\sqrt[2]{{\sigma }_2^2+ϵ}}& & \\ & & ...& \\ & & & \frac{{\gamma }_j}{\sqrt[2]{{\sigma }_j^2+ϵ}}\end{array}\right]b_{bn}=\left[\begin{array}{c}{\beta }_1-\frac{{\gamma }_1{\mu }_1}{\sqrt[2]{{\sigma }_1^2+ϵ}}\\ {\beta }_2-\frac{{\gamma }_2{\mu }_2}{\sqrt[2]{{\sigma }_2^2+ϵ}}\\ ...\\ {\beta }_j-\frac{{\gamma }_j{\mu }_j}{\sqrt[2]{{\sigma }_j^2+ϵ}}\end{array}\right]$$

将卷积层与归一化层融合

  令卷积层的输入为$F_0$，卷积层的输出或归一化层的输入为$F_1$，归一化层的输出为$F_2$；卷积层的权重为$W_{conv}$，偏移为$b_{conv}$，归一化层的权重为$W_{bn}$，偏移为$b_{bn}$，则有：
$$F_1=W_{conv}{F}_0+b_{conv}$$

$$F_2=W_{bn}{F}_1+b_{bn}$$

利用矩阵乘法的交换律，有：
$$\begin{array}{rl}{F}_2& =W_{bn}(W_{conv}{F}_0+b_{conv})+b_{bn}\\ & =(W_{bn}{W}_{conv})F_0+(W_{bn}{b}_{conv}+b_{bn})\end{array}$$

其中$W_{conv},b_{conv},W_{bn},b_{bn}$均需要reshape成$(num\_out\_channels,-1)$

代码实现

参考博客代码

  参考博客代码如下：

def fuse_conv_and_bn(conv, bn):
    #
    # init
    fusedconv = torch.nn.Conv2d(
        conv.in_channels,
        conv.out_channels,
        kernel_size=conv.kernel_size,
        stride=conv.stride,
        padding=conv.padding,
        bias=True
    )
    #
    # prepare filters
    w_conv = conv.weight.clone().view(conv.out_channels, -1)
    w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps+bn.running_var)))
    fusedconv.weight.copy_( torch.mm(w_bn, w_conv).view(fusedconv.weight.size()) )
    #
    # prepare spatial 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))
    fusedconv.bias.copy_( b_conv + b_bn )
    #
    # we're done
    return fusedconv

注意到变量b_conv直接与b_bn相加，并没有乘上w_bn。然而，该博客的测试代码所测得的结果并没有反应这个问题。下面是测试代码：

import torch
import torchvision
torch.set_grad_enabled(False)
x = torch.randn(16, 3, 256, 256)
rn18 = torchvision.models.resnet18(pretrained=True)
rn18.eval()
net = torch.nn.Sequential(
    rn18.conv1,
    rn18.bn1
)
y1 = net.forward(x)
fusedconv = fuse_conv_and_bn(net[0], net[1])
y2 = fusedconv.forward(x)
d = (y1 - y2).norm().div(y1.norm()).item()
print("error: %.8f" % d)

代码测试均值的偏差占原均值的比例，运行结果为：

error: 0.00000022

确实是很小的，说明融合后的新卷积层和原来的双层结构等效。但前面提到，b_conv并没有乘上w_bn。事实上，通过打印发现，原始的卷积层没有bias，因此b_conv全为0，乘与不成都没有区别。为了验证这一说法：

1. 重写融合函数，将b_conv乘上w_bn后再与b_bn相加。
2. 重写测试代码，为初始卷积层添加bias（打印可以看到其初始化值不全为0）。

修正代码

  重写代码如下：

1. 对于融合函数，在b_conv初始化后增加一行：

b_conv = torch.mm(w_bn, b_conv.view(conv.out_channels, -1)).view(fusedconv.bias.size())

2. 对于测试代码，在rn18.eval()后添加如下几行：

print(rn18.conv1.bias)
# # rn18.conv1.bias = torch.nn.parameter.Parameter(torch.ones(rn18.conv1.weight.size(0)))
rn18.conv1 = torch.nn.Conv2d(
    rn18.conv1.in_channels,
    rn18.conv1.out_channels,
    kernel_size=rn18.conv1.kernel_size,
    stride=rn18.conv1.stride,
    padding=rn18.conv1.padding,
    bias=True
)
print(rn18.conv1.bias)

将重写前后的代码分别测试，结果如下：

error: 0.10046744
error2: 0.00000011

这说明，原代码确实存在问题，并且新代码修正了该问题。

  经同学提醒，一般情况下，卷积层不使用bias。因此，从实际应用而言，这个错误的修正并没有多大意义，倒可以更多地被用于学习卷积层的操作以及卷积层与归一化层的融合过程。