YOLO算法改进2【中阶改进篇】：添加CBAM注意力机制

一、理论储备

卷积神经网络 (CNN)在计算机视觉任务中取得了显著的成功。为了进一步提高CNN的性能，研究人员提出了一种名为“卷积神经网络注意力模块”(Convolutional Block Attention Module, CBAM)的注意力机制。CBAM旨在自动学习输入特征图的空间和通道注意力权重，从而更好地捕捉图像中的局部信息和全局上下文。
CBAM包括两个主要组成部分: 空间注意力模块和通道注意力模块。它们可以作为插件，轻松地集成到现有的 CNNQ 架构中，以提高性能。

空间注意力模块 (Spatial Attention Module) : 空间注意力模块的目的是为输入特征图的每个位置分配一个注意力权重。这些权重可帮助网络集中注意力于感兴趣的区域。空间注意力模块主要包括以下步骤:

• a.对输入特征图进行全局平均池化和全局最大池化操作
• b.将池化后的特征图按通道相加，得到两个1维向量。
• c.对这两个向量进行点积，形成一个注意力权重矩阵
• d.将注意力权重矩阵应用于输入特征图，得到空间注意力调整后的特征图

通道注意力模块 (ChannelAttention Module) : 通道注意力模块旨在为输入特征图的每个通道分配一个权重，从而强调重要的通道并抑制不太重要的通道。通道注意力模块主要包括以下步骤:

• a.对输入特征图进行全局平均池化和全局最大池化操作.
• b.使用全连接层 (MLP) 学习每个通道的权重.
  c.将学到的权重应用于输入特征图，得到通道注意力调整后的特征图,

最后，将空间注意力模块和通道注意力模块的输出相加，形成CBAM调整后的特征图。这个特征图可以继续传递到CNN的下一层。
实验证明，CBAM在各种 计算机视觉Q 任务上都能显著提高性能，如图像分类、目标检测和语义分割等。它是一种通用的注意力机制，可以与任何卷积神经网络架构相结合。

论文地址：CBAM: Convolutional Block Attention Module

CBAM注意力机制原理图

二、代码实现

1.增加CBAM.yaml文件
添加至/models/中

# Parameters
nc: 80  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
anchors:
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

# YOLOv5 v6.0 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
   [-1, 3, C3, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
   [-1, 6, C3, [256]],
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
   [-1, 9, C3, [512]],
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
   [-1, 3, C3, [1024]],
   [-1, 1, SPPF, [1024, 5]],  # 9
  ]

# YOLOv5 v6.0 head
head:
  [[-1, 1, Conv, [512, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
   [-1, 3, C3, [512, False]],  # 13

   [-1, 1, Conv, [256, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)

   [-1, 1, Conv, [256, 3, 2]],
   [[-1, 14], 1, Concat, [1]],  # cat head P4
   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)

   [-1, 1, Conv, [512, 3, 2]],
   [[-1, 10], 1, Concat, [1]],  # cat head P5
   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
   [-1, 1, CBAM, [1024]],

   [[17, 20, 24], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

2.common.py配置
在./models/common.py文件中增加以下模块

class ChannelAttentionModule(nn.Module):
    def __init__(self, c1, reduction=16):
        super(ChannelAttentionModule, self).__init__()
        mid_channel = c1 // reduction
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.shared_MLP = nn.Sequential(
            nn.Linear(in_features=c1, out_features=mid_channel),
            nn.LeakyReLU(0.1, inplace=True),
            nn.Linear(in_features=mid_channel, out_features=c1)
        )
        self.act = nn.Sigmoid()
   
    def forward(self, x):
        avgout = self.shared_MLP(self.avg_pool(x).view(x.size(0),-1)).unsqueeze(2).unsqueeze(3)
        maxout = self.shared_MLP(self.max_pool(x).view(x.size(0),-1)).unsqueeze(2).unsqueeze(3)
        return self.act(avgout + maxout)
        
class SpatialAttentionModule(nn.Module):
    def __init__(self):
        super(SpatialAttentionModule, self).__init__()
        self.conv2d = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=7, stride=1, padding=3)
        self.act = nn.Sigmoid()
    def forward(self, x):
        avgout = torch.mean(x, dim=1, keepdim=True)
        maxout, _ = torch.max(x, dim=1, keepdim=True)
        out = torch.cat([avgout, maxout], dim=1)
        out = self.act(self.conv2d(out))
        return out

class CBAM(nn.Module):
    def __init__(self, c1,c2):
        super(CBAM, self).__init__()
        self.channel_attention = ChannelAttentionModule(c1)
        self.spatial_attention = SpatialAttentionModule()

    def forward(self, x):
        out = self.channel_attention(x) * x
        out = self.spatial_attention(out) * out
        return out

3.yolo.py配置
找到models/yolo.py文件中parse_model()函数的for i，（f,n,m,args）in enumerate(d[‘backjbone’]+d[‘head’])(258行上下)并其循环内添加如下代码。

        elif m is CBAM:
            c1, c2 = ch[f], args[0]
            if c2 != no:
                c2 = make_divisible(c2 * gw, 8)
            args = [c1, c2]

4.训练模型

python train.py --cfg CBAM.yaml

在原有YOLO中运行python train.py --cfg CBAM.yaml，会产生报错，报错内容如下：

RuntimeError: adaptive_max_pool2d_backward_cuda does not have a deterministic implementation, but you set ‘torch.use_deterministic_algorithms(True)’. You can turn off determinism just for this operation, or you can use the ‘warn_only=True’ option, if that’s acceptable for your application. You can also file an issue at https://github.com/pytorch/pytorch/issues to help us prioritize adding deterministic support for this operation.

解决方法：
打开train.py文件，在scaler.scale(loss).backward()前添加torch.use_deterministic_algorithms(False)

三、训练结果对比

未添加注意力机制的YOLO_V5的结果：

添加注意力原YOLO_V5添加CBAM注意力机制后的结果：

本文摘抄自K同学博客，仅为个人学习使用