Day49 - CBAM注意力机制-程序员充电站

1. 简介

CBAM (Convolutional Block Attention Module) 是一种轻量级的注意力模块，它可以无缝集成到任何CNN架构中，通过引入额外的开销来显著提升模型的性能。

与SE (Squeeze-and-Excitation) 模块主要关注通道注意力不同，CBAM 同时结合了通道注意力 (Channel Attention)和空间注意力 (Spatial Attention)。

这种串联的注意力机制使得网络能够依次学习"关注什么" (What to focus on) 和 "关注哪里" (Where to focus on)。

2. 核心原理

CBAM 包含两个子模块，通常采用串联方式连接（先通道后空间）：

2.1 通道注意力模块 (Channel Attention Module, CAM)

通道注意力旨在探索通道之间的依赖关系。CBAM 的通道注意力改进了 SE 模块：

不仅使用全局平均池化 (Average Pooling)，还引入了全局最大池化 (Max Pooling)。
认为最大池化能收集到更独特的对象特征，与平均池化互补。
两个池化后的特征向量共享同一个多层感知机 (MLP) 网络。
最终将两个输出相加并通过 Sigmoid 激活函数生成通道权重。

2.2 空间注意力模块 (Spatial Attention Module, SAM)

空间注意力旨在探索特征图在空间维度上的重要性（即哪些区域更重要）。

在通道维度上进行平均池化和最大池化，得到两个 2D 特征图。
将这两个特征图在通道维度拼接 (Concat)。
通过一个 7x7 的卷积层进行特征融合。
通过 Sigmoid 激活函数生成空间权重图。

3. 代码实现

以下是基于 PyTorch 的 CBAM 完整实现，包括通道注意力、空间注意力及其在 CNN 中的集成。

3.1 通道注意力 (ChannelAttention)

import torch import torch.nn as nn class ChannelAttention(nn.Module): def __init__(self, in_channels, ratio=16): super().__init__() # 平均池化和最大池化 self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) # 共享的全连接层 (MLP) # 使用1x1卷积代替全连接层，减少参数量并保持输入形状 self.fc = nn.Sequential( nn.Linear(in_channels, in_channels // ratio, bias=False), nn.ReLU(), nn.Linear(in_channels // ratio, in_channels, bias=False) ) self.sigmoid = nn.Sigmoid() def forward(self, x): b, c, h, w = x.shape # 平均池化分支 avg_out = self.fc(self.avg_pool(x).view(b, c)) # 最大池化分支 max_out = self.fc(self.max_pool(x).view(b, c)) # 结果相加后经过Sigmoid attention = self.sigmoid(avg_out + max_out).view(b, c, 1, 1) # 权重作用于原特征图 return x * attention

3.2 空间注意力 (SpatialAttention)

class SpatialAttention(nn.Module): def __init__(self, kernel_size=7): super().__init__() # padding计算保证输出大小不变 padding = kernel_size // 2 # 输入通道为2 (AvgPool 1 + MaxPool 1) self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): # 在通道维度上求平均 (b, 1, h, w) avg_out = torch.mean(x, dim=1, keepdim=True) # 在通道维度上求最大 (b, 1, h, w) max_out, _ = torch.max(x, dim=1, keepdim=True) # 拼接 (b, 2, h, w) pool_out = torch.cat([avg_out, max_out], dim=1) # 卷积 + Sigmoid attention = self.conv(pool_out) return x * self.sigmoid(attention)

3.3 CBAM 模块组合

class CBAM(nn.Module): def __init__(self, in_channels, ratio=16, kernel_size=7): super().__init__() self.channel_attention = ChannelAttention(in_channels, ratio) self.spatial_attention = SpatialAttention(kernel_size) def forward(self, x): # 串联结构：先通道后空间 x = self.channel_attention(x) x = self.spatial_attention(x) return x

3.4 集成到 CNN 模型

在经典的卷积神经网络中，CBAM 模块通常被放置在卷积层和激活函数之后，或者池化层之前。以下是一个简单的 CBAM-CNN 示例：

class CBAM_CNN(nn.Module): def __init__(self): super().__init__() # 第一层卷积块 self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) self.bn1 = nn.BatchNorm2d(32) self.relu1 = nn.ReLU() self.pool1 = nn.MaxPool2d(kernel_size=2) self.cbam1 = CBAM(in_channels=32) # 集成CBAM # ... 后续层省略 ... # 假设这里还有更多层 # 全连接层 self.fc1 = nn.Linear(128 * 4 * 4, 512) self.dropout = nn.Dropout(p=0.5) self.fc2 = nn.Linear(512, 10) def forward(self, x): x = self.conv1(x) x = self.bn1(x) x = self.relu1(x) x = self.pool1(x) x = self.cbam1(x) # 应用注意力机制 # ... 后续前向传播 ... x = x.view(-1, 128 * 4 * 4) x = self.fc1(x) x = self.relu1(x) # 注意这里应该是对应的激活函数 x = self.dropout(x) x = self.fc2(x) return x