CBAM卷积注意力网络结构图解析与高效实现方案

1次阅读

共计 2429 个字符，预计需要花费 7 分钟才能阅读完成。

在计算机视觉任务中，传统的卷积神经网络（CNN）通过堆叠卷积层来提取特征。然而，这种设计存在一个明显的缺陷：缺乏对特征通道和空间位置的自适应选择能力。具体来说：

所有特征通道被平等对待，无法自动区分哪些通道更重要
全局平均池化（GAP）虽然常用于分类任务，但会丢失重要的空间信息
网络难以自动聚焦于图像中最具判别性的区域

为了解决上述问题，研究人员提出了多种注意力机制：

SE 模块 （Squeeze-and-Excitation）：
仅关注通道维度注意力
通过全局平均池化获取通道权重
无法处理空间位置的重要性差异
Non-local Networks：
计算所有位置之间的相关性
计算复杂度高（O(N^2)）
难以部署到实时应用中
CBAM（Convolutional Block Attention Module）：
同时考虑通道和空间注意力
轻量级设计，计算开销小
可即插即用到现有网络架构

import torch
import torch.nn as nn

class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=8):
        super(ChannelAttention, self).__init__()
        # 共享的全连接层替代方案（使用 1D 卷积实现）self.shared_mlp = nn.Sequential(nn.Conv1d(in_planes, in_planes // ratio, 1, bias=False),  # [B, C/r, 1]
            nn.ReLU(),
            nn.Conv1d(in_planes // ratio, in_planes, 1, bias=False)   # [B, C, 1]
        )

    def forward(self, x):
        # 输入 x: [B, C, H, W]
        avg_pool = torch.mean(x, dim=[2, 3], keepdim=True)  # [B, C, 1, 1]
        max_pool = torch.max(x, dim=2, keepdim=True)[0]      # [B, C, H, 1]
        max_pool = torch.max(max_pool, dim=3, keepdim=True)[0] # [B, C, 1, 1]

        # 合并两种池化结果
        pooled = avg_pool + max_pool  # [B, C, 1, 1]

        # 1D 卷积处理（先去掉最后两个维度）channel_att = self.shared_mlp(pooled.squeeze(-1).squeeze(-1))  # [B, C]
        channel_att = channel_att.unsqueeze(-1).unsqueeze(-1)  # [B, C, 1, 1]

        # Sigmoid 激活
        return torch.sigmoid(channel_att)

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()
        assert kernel_size % 2 == 1, "核大小应为奇数"
        padding = kernel_size // 2

        self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)

    def forward(self, x):
        # 输入 x: [B, C, H, W]
        avg_pool = torch.mean(x, dim=1, keepdim=True)  # [B, 1, H, W]
        max_pool = torch.max(x, dim=1, keepdim=True)[0] # [B, 1, H, W]

        # 通道拼接
        concat = torch.cat([avg_pool, max_pool], dim=1) # [B, 2, H, W]

        # 卷积处理
        spatial_att = self.conv(concat)  # [B, 1, H, W]

        return torch.sigmoid(spatial_att)

class CBAM(nn.Module):
    def __init__(self, planes):
        super(CBAM, self).__init__()
        self.ca = ChannelAttention(planes)
        self.sa = SpatialAttention()

    def forward(self, x):
        # 通道注意力
        x = x * self.ca(x)

        # 空间注意力
        x = x * self.sa(x)

        return x

在 ResNet50 上集成 CBAM 后：