YOLOv5改进系列（1）——添加CBAM注意力机制

一、如何理解注意力机制

假设你正在阅读一本书，同时有人在你旁边说话。当你听到某些关键字时，比如“你的名字”或者“你感兴趣的话题”，你会自动把注意力从书上转移到他们的谈话上，尽管你并没有完全忽略书本的内容。这就是注意力机制的核心思想：动态地根据重要性来分配注意力，而不是对输入的信息一视同仁地处理。

在计算机视觉或自然语言处理（NLP）中，注意力机制让模型能够灵活地聚焦于输入数据的某些重要部分。例如，在图像分类任务中，模型不需要对整张图像的每一个像素都一视同仁，而是可以专注于那些关键区域，如目标物体的边缘或特征。在句子处理时，模型可以根据句子的上下文，专注于那些对语义理解最为关键的单词或短语，而忽略不太重要的部分。

注意力机制就像在复杂信息处理中自动筛选和重点关注的过程，帮助模型更智能地选择和处理最有用的信息。

作用

提升准确性：注意力机制聚焦关键信息，提升预测精度。
增强可解释性：能更清晰展示模型决策过程。
处理变长数据：适用于文本、语音等不定长序列数据。

不足

计算开销大：需要计算每个位置的权重，耗时长。
易过拟合：复杂权重可能导致模型在训练集上表现过好，泛化能力弱。
数据需求高：需要大量数据训练，否则效果不佳。

二、CBAM注意力机制

论文名称：《CBAM: Convolutional Block Attention Module》

论文地址：https://arxiv.org/pdf/1807.06521

论文代码：GitHub - Jongchan/attention-module: Official PyTorch code for "BAM: Bottleneck Attention Module (BMVC2018)" and "CBAM: Convolutional Block Attention Module (ECCV2018)"

CBAM从通道channel和空间spatial两个作用域出发，实现从通道到空间的顺序注意力结构。空间注意力可使神经网络更加关注在图像分类中决定性作用的像素区域而忽略无关紧要的区域，通道注意力则用于处理特征图通道的分配关系，同时对两个维度进行注意力分配加强了注意力机制对模型性能的提升效果。

2.1 CAM通道注意力模块

shared MLP

输入特征图（H×W×C）：
- 先将输入的特征图（H×W×C）分别经过基于宽度和高度的最大池化和平均池化，对特征图按两个维度压缩，得到两个1×1×C的特征图。
池化后的特征图进行处理：
- 将最大池化和平均池化的结果利用共享的全连接层（Shared MLP）进行处理：
  - 先通过一个全连接层下降通道数（C -> C/4）。
  - 然后再通过另一个全连接层恢复通道数（C/4 -> C）。
生成权重：
- 将共享的全连接层所得到的结果相加后再使用Sigmoid激活函数，生成最终的channel attention feature，得到每个通道的权重值（0～1之间）。
特征调整：
- 将权重通过逐通道相加到输入特征图上，生成最终调整后的特征图。

代码如下所示：

import torch
import torch.nn as nn

