目标检测 YOLOv5 - 卷积层和BN层的融合

目标检测 YOLOv5 - 卷积层和BN层的融合

即Conv2d和 BatchNorm2d融合

flyfish

为了减少模型推理时间，YOLOv5源码中attempt_load已经包括两层的合并，主要用在推理和导出模型提供给其他平台推理使用时。

函数调用链条是

attempt_load -》fuse-》fuse_conv_and_bn

fuse_conv_and_bn函数就是Conv2d层和BatchNorm2d层的合并。

在模型训练完成后，YOLOv5在推理阶段和导出模型时，将卷积层和BN层进行融合。

融合过程

（1）卷积层公式简写如下:
$$x_i=w\ast x_{i-1}+b$$

（2）Batch Normalization层公式简写如下
这里简单复述稍微详细一点看介绍BatchNorm2d
$${\mu }_{\mathcal{B}}\leftarrow \frac{1}{m}\sum _{i=1}^m{x}_i$$
$${\sigma }_{\mathcal{B}}^2\leftarrow \frac{1}{m}\sum _{i=1}^m(x_i-{\mu }_{\mathcal{B}}{)}^2$$
$${\hat{x}}_i\leftarrow \frac{x_i-{\mu }_{\mathcal{B}}}{\sqrt{{\sigma }_{\mathcal{B}}^2+ϵ}}$$

$$y_i=\frac{x_i-\mu }{\sqrt{{\sigma }^2+ϵ}}\cdot \gamma +\beta$$

4个参数

$\gamma$ 和 $\beta$ 在训练时，是可学习参数。
推理阶段均值和方差是整个训练集的均值和方差。假如硬件资源无限均值和方差训练时全部储存，最后计算推理时所用的均值和方差。硬件资源有限，训练时mini-batch的均值和方差不需要全部储存，太多了不好算，是通过滑动平均计算出整个训练集的均值和方差 。
训练完成后这四个值都是固定值即常量。
（3）将卷积层的式子带入到 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+(\frac{b-\mu }{\sqrt{{\sigma }^2+ϵ}}\cdot \gamma +\beta )\end{array}$$

融合后

上述式子可以简单看成直线方程
$$y=kx+b$$为了防止跟上面式子字母重复换一个字母
$$y=k_{1}x+k_2$$
$$k_1=\frac{w\ast \gamma }{\sqrt{{\sigma }^2+ϵ}}$$
$$k_2=\frac{b-\mu }{\sqrt{{\sigma }^2+ϵ}}+\beta$$
$$\begin{array}{rl}& w_{new}=\frac{w\ast \gamma }{\sqrt{{\sigma }^2+ϵ}}\\ & b_{new}=(\frac{b-\mu }{\sqrt{{\sigma }^2+ϵ}}\cdot \gamma +\beta )\end{array}$$

YOLOv5版本的融合代码实现

import torch
import torchvision
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
	print("conv.out_channels:",conv.out_channels)
	w_conv = conv.weight.clone().view(conv.out_channels, -1)
	print("w_conv:",w_conv.shape)
	w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps+bn.running_var)))
	print("w_bn:", w_bn.shape)
	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

#以下代码片段在ResNet18的前两层测试上述函数:

测试代码

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)

测试结果

conv.out_channels: 64
w_conv: torch.Size([64, 147])
w_bn: torch.Size([64, 64])
error: 0.00000030

PyTorch版本的融合代码实现

import copy
import torch

def fuse_conv_bn_eval(conv, bn):
    assert(not (conv.training or bn.training)), "Fusion only for eval!"
    fused_conv = copy.deepcopy(conv)

    fused_conv.weight, fused_conv.bias = \
        fuse_conv_bn_weights(fused_conv.weight, fused_conv.bias,
                             bn.running_mean, bn.running_var, bn.eps, bn.weight, bn.bias)

    return fused_conv

def fuse_conv_bn_weights(conv_w, conv_b, bn_rm, bn_rv, bn_eps, bn_w, bn_b):
    if conv_b is None:
        conv_b = torch.zeros_like(bn_rm)
    if bn_w is None:
        bn_w = torch.ones_like(bn_rm)
    if bn_b is None:
        bn_b = torch.zeros_like(bn_rm)
    bn_var_rsqrt = torch.rsqrt(bn_rv + bn_eps)

    conv_w = conv_w * (bn_w * bn_var_rsqrt).reshape([-1] + [1] * (len(conv_w.shape) - 1))
    conv_b = (conv_b - bn_rm) * bn_var_rsqrt * bn_w + bn_b

    return torch.nn.Parameter(conv_w), torch.nn.Parameter(conv_b)

我们自己添加如下代码测试

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_bn_eval(net[0], net[1])
y2 = fusedconv.forward(x)
d = (y1 - y2).norm().div(y1.norm()).item()
print("error: %.8f" % d)

输出结果

error: 0.00000030

YOLOv5版本的两层融合与 PyTorch提供的两层融合结果是一模一样的。

阅读此文章的人还阅读了以下内容
目标检测 YOLOv5 - 卷积层和BN层的融合
目标检测 YOLOv5 - Sample Assignment
目标检测 YOLOv5 - 数据增强
目标检测 YOLOv5 - 学习率
目标检测 YOLOv5 - 多机多卡训练
目标检测 YOLOv5 - 浮点取模
目标检测 YOLOv5 - 在多类别中应用NMS（非极大值抑制）
目标检测 YOLOv5 - loss for objectness and classification
目标检测 YOLOv5 - loss for bounding box regression
目标检测 YOLOv5 - 指标计算
目标检测 YOLOv5 - anchor设置
目标检测 YOLOv5 - SPP模块
目标检测 YOLOv5 - 边框预测（bounding box prediction）
目标检测 YOLOv5 - 自定义网络结构（YOLOv5-ShuffleNetV2）
目标检测 YOLOv5 - 常见的边框（bounding box ）坐标表示方法
目标检测 YOLOv5 - 图像大小与loss权重的关系
目标检测 YOLOv5 - 根据配置改变网络的深度和宽度
目标检测 YOLOv5 - 转ncnn移动端部署
目标检测 YOLOv5 - Backbone中的Focus
目标检测 YOLOv5 - 模型训练、推理、导出命令
目标检测 YOLOv5 - 人脸数据集widerface转YOLOv5格式
目标检测 YOLOv5 - 使用的数据集格式
目标检测 YOLOv5 - 中使用COCO数据集
目标检测 YOLOv5 - CrowdHuman数据集格式转YOLOv5格式