YOLOv8改进 - 注意力篇 - 引入CBAM注意力机制

一、本文介绍

作为入门性第一篇，这里介绍了CBAM注意力在YOLOv8中的使用。包含CBAM原理分析，CBAM的代码、CBAM的使用方法、以及添加以后的yaml文件及运行记录。

二、CBAM原理分析

CBAM官方论文地址：CBAM论文

CBAM的pytorch版代码：CBAM的pytorch版代码

CBAM:卷积块注意力模块，由通道注意力和空间注意力组成。其中通道注意力机制主要检测目标的内容信息，空间注意力主要检测目标位置信息。模块先应用通道注意力，再利用空间注意力；其原理结构如下图所示。

相关代码：

在YOLOv8中，作者已经集成了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:
        """Initializes the class and sets the basic configurations and instance variables required."""
        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:
        """Applies forward pass using activation on convolutions of the input, optionally using batch normalization."""
        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):
        """Initialize CBAM with given input channel (c1) and kernel size."""
        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中CBAM使用方法

YOLOv8中CBAM模块，作者存于ultralytics/nn/modules/conv.py中。

我们使用CBAM模块，仅需在ultralytics/nn/tasks.py进行CBAM注意力机制的注册，以及在YOLOv8的yaml配置文件中添加CBAM即可。

首先打开task.py文件，按住Ctrl+F，输入parse_model进行搜索。找到parse_model函数。在其最后一个else前面添加以下注册代码：

        elif m in {CBAM,MHSA,SEAttention,ECA,ShuffleAttention,ECA_SA,SE_SA,SA_ECA,SA_ShuffleAttention,CBAM_base,PECA_SA,CPAM, CPAM_SA, MSCA_SA, SKAttention, DoubleAttention,PCPAM_SA,SA_CPAM,CoordAtt}:#自己加的注意力模块
            c1, c2 = ch[f], args[0]
            if c2 != nc:
                c2 = make_divisible(min(c2, max_channels) * width, 8)
            args = [c1, *args[1:]]

然后，就是新建一个名为YOLOv8_CBAM.yaml的配置文件：（路径:ultralytics/cfg/models/v8/YOLOv8_CBAM.yaml）

# Ultralytics YOLO 

标签：__,nn,CBAM,self,YOLOv8,注意力	

From： https://blog.csdn.net/2301_79619145/article/details/142319327