YOLOv5的Tricks | 【Trick4】参数重结构化(融合Conv+BatchNorm2d)

如有错误,恳请指出。

文章目录

这篇文章是想要记录yolov5在模型搭建过程中的一个融合模块,就是把卷积与批归一化的参数进行融合,想卷积带有批归一化的性质,使得推理过程中可以加快模型推理速度,简化整个模型结构,实现训练与推理两个阶段的解耦。

  1. 参数融合概念介绍

我最早接触参数重结构化这个词是看见了大佬丁霄汉发表的几篇论文:RepVGG,RepMLP,RepLKNet,这些构建新backbone的论文无一例外的全部使用了参数重结构化的思想。

RepVGG将3×3,1×1,identity分支的残差结果利用数学计算方法等价为一个3×3的卷积结构,实现训练与推断过程的解耦;RepMLP将局部的CNN先验信息加进了全连接层,使得其与MLP相结合等等。这里需要注意,重结构化层MLP结构也不是说变成Linear层,而是简化为1×1的卷积。(后续有机会把这几篇文章介绍一下,或者直接看大佬的知乎:https://www.zhihu.com/people/ding-xiao-yi-93/posts)

BN(批归一化)层常用于在卷积层之后,对feature maps进行归一化,从而加速网络学习,也具有一定的正则化效果。 训练时,BN需要学习一个minibatch数据的均值、方差,然后利用这些信息进行归一化。而在推理过程,通常为了加速,都会把BN融入到其上层卷积中,这样就将两步运算变成了一步,也就达到了加速目的。

那么这里yolov5所实现的, 是参数重结构化的一个小内容,就是把卷积与批归一化进行融合,变成一个新的卷积,但是包含BN层的特性。所以相比之下,算是参数重结构化系列的一个小小idea,可以稍微的加快推理速度。因为使用的是csp结构,所以没有涉及多并联分支的卷积模块(所以这一点其实也可以魔改下yolov5试试)。

  1. 参数融合详细推导

在yolov5的注释中给了一个推导的参考资料:Fusing batch normalization and convolution in runtime,代码也是基于这篇文章来稍微修改的。

其实现的主要思想就是将bn层转化为一个1×1的卷积:

YOLOv5的Tricks | 【Trick4】参数重结构化(融合Conv+BatchNorm2d)

然后就变成了两个卷积层的迭代处理,公式为:
f ^ i , j = W B N ⋅ ( W c o n v ⋅ f i , j + b c o n v ) + b B N = ( W B N ⋅ W c o n v ) ⋅ f i , j + ( W B N ⋅ b c o n v + b B N ) = W ⋅ f i , j + b \begin{aligned} \hat{f}{i,j} &=W{BN}·(W_{conv}·f_{i,j}+b_{conv})+b_{BN} \ &=(W_{BN}·W_{conv})·f_{i,j}+(W_{BN}·b_{conv}+b_{BN}) \ &=W·f_{i,j}+b \end{aligned}f ^​i ,j ​​=W B N ​⋅(W c o n v ​⋅f i ,j ​+b c o n v ​)+b B N ​=(W B N ​⋅W c o n v ​)⋅f i ,j ​+(W B N ​⋅b c o n v ​+b B N ​)=W ⋅f i ,j ​+b ​

其中:

  • filter weights:W = W B N ⋅ W c o n v W=W_{BN}·W_{conv}W =W B N ​⋅W c o n v ​
  • bias:b = W B N ⋅ b c o n v + b B N b=W_{BN}·b_{conv}+b_{BN}b =W B N ​⋅b c o n v ​+b B N ​

在pytorch实现中,每个BN层都有以下几个:

  • scaling γ \gamma γ:bn.weight
  • shift β \beta β:bn.bias
  • mean estiamte μ ^ \hat{μ}μ^​:bn.running_mean
  • variance estimate σ ^ 2 \hat{\sigma}^2 σ^2:bn.running_var
  • ϵ ϵϵ (for numerical stability):bn.eps

但是在批归一化转换为1×1的卷积那里其实没有给出太多解释,后来我看了另一篇博客介绍:卷积层与BN层的融合方式,其实我是没有完全弄到batchnormalization的过程。为了搞清楚如何融合卷积和BN,需要先搞懂卷积和BN的过程。

对于卷积层,卷积核的权重为W W W,卷积过程就是利用W W W在其输入feature map中滑窗计算,公式表达为:y c o n v = w ⋅ x + b y_{conv}=w·x+b y c o n v ​=w ⋅x +b

对于BN层,需要计算一个minibatch中元素的均值方差,然后对于x需要减去均值除以标准差,最后利用γ , β γ,βγ,β进行仿射变换,即可得到最终的BN输出

YOLOv5的Tricks | 【Trick4】参数重结构化(融合Conv+BatchNorm2d)
其中,第一个公式为求均值、第二个公式为求方差、第三个公式为归一化、第四个公式为仿射变换

这里直接将卷积公式:y c o n v = w ⋅ x + b y_{conv}=w·x+b y c o n v ​=w ⋅x +b代入到BN的公式中,如下所示:

YOLOv5的Tricks | 【Trick4】参数重结构化(融合Conv+BatchNorm2d)
这里重新将其变化为w ⋅ x + b w·x+b w ⋅x +b的形式:
w ^ = γ w σ B 2 + ϵ = γ σ B 2 + ϵ ⋅ w = W B N ⋅ W c o n v \begin{aligned} \hat{w} &= \frac{\gamma w}{\sqrt{\sigma^2_{B}+ϵ}} \ &= \frac{\gamma }{\sqrt{\sigma^2_{B}+ϵ}} ·w \ &= W_{BN}·W_{conv} \end{aligned}w ^​=σB 2 ​+ϵ​γw ​=σB 2 ​+ϵ​γ​⋅w =W B N ​⋅W c o n v ​​

b ^ = γ σ B 2 + ϵ ⋅ ( b − μ B ) + β = γ σ B 2 + ϵ ⋅ b − γ σ B 2 + ϵ ⋅ μ B + β = γ σ B 2 + ϵ ⋅ b + ( β − γ σ B 2 + ϵ ⋅ μ B ) = W B N ⋅ b c o n v + b B N \begin{aligned} \hat{b} &= \frac{\gamma}{\sqrt{\sigma^2_{B}+ϵ}}·(b-μ_{B})+\beta \ &= \frac{\gamma}{\sqrt{\sigma^2_{B}+ϵ}}·b – \frac{\gamma}{\sqrt{\sigma^2_{B}+ϵ}}·μ_{B} + \beta \ &= \frac{\gamma}{\sqrt{\sigma^2_{B}+ϵ}}·b + (\beta – \frac{\gamma}{\sqrt{\sigma^2_{B}+ϵ}}·μ_{B} ) \ &= W_{BN}·b_{conv} + b_{BN} \end{aligned}b ^​=σB 2 ​+ϵ​γ​⋅(b −μB ​)+β=σB 2 ​+ϵ​γ​⋅b −σB 2 ​+ϵ​γ​⋅μB ​+β=σB 2 ​+ϵ​γ​⋅b +(β−σB 2 ​+ϵ​γ​⋅μB ​)=W B N ​⋅b c o n v ​+b B N ​​

所以上面的原结果转换为:

YOLOv5的Tricks | 【Trick4】参数重结构化(融合Conv+BatchNorm2d)
至此完成了卷积层和BN层的融合,可以和代码的参考文档一一对上了。

; 3. 参数融合代码实现

这里的融合代码集成了模型中,然后再另外的调用其他打码,yolov5代码如下所示:

class Model(nn.Module):
    def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None):
        super().__init__()
        ...

        with open(cfg, errors='ignore') as f:
            self.yaml = yaml.safe_load(f)

        self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch])
        ...

    def forward(self, x, augment=False, profile=False, visualize=False):
        if augment:
            return self._forward_augment(x)
        return self._forward_once(x, profile, visualize)

    def fuse(self):
        LOGGER.info('Fusing layers... ')
        for m in self.model.modules():

            if isinstance(m, (Conv, DWConv)) and hasattr(m, 'bn'):

                m.conv = fuse_conv_and_bn(m.conv, m.bn)

                delattr(m, 'bn')

                m.forward = m.forward_fuse
        self.info()
        return self
    ...

融合Conv+BatchNorm2d的具体实现代码如下所示:

def fuse_conv_and_bn(conv, bn):

    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)

    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.shape))

    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))

    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):

    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
        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):

    def __init__(self, c1, c2, k=1, s=1, act=True):

        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层的融合方式

Original: https://blog.csdn.net/weixin_44751294/article/details/125095552
Author: Clichong
Title: YOLOv5的Tricks | 【Trick4】参数重结构化(融合Conv+BatchNorm2d)

原创文章受到原创版权保护。转载请注明出处:https://www.johngo689.com/623245/

转载文章受原作者版权保护。转载请注明原作者出处!

(0)

大家都在看

亲爱的 Coder【最近整理,可免费获取】👉 最新必读书单  | 👏 面试题下载  | 🌎 免费的AI知识星球