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

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

引用

CSDN

1.

https://blog.csdn.net/weixin_43752269/article/details/137374844

YOLOv8是目前最流行的实时目标检测模型之一，而CBAM（Convolutional Block Attention Module）是一种有效的注意力机制，可以进一步提升模型的检测性能。本文将详细介绍如何在YOLOv8中添加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))
  

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

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

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

然后在下图中指定位置，加入‘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:]]

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