Yolov5中添加CBAM注意力机制的实现方法

CBAM注意力机制

CBAM注意力机制是一种用于计算机视觉任务的注意力机制，它通过对图像的通道和空间维度进行注意力加权来提升模型的性能。它由两个关键组件组成：通道注意力和空间注意力。

通道注意力模块主要用于对输入的通道特征图进行加权。它通过使用全局平均池化来获取全局的通道特征表示，然后经过一层全连接网络，生成通道注意力权重。这些权重用于调整每个通道的重要性，使得模型能够自动学习到不同通道之间的相关性和重要性，从而更好地利用通道信息。

空间注意力模块主要用于对输入的空间特征图进行加权。它通过使用逐通道的最大池化和平均池化来提取不同尺度的空间信息，然后经过一层全连接网络，生成空间注意力权重。这些权重用于调整每个空间位置的重要性，使得模型能够自适应地关注不同位置的细节信息，从而提升模型对空间变化的感知能力。

通过结合通道注意力和空间注意力，CBAM注意力机制能够自适应地调整不同通道和空间位置的重要性，从而提升模型对图像特征的表达能力。这使得模型能够更好地捕捉图像中的关键信息，提高任务的准确性和鲁棒性。

引入CBAM后，YOLOv5在基础骨干网的每个残差单元的最后添加了CBAM模块，并在中间层的SPP（Spatial Pyramid Pooling）层后也添加了CBAM模块。这个修改使得模型能够更好地捕捉特征之间的依赖关系，并提高特征表示能力和检测精度。此外，YOLOv5的检测头也添加了CBAM模块，从而使模型能够自适应地学习不同尺度特征之间的依赖关系，进一步提高检测精度。总之，通过引入CBAM注意力机制，YOLOv5改进了网络的基础骨干网、中间层和检测头，使得模型能够更好地捕捉特征之间的依赖关系，提高特征表示能力和检测精度。

CBAM注意力机制代码

步骤1

在common.py中添加CBAM模块

将下面的CBAM模块的代码复制粘贴到common.py文件的末尾。

# CBAM
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):
    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