1. 参数融合概念介绍
我最早接触参数重结构化这个词是看见了大佬丁霄汉发表的几篇论文:RepVGG,RepMLP,RepLKNet,这些构建新backbone的论文无一例外的全部使用了参数重结构化的思想。
RepVGG将3x3,1x1,identity分支的残差结果利用数学计算方法等价为一个3x3的卷积结构,实现训练与推断过程的解耦;RepMLP将局部的CNN先验信息加进了全连接层,使得其与MLP相结合等等。这里需要注意,重结构化层MLP结构也不是说变成Linear层,而是简化为1x1的卷积。(后续有机会把这几篇文章介绍一下,或者直接看大佬的知乎:https://www.zhihu.com/people/ding-xiao-yi-93/posts)
BN(批归一化)层常用于在卷积层之后,对feature maps进行归一化,从而加速网络学习,也具有一定的正则化效果。训练时,BN需要学习一个minibatch数据的均值、方差,然后利用这些信息进行归一化。而在推理过程,通常为了加速,都会把BN融入到其上层卷积中,这样就将两步运算变成了一步,也就达到了加速目的。
那么这里yolov5所实现的,是参数重结构化的一个小内容,就是把卷积与批归一化进行融合,变成一个新的卷积,但是包含BN层的特性。所以相比之下,算是参数重结构化系列的一个小小idea,可以稍微的加快推理速度。因为使用的是csp结构,所以没有涉及多并联分支的卷积模块(所以这一点其实也可以魔改下yolov5试试)。
2. 参数融合详细推导
在yolov5的注释中给了一个推导的参考资料:Fusing batch normalization and convolution in runtime,代码也是基于这篇文章来稍微修改的。
其实现的主要思想就是将bn层转化为一个1x1的卷积:
然后就变成了两个卷积层的迭代处理,公式为:
其中:
在pytorch实现中,每个BN层都有以下几个:
但是在批归一化转换为1x1的卷积那里其实没有给出太多解释,后来我看了另一篇博客介绍:卷积层与BN层的融合方式,其实我是没有完全弄到batchnormalization的过程。为了搞清楚如何融合卷积和BN,需要先搞懂卷积和BN的过程。
至此完成了卷积层和BN层的融合,可以和代码的参考文档一一对上了。
3. 参数融合代码实现
这里的融合代码集成了模型中,然后再另外的调用其他打码,yolov5代码如下所示:
class Model(nn.Module): def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None): # model, input channels, number of classes super().__init__() ... # 如果配置文件中有中文,打开时要加encoding参数 with open(cfg, errors='ignore') as f: self.yaml = yaml.safe_load(f) # model dict # 创建网络模型 # self.model: 初始化的整个网络模型(包括Detect层结构) # self.save: 所有层结构中from不等于-1的序号,并排好序 [4, 6, 10, 14, 17, 20, 23] self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist ... def forward(self, x, augment=False, profile=False, visualize=False): # debug同样需要第三次才能正常跳进来 if augment: # use Test Time Augmentation(TTA), 如果打开会对图片进行scale和flip return self._forward_augment(x) # augmented inference, None return self._forward_once(x, profile, visualize) # single-scale inference, train # 参数重结构化: 融合conv2d + batchnorm2d (推理的时候用, 可以加快模型的推理速度) def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers LOGGER.info('Fusing layers... ') for m in self.model.modules(): # 对Conv与DWConv(继承Conv)的结构进行卷积与bn的融合 if isinstance(m, (Conv, DWConv)) and hasattr(m, 'bn'): # 融合Conv模块中的conv与bn层(不包含激活函数), 返回的是参数融合后的卷积 m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv # 融合后conv的参数就包含了bn的用途, 所以可以删除bn层 delattr(m, 'bn') # remove batchnorm # 由于不需要bn层, 所以forward函数需要改写: # self.act(self.bn(self.conv(x))) -> self.act(self.conv(x)) m.forward = m.forward_fuse # update forward self.info() return self ...
融合Conv+BatchNorm2d的具体实现代码如下所示:
def fuse_conv_and_bn(conv, bn): # Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/ # 设置为no grad(推理不需要反向传播), 指定在相同的设备上 fusedconv = nn.Conv2d(conv.in_channels, conv.out_channels, kernel_size=conv.kernel_size, stride=conv.stride, padding=conv.padding, groups=conv.groups, bias=True).requires_grad_(False).to(conv.weight.device) # prepare filters: W = W_BN * W_conv # (32,3,6,6) -> (32,3*6*6) -> (32,108) w_conv = conv.weight.clone().view(conv.out_channels, -1) w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var))) # (32,32)*(32,108) -> (32,108) -> (32,3,6,6) fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape)) # prepare spatial bias: b = W_BN * b_conv + b_BN b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps)) # (32,32)*(32,1) -> (32,1) -> (32) fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn) return fusedconv
同时可以注意到,在fuse模块代码中还改变了模型模块的前向传播函数:m.forward = m.forward_fuse,这里是因为在Conv中是卷积+bn连续使用的,其作为整个模型的一个基础卷积模块,使用只需要改变其前向传播过程就可以融合卷积+BN层,在推理的时候可以加快速度。其代码如下:
class Conv(nn.Module): # Standard convolution def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups super().__init__() self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False) self.bn = nn.BatchNorm2d(c2) self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity()) def forward(self, x): return self.act(self.bn(self.conv(x))) def forward_fuse(self, x): return self.act(self.conv(x)) class DWConv(Conv): # Depth-wise convolution class def __init__(self, c1, c2, k=1, s=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups # math.gcd最大公约数 super().__init__(c1, c2, k, s, g=math.gcd(c1, c2), act=act)
后续:如果还想了解其他不同卷积和之间的融合或者与全连接的融合可以了解重参数化系列文章:RepVGG,RepMLP等
参考资料:
1. Fusing batch normalization and convolution in runtime
2. 卷积层与BN层的融合方式