别再死记公式了！用PyTorch手写SENet和CBAM，5分钟搞懂通道与空间注意力

从零实现SENet与CBAM：用PyTorch代码拆解注意力机制的本质

在计算机视觉领域，注意力机制已经成为提升模型性能的关键组件。但很多初学者在理解通道注意力和空间注意力时，常常陷入公式推导的泥潭而忽略了其工程实现的本质。本文将带你用PyTorch从零实现两种经典注意力模块——SENet（通道注意力）和CBAM（混合注意力），通过代码层面的拆解，直观感受神经网络"关注什么"（What）和"关注哪里"（Where）的差异。

1. 注意力机制的核心思想

注意力机制的本质是让神经网络学会"选择性聚焦"。想象人类观察一幅画时，会自然地关注重要区域而忽略背景——这正是注意力机制要模拟的认知过程。在深度学习中，这种机制通过权重分配来实现：

• 通道注意力（如SENet）：决定"哪些特征通道更重要"
• 空间注意力（如CBAM中的SAM）：决定"特征图的哪些空间位置更重要"

# 伪代码展示注意力机制的核心操作 def attention_mechanism(features): # 生成注意力权重（范围0-1） attention_weights = generate_weights(features) # 特征图与权重逐元素相乘 return features * attention_weights

提示：注意力权重不是预先设定的，而是通过子网络从数据中学习得到的，这正是其强大之处

2. 实现SENet通道注意力模块

SENet（Squeeze-and-Excitation Network）是通道注意力的经典实现，其核心分为三步：

1. Squeeze：全局平均池化压缩空间维度
2. Excitation：全连接层学习通道间关系
3. Scale：权重与原始特征相乘

2.1 完整PyTorch实现