class ChannelAttentionModule(nn.Module):
    def __init__(self, in_channels, reduction=4):
        super(ChannelAttentionModule, self).__init__()
        # 使用最大池化和平均池化
        self.avg_pool = nn.AdaptiveAvgPool2d(1)  # 输出形状 1x1
        self.max_pool = nn.AdaptiveMaxPool2d(1)  # 输出形状 1x1
        
        # 全连接层用于减少通道数再增加回来
        self.fc1 = nn.Conv2d(in_channels, in_channels // reduction, 1, bias=False)  # 1x1卷积，降维
        self.relu = nn.ReLU()  # 激活函数ReLU
        self.fc2 = nn.Conv2d(in_channels // reduction, in_channels, 1, bias=False)  # 1x1卷积，升维
        
        # 使用 Sigmoid 激活函数
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # 输入特征图的形状为 (B, C, H, W)
        avg_out = self.avg_pool(x)  # 平均池化
        max_out = self.max_pool(x)  # 最大池化
        
        # 使用共享的全连接层（MLP）进行处理
        avg_out = self.fc2(self.relu(self.fc1(avg_out)))
        max_out = self.fc2(self.relu(self.fc1(max_out)))
        
        # 将池化结果加起来
        out = avg_out + max_out
        
        # 使用Sigmoid激活
        out = self.sigmoid(out)
        
        # 通过逐通道相乘的方式调整输入特征图
        return x * out

# 测试模块
if __name__ == "__main__":
    # 输入张量 (batch_size, channels, height, width)
    input_tensor = torch.randn(1, 64, 32, 32)  # 示例输入 (B=1, C=64, H=32, W=32)
    
    # 实例化通道注意力模块
    cam = ChannelAttentionModule(in_channels=64)
    
    # 前向传播
    output = cam(input_tensor)
    
    print(output.shape)  # 输出特征图的形状

2.2 SAM空间注意力模块

具体流程如下：

将上面CAM模块输出的特征图F'作为本模块的输入特征图。

首先，对输入特征图在通道维度下做最大池化和平均池化，将池化后的两张特征图在通道维度堆叠（concat）。

然后，经过一个7×7卷积（7×7比3×3效果更好）操作，降维为1个channel，即叠积核融合通道信息，特征图的shape从b,2,h,w - >b,1,h,w。

最后，将卷积后的结果经过sigmoid函数对特征图的空间权重归一化，再将输入特征图和权重相乘。

class spatial_attention(nn.Module):
    def __init__(self, kernel_size=7):
        super(spatial_attention, self).__init__()

        # 为了保持卷积前后的特征图shape相同，卷积时需要padding
        padding = kernel_size // 2  # 确保7x7卷积后输出形状与输入一致

        # 7×7卷积融合通道信息，输入[b, 2, h, w]，输出[b, 1, h, w]
        self.conv = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=kernel_size, padding=padding, bias=False)

        # sigmoid激活函数
        self.sigmoid = nn.Sigmoid()

    def forward(self, inputs):
        # 在通道维度上最大池化 [b, 1, h, w]，keepdim保留原有深度
        x_maxpool, _ = torch.max(inputs, dim=1, keepdim=True)

        # 在通道维度上平均池化 [b, 1, h, w]
        x_avgpool = torch.mean(inputs, dim=1, keepdim=True)

        # 池化后的结果在通道维度上堆叠 [b, 2, h, w]
        x = torch.cat([x_maxpool, x_avgpool], dim=1)

        # 卷积融合通道信息， [b, 2, h, w] ==> [b, 1, h, w]
        x = self.conv(x)

        # 空间权重归一化
        x = self.sigmoid(x)

        # 输入特征图和空间权重相乘
        outputs = inputs * x

        return outputs

输入特征图：假设输入一个中间特征图 F∈RC×H×W，其中 C、H、W 分别表示通道数、高度和宽度。
通道注意力模块：
- 对输入特征图 F 进行全局最大池化（Max Pooling）和全局平均池化（Average Pooling）操作，生成两个不同的空间信息描述。即通过最大池化和平均池化将特征图 F 沿空间维度压缩为1维向量。
- 将池化后的结果通过一个共享的多层感知机（MLP），以融合不同空间信息并输出通道维度的权重 Mc(F)。
- 最后通过公式 (3) 计算通道注意力：

空间注意力模块：
- 将通道加权后的特征图 F′ 进一步经过空间注意力机制。首先，再次对特征图 F′ 进行全局最大池化和全局平均池化操作，但这次是在通道维度进行操作。
- 将池化结果拼接后，经过一个卷积操作生成空间注意力图 Ms(F′)，公式 (4) 表示为：

最终将空间注意力图 Ms(F′) 与特征图 F′ 按元素相乘，得到最终的输出特征图 F′′。

CBAM模块的完整过程：
- 通道注意力模块和空间注意力模块可以串联使用，先通过通道注意力调整特征，再通过空间注意力调整，完成特征图的两次加权。

三、CBAM注意力机制添加过程

1.在common.py中添加网络结构

将下面代码复制到common.py文件最下面

class ChannelAttentionModule(nn.Module):
    def __init__(self, in_channels, reduction=4):
        super(ChannelAttentionModule, self).__init__()
        # 使用最大池化和平均池化
        self.avg_pool = nn.AdaptiveAvgPool2d(1)  # 输出形状 1x1
        self.max_pool = nn.AdaptiveMaxPool2d(1)  # 输出形状 1x1

        # 全连接层用于减少通道数再增加回来
        self.fc1 = nn.Conv2d(in_channels, in_channels // reduction, 1, bias=False)  # 1x1卷积，降维
        self.relu = nn.ReLU()  # 激活函数ReLU
        self.fc2 = nn.Conv2d(in_channels // reduction, in_channels, 1, bias=False)  # 1x1卷积，升维

        # 使用 Sigmoid 激活函数
        self.sigmoid = nn.Sigmoid()
    def forward(self, x):
        b, c, h, w = x.size()  # 获取输入的形状 (batch_size, channels, height, width)

        # 最大池化和平均池化，输出维度 [b, c, 1, 1]
        max_pool = self.max_pool(x).view(b, c)  # 调整池化结果维度为 [b, c]
        avg_pool = self.avg_pool(x).view(b, c)  # 调整池化结果维度为 [b, c]

        # 第一个全连接层降通道数 [b, c] => [b, c/4]
        x_maxpool = self.fc1(max_pool)
        x_avgpool = self.fc1(avg_pool)

        # 激活函数
        x_maxpool = self.relu(x_maxpool)
        x_avgpool = self.relu(x_avgpool)

        # 第二个全连接层恢复通道数 [b, c/4] => [b, c]
        x_maxpool = self.fc2(x_maxpool)
        x_avgpool = self.fc2(x_avgpool)

        # 将最大池化和平均池化的结果相加 [b, c]
        x = x_maxpool + x_avgpool

        # Sigmoid函数权重归一化
        x = self.sigmoid(x)

        # 调整维度 [b, c] => [b, c, 1, 1]
        x = x.view(b, c, 1, 1)

        # 输入特征图和通道权重相乘 [b, c, h, w]
        outputs = x * x

        return outputs
    # def forward(self, x):
    #     # 输入特征图的形状为 (B, C, H, W)
    #     b,c,h,w = x.shape
    #     max_out = self.max_pool(x)  # 最大池化
    #     avg_out = self.avg_pool(x)  # 平均池化
    #
    #     # 使用共享的全连接层（MLP）进行处理
    #     avg_out = self.fc2(self.relu(self.fc1(avg_out)))
    #     max_out = self.fc2(self.relu(self.fc1(max_out)))
    #
    #     # 将池化结果加起来
    #     out = avg_out + max_out
    #
    #     # 使用Sigmoid激活
    #     out = self.sigmoid(out)
    #
    #     # 通过逐通道相乘的方式调整输入特征图
    #     return x * out


class spatial_attention(nn.Module):
    def __init__(self, kernel_size=7):
        super(spatial_attention, self).__init__()

        # 为了保持卷积前后的特征图shape相同，卷积时需要padding
        padding = kernel_size // 2  # 确保7x7卷积后输出形状与输入一致

        # 7×7卷积融合通道信息，输入[b, 2, h, w]，输出[b, 1, h, w]
        self.conv = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=kernel_size, padding=padding, bias=False)

        # sigmoid激活函数
        self.sigmoid = nn.Sigmoid()

    def forward(self, inputs):
        # 在通道维度上最大池化 [b, 1, h, w]，keepdim保留原有深度
        x_maxpool, _ = torch.max(inputs, dim=1, keepdim=True)

        # 在通道维度上平均池化 [b, 1, h, w]
        x_avgpool = torch.mean(inputs, dim=1, keepdim=True)

        # 池化后的结果在通道维度上堆叠 [b, 2, h, w]
        x = torch.cat([x_maxpool, x_avgpool], dim=1)

        # 卷积融合通道信息， [b, 2, h, w] ==> [b, 1, h, w]
        x = self.conv(x)

        # 空间权重归一化
        x = self.sigmoid(x)

        # 输入特征图和空间权重相乘
        outputs = inputs * x

        return outputs


class CBAM(nn.Module):
    def __init__(self, c1, c2, ratio=16, kernel_size=7):  # ch_in
        super(CBAM, self).__init__()
        self.channel_attention = ChannelAttentionModule(c1, ratio)
        self.spatial_attention = spatial_attention(kernel_size)

    def forward(self, x):
        out = self.channel_attention(x) * x  # 通道注意力加权
        # c*h*w (通道数、高度、宽度)
        out = self.spatial_attention(out) * out  # 空间注意力加权
        # c*h*w * 1*h*w（空间维度权重）
        return out