pytorch实现卷积层和BN层融合

批归一化

数据的规范化也即$(x-mean(x)/var(x))$ 可以将数据的不同特征规范到均值为0，方差相同的标准正态分布,这就可以避免大数造成得数据溢出。可以使网络以更高得学习率更快得收敛。
如果仅仅对输入的图片进行归一化，后面其余网络层的输出不进行相关操作，那么随着网络深度得增加，数据得分布会逐渐偏离标准正态分布，由于梯度得反向传播是不同层得梯度不断相乘得结果，数据过小或过大容易造成梯度消失或梯度爆炸。
为了解决上述问题，有人提出了对网络层的中间输出进行规范化得方法，也就是批规范化（Batch Normalization）,对某个网络层的批规范化，通常是在加权和与激活函数运算之间进行的。也就是：
$$sigmoid(BN(xW+b))=sigmoid(BN(z))$$ 但是如果将$z$ 规范到$N(0,1)$ 的标准正态分布，那么无论前面的网络如何变化，经过这个这个层后都将服从标准正态分布，批规范化引入了可学习的参数$\beta 、\gamma$ ,将规范到$N(0,1)$ 的标准正态分布变换到$N(\beta ,\gamma )$ 由于$\beta 、\gamma$ 是可学习的，所以就避免了模型能力降低的问题：
$${\hat{x}}_i=\gamma \frac{x_i-\mu }{\sqrt{{\sigma }^2+ϵ}}+\beta$$其中$$\begin{array}{c}\mu =\frac{1}{n}\sum x_i\\ {\sigma }^2=\frac{1}{n}\sum {\left(x_i-\mu \right)}^2\end{array}$$

BN和卷积层融合

在图像处理过程中，BN经常用于卷积之后，可以考虑将两个融合在在一起 ，简单点就是将是合二为一（融合层通常用测试阶段）。那么合二为一新的网络层的网络权重和偏置该怎么初始化呢？根据卷积层和BN层的相关参数进行计算得到新的权重和偏置，公式入下：$$\begin{array}{rl}{y}_i& =\frac{x_i-\mu }{\sqrt{{\sigma }^2+ϵ}}\cdot \gamma +\beta \\ & =\frac{w\ast x_i+b-\mu }{\sqrt{{\sigma }^2+ϵ}}\cdot \gamma +\beta \\ & =\frac{w\ast \gamma }{\sqrt{{\sigma }^2+ϵ}}\cdot x+\left(\frac{b-\mu }{\sqrt{{\sigma }^2+ϵ}}\cdot \gamma +\beta \right)\end{array}(1)$$
简化公式结构为$y_i=w_{new}x+b_{new}$.
$$\begin{array}{l}{w}_{\text{new }}=\frac{w\ast \gamma }{\sqrt{{\sigma }^2+ϵ}}\\ b_{\text{new }}=\left(\frac{b-\mu }{\sqrt{{\sigma }^2+ϵ}}\cdot \gamma +\beta \right)\end{array}(2)$$ 其中$w，b$为卷积层的权重和偏置，$\gamma 、\beta ，\mu 、{\sigma }^2$为$BN$层的权重、偏置、$mean、var$对应的参数

pytorch代码

pytorch 中 Conv2 和 BN 层中参数的的具体意义

import torch
def fuse_conv_and_bn(conv, bn):
	# 初始化
	fusedconv = torch.nn.Conv2d(
		conv.in_channels,
		conv.out_channels,
		kernel_size=conv.kernel_size,
		stride=conv.stride,
		padding=conv.padding,
		bias=True
	)
	w_conv = conv.weight.clone().view(conv.out_channels, -1)
	w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps+bn.running_var)))
	# 融合层的权重初始化(W_bn*w_conv(卷积的权重))
	fusedconv.weight.copy_( torch.mm(w_bn, w_conv).view(fusedconv.weight.size()) )

	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_( torch.matmul(w_bn, b_conv) + b_bn )

	return fusedconv

那么，融合层计算出来的结果和分层计算出来的结果有多大差距呢？？从下面代码可以看到用人家预训练好的相关参数进行融合，融合层和分层计算得出的结果所差无几嘛。

import torchvision
torch.set_grad_enabled(False)
x = torch.randn(32, 3, 256, 256)   # batch,channel,w,h
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)
# torch.norm()是求范数，没有指定维度则是对所有元素求二范数
d = (y1 - y2).norm().div(y1.norm()).item()
print("error: %.8f" % d)

>>>error: 0.00000022

那么问题又出来了，既然所查无几，为什么要费劲的去融合他呢，那我们再想想深度网络运行时间是不是我们担心的问题呢，走起，去验证一下时间上有无改进吧。
这里就假设一次输入32张图片，测试了300次。可以看出，在速度提升上还是蛮大的。

def cal_time(fun):
    def inner(*args,**kwargs):
        t1 = time.time()
        fun(*args,**kwargs)
        t2 = time.time()
        print(f"运行时间:{t2-t1}")    
    return inner

@cal_time
def main(net,x):
    for i in tqdm(range(300)): # 假设一次输入32张图片，测试了300次。
        y = net.forward(x)
        
main(net,x)
main(fusedconv,x)

>>>93.39383339
>>>69.93862175

那么事在训练的时候融合还是测试的时候融合呢，那肯定是用在测试阶段，前面说过融合层的权重和偏置要根据训练好的卷积、BN层的参数确定。这样测试阶段使用融合后的网络进行推理就可以加快推理速度。对于有些实时检测场所，速度越快岂不是更好。
参考链接: Fusing batch normalization and convolution in runtime
小白一枚，如果有问题，希望大家提出来。谢谢。