YOLOv8添加CBAM注意力机制(模型改进)_yolo注意力cbam

一、CBAM注意力机制

CBAM（Convolutional Block Attention Module）是一种用于增强卷积神经网络（CNN）性能的注意力机制模块。它由Sanghyun Woo等人在2018年的论文（https://arxiv.org/abs/1807.06521）中提出。CBAM的主要目标是通过在CNN中引入通道注意力和空间注意力来提高模型的感知能力，从而在不增加网络复杂性的情况下改善性能。

CBAM由两个关键部分组成：通道注意力模块（C-channel）和空间注意力模块（S-channel）。这两个模块可以分别嵌入到CNN中的不同层，以增强特征表示。

二、通道注意力模块

通道注意力模块的目标是增强每个通道的特征表达

三、空间注意力模块

空间注意力模块的目标是强调图像中不同位置的重要性

四、CBAM模块

CBAM就是将通道注意力模块和空间注意力模块的输出特征逐元素相乘，得到最终的注意力增强特征。这个增强的特征将用作后续网络层的输入，以在保留关键信息的同时，抑制噪声和无关信息。原文实验证明先进行通道维度的整合，再进行空间维度的整合，模型效果更好

五、如何在YOLOv8中添加CBAM注意力模块

CBAM模块代码(pytorch)：

#CBAM注意力
class ChannelAttention(nn.Module):
    """Channel-attention module https://github.com/open-mmlab/mmdetection/tree/v3.0.0rc1/configs/rtmdet."""

    def __init__(self, channels: int) -> None:
        super().__init__()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Conv2d(channels, channels, 1, 1, 0, bias=True)
        self.act = nn.Sigmoid()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * self.act(self.fc(self.pool(x)))


class SpatialAttention(nn.Module):
    """Spatial-attention module."""

    def __init__(self, kernel_size=7):
        """Initialize Spatial-attention module with kernel size argument."""
        super().__init__()
        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1
        self.cv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.act = nn.Sigmoid()

    def forward(self, x):
        """Apply channel and spatial attention on input for feature recalibration."""
        return x * self.act(self.cv1(torch.cat([torch.mean(x, 1, keepdim=True), torch.max(x, 1, keepdim=True)[0]], 1)))

class CBAM(nn.Module):
    """Convolutional Block Attention Module."""

    def __init__(self, c1, kernel_size=7):  # ch_in, kernels
        super().__init__()
        self.channel_attention = ChannelAttention(c1)
        self.spatial_attention = SpatialAttention(kernel_size)

    def forward(self, x):
        """Applies the forward pass through C1 module."""
        return self.spatial_attention(self.channel_attention(x))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

首先将上述代码复制，粘贴到YOLOv8项目的 ultralytics-main/ultralytics/nn/modules/conv.py 中;

然后将下面内容加到conv.py指定位置中，

'ChannelAttention', 'SpatialAttention', 'CBAM',
1

然后在下图中指定位置，加入‘CBAM’的引用。

然后，在task.py中指定位置import导入CBAM模块；

在task.py中，加入以下代码：

        elif m is CBAM:
            c1,c2 = ch[f],args[0]
            if c2 != nc:
                c2 = make_divisible(min(c2, max_channels) * width, 8)
            args = [c1, *args[1:]]
1
2
3
4
5

至此，CBAM就成功添加到了YOLOv8中。如果想要使用，就可以在yolov8.yaml中添加使用。