代码收藏家技术教程 2022-09-13

YOLOv5改进之二：添加CBAM注意力机制

前言：作为当前先进的深度学习目标检测算法YOLOv5，已经集合了大量的trick，但是在处理一些复杂背景问题的时候，还是容易出现错漏检的问题。此后的系列文章，将重点对YOLOv5的如何改进进行详细的介绍，目的是为了给那些搞科研的同学需要创新点或者搞工程项目的朋友需要达到更好的效果提供自己的微薄帮助和参考。

解决问题：加入CBAM双通道注意力机制，可以让网络更加关注待检测目标，提高检测效果

添加方法：

第一步：确定添加的位置，作为即插即用的注意力模块，可以添加到YOLOv5网络中的任何地方。本文以添加进卷积Conv模块中为例。

第二步：common.py构建融入CBAM模块的Conv_CBAM,与原Conv模块不同的是：在该模块中的激活函数后加入CBAM模块。

class Conv_CBAM(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(Conv_CBAM, self).__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = nn.Hardswish() if act else nn.Identity()
        self.ca = ChannelAttention(c2)
        self.sa = SpatialAttention()

    def forward(self, x):
        x = self.act(self.bn(self.conv(x)))
        x = self.ca(x) * x
        x = self.sa(x) * x
        return x

    def fuseforward(self, x):
        return self.act(self.conv(x))

第三步：yolo.py中注册我们进行修改的Conv_CBAM模块

第四步：修改yaml文件，本文以修改主干特征提取网络为例，将原Conv模块改为Conv_CBAM即可。