PyTorch实战：手把手教你给U-Net加上CBAM注意力模块（附完整代码）

PyTorch实战：手把手教你给U-Net加上CBAM注意力模块（附完整代码）

在医学图像分割领域，U-Net凭借其独特的编码器-解码器结构和跳跃连接，一直是许多研究者的首选架构。但当我们面对复杂场景时，原始U-Net可能会忽略一些关键特征。这时，注意力机制就像给模型装上了"智能聚光灯"，让网络学会聚焦重要区域。本文将带你从零实现CBAM模块的集成，通过残差连接方式让注意力机制真正实现"即插即用"。

1. 环境准备与基础理解

首先确保你的环境已安装PyTorch 1.8+和torchvision。推荐使用Python 3.8+环境，可以通过以下命令快速验证：

python -c "import torch; print(torch.__version__)"

CBAM（Convolutional Block Attention Module）由两个子模块组成：

• 通道注意力：学习不同特征通道的重要性权重
• 空间注意力：关注特征图的空间位置关系

二者的结合方式不是简单的串联，而是先通道后空间的级联结构。这种设计在ImageNet分类任务上已经证明了其有效性，但我们更关心它在分割任务中的表现。

提示：虽然原论文建议先通道后空间，但在实际分割任务中，有些开发者尝试调换顺序，效果可能因数据集而异

2. CBAM模块的PyTorch实现

我们先拆解CBAM的核心代码。新建一个cbam.py文件，写入以下完整实现：

import torch import torch.nn as nn class ChannelAttention(nn.Module): def __init__(self, in_planes, ratio=16): super().__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) self.fc = nn.Sequential( nn.Conv2d(in_planes, in_planes//ratio, 1, bias=False), nn.ReLU(), nn.Conv2d(in_planes//ratio, in_planes, 1, bias=False) ) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = self.fc(self.avg_pool(x)) max_out = self.fc(self.max_pool(x)) out = avg_out + max_out return self.sigmoid(out) class SpatialAttention(nn.Module): def __init__(self, kernel_size=7): super().__init__() assert kernel_size in (3,7), "kernel size must be 3 or 7" padding = 3 if kernel_size == 7 else 1 self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): 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) x = self.conv(x) return self.sigmoid(x) class CBAM(nn.Module): def __init__(self, channels, reduction_ratio=16): super().__init__() self.ca = ChannelAttention(channels, reduction_ratio) self.sa = SpatialAttention() def forward(self, x): x = self.ca(x) * x # 通道注意力加权 x = self.sa(x) * x # 空间注意力加权 return x

关键改进点说明：

1. 将原论文的共享MLP改为更直观的fc命名
2. 为空间注意力添加了kernel_size可选参数
3. 使用更规范的变量命名(in_planes替代channel)

3. U-Net架构改造实战

现在我们改造标准的U-Net。关键是在下采样路径的每个阶段后插入CBAM模块，这里采用残差连接来保留原始特征：

class DoubleConv(nn.Module): """(convolution => [BN] => ReLU) * 2""" def __init__(self, in_channels, out_channels): super().__init__() self.double_conv = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True), nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True) ) def forward(self, x): return self.double_conv(x) class DownWithCBAM(nn.Module): """下采样模块，包含MaxPool和双卷积，后接CBAM""" def __init__(self, in_channels, out_channels): super().__init__() self.maxpool_conv = nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_channels, out_channels) ) self.cbam = CBAM(out_channels) def forward(self, x): features = self.maxpool_conv(x) attended = self.cbam(features) return features + attended # 残差连接

完整的U-Net架构整合如下：

class UNetWithCBAM(nn.Module): def __init__(self, n_channels=3, n_classes=2): super().__init__() # 下采样路径 self.inc = DoubleConv(n_channels, 64) self.down1 = DownWithCBAM(64, 128) self.down2 = DownWithCBAM(128, 256) self.down3 = DownWithCBAM(256, 512) self.down4 = DownWithCBAM(512, 1024) # 上采样路径 self.up1 = Up(1024, 512) self.up2 = Up(512, 256) self.up3 = Up(256, 128) self.up4 = Up(128, 64) self.outc = OutConv(64, n_classes) def forward(self, x): x1 = self.inc(x) x2 = self.down1(x1) x3 = self.down2(x2) x4 = self.down3(x3) x5 = self.down4(x4) x = self.up1(x5, x4) x = self.up2(x, x3) x = self.up3(x, x2) x = self.up4(x, x1) return self.outc(x)

4. 训练技巧与效果验证

在实际训练时，有几个关键点需要注意：

1. 学习率调整：
   • CBAM模块引入的新参数需要适当的学习率
   • 推荐使用分层学习率策略

optimizer = torch.optim.Adam([ {'params': model.inc.parameters(), 'lr': 1e-4}, {'params': model.down1.parameters(), 'lr': 1e-4}, {'params': model.cbam_modules.parameters(), 'lr': 5e-4} # CBAM参数使用更大学习率 ], weight_decay=1e-5)

2. 注意力可视化： 添加以下代码可视化注意力权重：

def visualize_attention(self, x): # 获取各层CBAM权重 ca_weights = [] sa_weights = [] x1 = self.inc(x) x2, ca1, sa1 = self.down1(x1) ca_weights.append(ca1) sa_weights.append(sa1) # ... 其他层类似 return ca_weights, sa_weights

3. 效果对比指标：

从实验数据可以看出，CBAM带来了明显的精度提升，而计算开销增加有限。特别是在小目标分割场景下，改进更为显著。