import torch import torch.nn as nn class SEBlock(nn.Module): def __init__(self, channels, reduction=16): super().__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Sequential( nn.Linear(channels, channels // reduction), nn.ReLU(inplace=True), nn.Linear(channels // reduction, channels), nn.Sigmoid() ) def forward(self, x): b, c, _, _ = x.size() # Squeeze y = self.avg_pool(x).view(b, c) # Excitation y = self.fc(y).view(b, c, 1, 1) # Scale return x * y.expand_as(x)

2.2 关键实现细节解析

1. 降维比例选择：

   • reduction参数控制中间层维度（通常取16）
   • 过大导致信息损失，过小则参数量剧增

2. 池化操作对比：

   池化类型	计算方式	特点
   全局平均池化	取每个通道平均值	稳定但可能平滑过度
   全局最大池化	取每个通道最大值	突出显著特征但易受噪声影响

3. 常见问题排查：

   • 维度不匹配：确保view操作与张量形状一致
   • 梯度消失：检查Sigmoid输出是否饱和（可尝试替换为Hard-Sigmoid）

注意：SEBlock的输出维度与输入完全相同，可以无缝嵌入任何CNN架构

3. 实现CBAM混合注意力模块

CBAM（Convolutional Block Attention Module）创新性地将通道注意力和空间注意力串联，形成更强大的混合注意力机制。

3.1 通道注意力模块改进

CBAM的通道注意力在SENet基础上增加了并行分支：

class ChannelAttention(nn.Module): def __init__(self, channels, reduction=16): super().__init__() self.max_pool = nn.AdaptiveMaxPool2d(1) self.avg_pool = nn.AdaptiveAvgPool2d(1) self.mlp = nn.Sequential( nn.Conv2d(channels, channels//reduction, 1), nn.ReLU(), nn.Conv2d(channels//reduction, channels, 1) ) self.sigmoid = nn.Sigmoid() def forward(self, x): max_out = self.mlp(self.max_pool(x)) avg_out = self.mlp(self.avg_pool(x)) return self.sigmoid(max_out + avg_out)

3.2 空间注意力模块实现

空间注意力关注"在哪里"的问题：

class SpatialAttention(nn.Module): def __init__(self, kernel_size=7): super().__init__() self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2) self.sigmoid = nn.Sigmoid() def forward(self, x): max_out, _ = torch.max(x, dim=1, keepdim=True) avg_out = torch.mean(x, dim=1, keepdim=True) combined = torch.cat([max_out, avg_out], dim=1) return self.sigmoid(self.conv(combined))

3.3 完整CBAM集成

class CBAM(nn.Module): def __init__(self, channels, reduction=16, kernel_size=7): super().__init__() self.channel_att = ChannelAttention(channels, reduction) self.spatial_att = SpatialAttention(kernel_size) def forward(self, x): x = x * self.channel_att(x) x = x * self.spatial_att(x) return x

4. 可视化分析与实战技巧

4.1 注意力权重可视化

理解注意力机制最直观的方式是可视化其生成的权重：

import matplotlib.pyplot as plt def visualize_attention(model, input_tensor): # 获取通道注意力权重 channel_weights = model.channel_att(input_tensor) # 获取空间注意力权重 spatial_weights = model.spatial_att(input_tensor) plt.figure(figsize=(12,4)) plt.subplot(131) plt.imshow(input_tensor[0,0].cpu().detach(), cmap='gray') plt.title('Input Feature') plt.subplot(132) plt.imshow(channel_weights[0,0].cpu().detach(), cmap='hot') plt.title('Channel Attention') plt.subplot(133) plt.imshow(spatial_weights[0,0].cpu().detach(), cmap='hot') plt.title('Spatial Attention') plt.show()

4.2 模型嵌入实践指南

将注意力模块嵌入现有架构时需考虑：

1. 插入位置：

   • 通常在卷积块之后插入
   • ResNet中可放在残差连接前

2. 计算开销控制：

   • 通道降维比例合理设置
   • 大模型中使用更经济的注意力变体

3. 训练技巧：

   • 初始阶段可冻结注意力模块
   • 配合学习率warmup策略

4.3 性能对比实验

在CIFAR-10上的对比实验结果：

5. 进阶应用与优化方向

5.1 轻量化注意力设计

针对移动设备的优化方案：

class EfficientChannelAttention(nn.Module): """ 使用1D卷积替代全连接层 """ def __init__(self, channels, gamma=2, b=1): super().__init__() t = int(abs((math.log2(channels) + b) / gamma)) k = t if t % 2 else t + 1 self.avg_pool = nn.AdaptiveAvgPool2d(1) self.conv = nn.Conv1d(1, 1, kernel_size=k, padding=k//2) self.sigmoid = nn.Sigmoid() def forward(self, x): y = self.avg_pool(x) y = self.conv(y.squeeze(-1).transpose(-1,-2)) y = y.transpose(-1,-2).unsqueeze(-1) return x * self.sigmoid(y)

5.2 注意力机制组合策略

不同注意力模块的组合方式对比：

1. 串行组合（CBAM方式）：

   输入 → 通道注意力 → 空间注意力 → 输出

2. 并行组合：

   # 并行处理后再融合 channel_out = channel_att(x) spatial_out = spatial_att(x) return x * channel_out * spatial_out

3. 混合组合：

   • 深层网络使用串行
   • 浅层网络使用并行

5.3 跨模态注意力扩展

注意力机制可自然扩展到多模态场景：

class CrossModalAttention(nn.Module): def __init__(self, channels): super().__init__() self.query = nn.Conv2d(channels, channels//8, 1) self.key = nn.Conv2d(channels, channels//8, 1) self.value = nn.Conv2d(channels, channels, 1) def forward(self, x1, x2): # x1和x2是不同模态的特征 q = self.query(x1) k = self.key(x2) v = self.value(x2) attn = torch.softmax((q @ k.transpose(-2,-1)) / math.sqrt(q.size(1)), dim=-1) return attn @ v

在实际项目中，注意力模块的调试往往需要结合具体任务特点。例如在图像分割中，空间注意力的效果通常比通道注意力更显著；而在细粒度分类任务中，二者结合往往能带来最大收益。