【YOLO系列改进NO.2】添加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))
class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.f1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)
        self.relu = nn.ReLU()
        self.f2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
        # 写法二,亦可使用顺序容器
        # self.sharedMLP = nn.Sequential(
        # nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False), nn.ReLU(),
        # nn.Conv2d(in_planes // rotio, in_planes, 1, bias=False))

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.f2(self.relu(self.f1(self.avg_pool(x))))
        max_out = self.f2(self.relu(self.f1(self.max_pool(x))))
        out = self.sigmoid(avg_out + max_out)
        return out


class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()

        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1

        self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x = torch.cat([avg_out, max_out], dim=1)
        x = self.conv(x)
        return self.sigmoid(x)

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

 if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP,GhostBottleneck,ghostC3,GhostConv,
                 C3,seC3,C3TR,CBAMC3,CoordAtt,Conv_CBAM

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

第五步：将train.py中改为本文的yaml文件即可，开始训练。

结 果：本人在多个数据集上做了大量实验，针对不同的数据集效果不同，同一个数据集的不同添加位置方法也是有差异，需要大家进行实验。有效果有提升的情况占大多数。

预告一下：下一篇内容分享Coordinate坐标注意力机制。有兴趣的朋友可以关注一下我，有问题可以留言或者私聊我哦

PS：CBAM双通道注意力机制，不仅仅是可以添加进YOLOv5，也可以添加进任何其他的深度学习网络，不管是分类还是检测还是分割，主要是计算机视觉领域，都可能会有不同程度的提升效果。

最后，希望能互粉一下，做个朋友，一起学习交流。

​