YOLOv5改进 | 注意力机制 | 通道和空间的双重作用的CBAM注意力机制

原创已于 2024-05-12 17:40:19 修改 · 3.3k 阅读

76 ·

CC 4.0 BY-SA版权

文章标签：

#YOLO #目标跟踪 #人工智能

于 2024-05-12 17:07:11 首次发布

YOLOv5入门 + 改进涨点专栏收录该内容

96 篇文章

订阅专栏

在深度学习目标检测领域，YOLOv5成为了备受关注的模型之一。本文给大家带来的是通道和空间的双重作用的CBAM注意力机制。文章在介绍主要的原理后，将手把手教学如何进行模块的代码添加和修改，并将修改后的完整代码放在文章的最后，方便大家一键运行，小白也可轻松上手实践。以帮助您更好地学习深度学习目标检测YOLO系列的挑战。

专栏地址：YOLOv5改进+入门——持续更新各种有效涨点方法

1.原理

论文地址：CBAM: Convolutional Block Attention Module点击即可跳转

实现代码：CBAM代码实现点击即可跳转

CBAM（Convolutional Block Attention Module）是一种引入了注意力机制的卷积神经网络模块，旨在增强CNN模型的表征能力和性能。它由两个关键组件组成：通道注意力模块（Channel Attention Module）和空间注意力模块（Spatial Attention Module）。

1. 通道注意力模块（CAM）：
CAM主要用于建模特征图在通道维度上的关系。它的目标是学习每个通道的重要性，并对不同通道的特征进行加权，以提升有用特征的影响力，抑制无用特征的干扰。
CAM首先对输入的特征图进行全局平均池化（Global Average Pooling），将每个通道的特征图压缩成一个标量，然后通过全连接层（FC）学习得到每个通道的权重。这些权重用于对每个通道的特征图进行加权，得到加权后的特征表示。

2. 空间注意力模块（SAM）：
SAM用于捕捉特征图在空间维度上的重要性。它的目标是学习不同空间位置的权重，使网络能够更好地关注图像中的重要区域。
SAM首先对特征图进行两种池化操作：最大池化和平均池化。这两种池化操作分别用于捕捉特征图中的局部显著性和全局分布信息。然后，将两种池化结果结合，并通过全连接层学习得到每个空间位置的权重，以产生最终的空间注意力图。

通过结合通道注意力和空间注意力，CBAM可以使网络更好地理解输入数据中的关键信息，并提高模型在各种视觉任务上的性能。这种注意力机制的引入使得网络能够自适应地调整特征图中不同通道和空间位置的重要性，从而有效地提升了模型的表现力和泛化能力。CBAM已经被成功应用于图像分类、目标检测、语义分割等多个计算机视觉任务中，取得了显著的性能提升。

2.YOLOv5添加CBAM注意力机制

2.1 CBAM注意力机制代码

关键步骤一：将下面代码添加到 yolov5/models/common.py中任意位置

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.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
        # (特征图的大小-算子的size+2*padding)/步长+1
        self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # 1*h*w
        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)
        #2*h*w
        x = self.conv(x)
        #1*h*w
        return self.sigmoid(x)


class CBAM(nn.Module):
    # CSP Bottleneck with 3 convolutions
    def __init__(self, c1, c2, ratio=16, kernel_size=7):  # ch_in, ch_out, number, shortcut, groups, expansion
        super(CBAM, self).__init__()
        self.channel_attention = ChannelAttention(c1, ratio)
        self.spatial_attention = SpatialAttention(kernel_size)

    def forward(self, x):
        out = self.channel_attention(x) * x
        # c*h*w
        # c*h*w * 1*h*w
        out = self.spatial_attention(out) * out
        return out

CBAM（Convolutional Block Attention Module）注意力机制的流程可以总结如下：

1. 输入特征图：接收来自上一层或输入图像的特征图作为输入。

2. 通道注意力模块（Channel Attention Module，CAM）：
对输入特征图进行全局平均池化（Global Average Pooling），将每个通道的特征进行降维，得到每个通道的全局描述。
通过全连接层（Fully Connected Layer）学习得到每个通道的权重向量，这些权重用于衡量每个通道的重要性。
将学习到的权重与原始特征图相乘，以加权增强有用特征和抑制无用特征。

3. 空间注意力模块（Spatial Attention Module，SAM）：
对输入特征图进行最大池化（Max Pooling）和平均池化（Average Pooling），分别捕获局部显著性和全局分布信息。
将两种池化结果进行组合（如相加），得到综合的空间注意力图。
通过激活函数（如sigmoid）对空间注意力图进行归一化，得到每个空间位置的权重。

4. 结合通道和空间注意力：
将通道注意力加权后的特征图与空间注意力加权后的特征图进行逐元素相乘，得到最终的注意力增强特征图。

5. 输出：最终的注意力增强特征图作为模块的输出，传递给下一层网络进行后续的处理，如分类、检测或分割等任务。

整个CBAM注意力机制的流程是将通道注意力和空间注意力相结合，使得网络能够自适应地调整不同通道和空间位置的重要性，从而提升模型的性能和泛化能力。

2.2新增yaml文件

关键步骤二：在 /yolov5/models/ 下新建文件 yolov5_cbam.yaml并将下面代码复制进去

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license

# 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, CBAM, [1024]],
   [-1, 1, SPPF, [1024, 5]],  # 10
  ]

# 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]],  # 14

   [-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]],  # 18 (P3/8-small)

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

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

   [[18, 21, 24], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

温馨提示：因为本文只是对yolov5s基础上添加CBAM模块，如果要对yolov5n/l/m/x进行添加则只需要修改对应的depth_multiple 和 width_multiple。

yolov5n/l/m/x对应的depth_multiple 和 width_multiple如下：

# YOLOv5n
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.25  # layer channel multiple

# YOLOv5l 
depth_multiple: 1.0  # model depth multiple
width_multiple: 1.0  # layer channel multiple

# YOLOv5m
depth_multiple: 0.67  # model depth multiple
width_multiple: 0.75  # layer channel multiple

# YOLOv5x
depth_multiple: 1.33  # model depth multiple
width_multiple: 1.25  # layer channel multiple

2.3 注册模块

关键步骤三：在yolov5/models/yolo.py中注册,大概在250行左右添加 ‘CBAM’

2.4 执行程序

在train.py中，将cfg的参数路径设置为yolov5_cbam.yaml的路径，如下图所示

建议大家写绝对路径，确保一定能找到

运行程序，如果出现下面的内容则说明添加成功🚀

我修改后的代码：链接: https://pan.baidu.com/s/1qoLGhu7t4noFMxvi7t0rzA?pwd=92im 提取码: 92im

3.总结

CBAM（Convolutional Block Attention Module）是一种用于增强卷积神经网络（CNN）性能的注意力机制。它由两个子模块组成：通道注意力模块和空间注意力模块。通道注意力模块通过全局平均池化和全连接层学习通道间的关系，并利用学到的权重对每个通道的特征图进行加权，以增强有用的特征并抑制无用的特征。空间注意力模块则通过对特征图在空间维度上进行最大池化和平均池化操作，结合两种池化结果通过全连接层学习得到每个空间位置的权重，使得网络能够更好地关注图像中的重要区域。CBAM的引入可以帮助网络更好地理解输入数据中的关键信息，从而提高了模型在各种视觉任务上的性能，如图像分类、目标检测和语义分割等。