Attention UNet在医学图像分割中的实践与PyTorch实现解析

吃瓜不吐籽595

470人浏览 · 2026-03-08 01:32:00

吃瓜不吐籽595 · 2026-03-08 01:32:00 发布

1. 从“看全”到“看准”：为什么医学图像分割需要注意力

大家好，我是老张，在AI和医疗影像这个行当里摸爬滚打了十来年。今天咱们不聊那些虚头巴脑的概念，就聊聊一个实实在在、能帮你把模型效果提上去的技术——Attention UNet。如果你正在做医学图像分割，比如从CT里勾画肿瘤，从MRI里分割器官，感觉模型总是“看”得不够准，边界糊糊的，那这篇文章可能就是你要找的“解药”。

咱们先想一个场景：医生看一张肺部CT片，他不会把整张图每个像素都同等用力地看。他的眼睛会先扫过整个图像，然后注意力自然而然地就聚焦到那些有结节、纹理异常的区域，仔细分析它们的边界、密度。这种“聚焦”能力，就是人类视觉系统中的注意力机制。传统的UNet模型，虽然通过跳跃连接把浅层细节和深层语义信息结合得很好，但它有个小缺点：它在融合不同层级特征时，是“一视同仁”的。浅层特征（比如器官边缘）里有很多有用的细节，但也混杂着大量无关的背景噪声。模型把这些信息全盘接收，结果就是学到的特征不够“纯净”，分割出来的边界自然就容易毛糙。

Attention UNet干了一件特别聪明的事：它给UNet的跳跃连接加了一个“智能开关”，也就是注意力门（Attention Gate）。这个门的作用，就是像医生的眼睛一样，去评估跳跃连接传过来的每一个特征图区域的重要性。对于当前分割任务有用的区域（比如肿瘤的边缘），它就放大信号，让这些特征“通过”；对于那些无关的背景或噪声区域，它就减弱甚至关闭信号。这么一来，解码器在重建高分辨率分割图时，接收到的都是“精挑细选”过的、富含任务相关信息的特征，分割精度，尤其是边界的精细度，就能得到显著提升。

我最早在胰腺肿瘤分割项目里试过这个思路。用普通UNet，模型经常把一些血管影或者邻近组织的边缘误判为肿瘤边界，导致分割结果“溢出去”一点。加了注意力机制之后，模型明显“学乖了”，它更懂得去关注肿瘤区域特有的纹理和对比度变化，分割结果和医生手工勾画的吻合度（Dice系数）直接涨了差不多5个百分点。这个提升在医疗场景里非常宝贵，因为哪怕边界优化一两个像素，都可能影响后续的体积计算和疗效评估。所以，Attention UNet不是什么炫技的玩具，而是一个能解决实际痛点的实用工具。

2. 庖丁解牛：Attention Gate是如何工作的

光说它好没用，咱们得把它拆开，看看里面的齿轮是怎么转的。理解了原理，调参和魔改的时候心里才有底。上面那张结构图很经典，但咱们用更“人话”的方式捋一遍。

我们把需要被“筛选”的特征叫做 x，它来自编码器的浅层，比如第2、3层，分辨率高，包含丰富的空间细节（边缘、纹理），但也可能有很多噪声。把作为“指导”或“上下文”的特征叫做 g，它来自编码器更深层（或者解码器的上一级），分辨率低，但语义信息强，大概“知道”要分割的目标在哪个区域。

注意力门的计算就像一场小型会议：

统一频道（1x1卷积）：首先，x和g可能通道数不一样，直接开会鸡同鸭讲。所以分别给它们过一个1x1卷积（通常跟着BatchNorm），把通道数统一到一个中间值（比如int_channels）。这个操作不改变特征图的高和宽，只是做一下特征变换，让它们能在同一个“频道”里交流。
对齐信息（上采样）：g的特征图尺寸比较小，需要把它上采样（比如用双线性插值）到和x一样的大小，这样才能进行逐元素的加减操作。
融合与激活（相加+ReLU）：把变换后的x和上采样后的g逐元素相加。这一步可以理解为让全局的语义信息（g）去“点亮”局部细节信息（x）中相关的部分。相加的结果通过一个ReLU激活函数，引入非线性。
生成注意力地图（1x1卷积 + Sigmoid）：将上一步的结果再通过一个1x1卷积（通常通道数压缩为1），接一个Sigmoid函数。这个操作会生成一张和x空间尺寸相同、但只有1个通道的“注意力图”。这张图里的每一个值都在0到1之间，你可以把它理解为一个“重要性权重”。越接近1，表示对应位置的特征越重要；越接近0，则表示越无关或可能是噪声。
应用权重（逐元素乘法）：最后，用这张注意力图去乘以原始的输入特征x（注意，是乘原始的x，不是变换后的）。这就完成了筛选：重要的特征被保留甚至增强，不重要的特征被抑制。

这个过程的核心思想是让模型自己学会“看图说话”，根据高级语义线索（g）来动态地调整对低级细节（x）的依赖程度。它不是我们人工设定的规则，而是从数据中学出来的。在实际的CT肺结节分割中，你会发现模型学到的注意力图，真的会高亮结节区域及其边缘，而抑制均匀的肺实质部分，非常直观。

3. 手把手搭建：用PyTorch实现Attention UNet

理论懂了，不敲代码都是纸上谈兵。下面我结合自己踩过的坑，给你一个更健壮、更易用的PyTorch实现。我们会从最基础的模块开始搭建。

3.1 搭建积木：注意力块与卷积块

首先，我们实现一个万金油的卷积块，它包含卷积、批归一化和激活函数，这是所有CNN的基础组件。

import torch
import torch.nn as nn
import torch.nn.functional as F

class ConvBlock(nn.Module):
    """一个简单的卷积块：Conv2d -> BatchNorm2d -> ReLU，重复两次。"""
    def __init__(self, in_channels, out_channels):
        super(ConvBlock, self).__init__()
        # 这里我习惯用3x3卷积，padding=1保证尺寸不变（stride=1时）
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True), # 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)

接下来是核心的注意力块（AttentionBlock）。这里的实现我优化了一下，更贴近原论文，并且加了详细的注释。

class AttentionBlock(nn.Module):
    """
    注意力门机制。
    参数：
        F_g (int): 门控信号g的通道数（来自深层）。
        F_l (int): 跳跃连接特征x的通道数（来自浅层）。
        F_int (int): 中间表示的通道数，通常取F_g或F_l的一半。
    """
    def __init__(self, F_g, F_l, F_int):
        super(AttentionBlock, self).__init__()

        # 用于门控信号g的权重层，1x1卷积降维到F_int
        self.W_g = nn.Sequential(
            nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(F_int)
        )

        # 用于跳跃连接特征x的权重层，1x1卷积降维到F_int
        self.W_x = nn.Sequential(
            nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(F_int)
        )

        # 用于生成注意力图的Psi函数：卷积+激活
        # 输出通道为1，经过Sigmoid得到0-1的注意力系数
        self.psi = nn.Sequential(
            nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(1),
            nn.Sigmoid()
        )

        # 激活函数，论文中使用ReLU
        self.relu = nn.ReLU(inplace=True)

    def forward(self, g, x):
        """
        前向传播。
        参数：
            g (Tensor): 门控信号，通常来自解码器深层，尺寸较小。
            x (Tensor): 跳跃连接特征，来自编码器浅层，尺寸较大。
        返回：
            out (Tensor): 经过注意力加权的特征，尺寸与x相同。
        """
        # 1. 对g进行权重变换
        g1 = self.W_g(g)
        # 2. 对x进行权重变换
        x1 = self.W_x(x)
        # 3. 将g上采样到x的尺寸，然后相加
        # 使用双线性插值，更平滑。注意对齐角点参数。
        g1_up = F.interpolate(g1, size=x1.shape[2:], mode='bilinear', align_corners=False)
        # 4. 相加并通过ReLU激活
        psi = self.relu(g1_up + x1)
        # 5. 通过Psi函数生成注意力图
        psi = self.psi(psi)
        # 6. 将注意力图应用于原始特征x（注意是x，不是x1）
        return x * psi

3.2 组装升级模块：带注意力的上采样块

有了注意力块，我们就可以构建解码器中的上采样模块了。这个模块需要完成：上采样、与跳跃连接做注意力融合、卷积细化。

class AttentionUpBlock(nn.Module):
    """解码器中的上采样块，集成了注意力机制。"""
    def __init__(self, in_channels, out_channels):
        """
        参数：
            in_channels: 输入通道数（来自上一解码层或瓶颈层）。
            out_channels: 输出通道数，也是跳跃连接的通道数。
        """
        super(AttentionUpBlock, self).__init__()
        # 上采样方式：转置卷积。也可以选择双线性插值+卷积，这里按论文来。
        self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)
        # 注意力模块，注意参数顺序：F_g对应来自上一层的特征，F_l对应跳跃连接特征
        self.att = AttentionBlock(F_g=in_channels, F_l=out_channels, F_int=out_channels//2)
        # 注意力融合后的卷积块
        # 输入通道：注意力加权的跳跃连接(out_channels) + 上采样后的特征(out_channels)
        self.conv = ConvBlock(out_channels * 2, out_channels)

    def forward(self, x, skip):
        """
        参数：
            x: 来自解码器上一层的输入。
            skip: 来自编码器的跳跃连接。
        """
        # 1. 对x进行上采样
        x_up = self.up(x)
        # 2. 应用注意力机制到跳跃连接上，x_up作为门控信号
        skip_att = self.att(g=x_up, x=skip)
        # 3. 拼接注意力加权的跳跃连接和上采样特征
        # 这里顺序很重要，通常把跳跃连接放前面
        x = torch.cat([skip_att, x_up], dim=1)
        # 4. 通过卷积块进行特征融合与细化
        return self.conv(x)

3.3 构建完整网络：Attention UNet

现在，我们把编码器、瓶颈层和解码器组装起来。编码器部分就是经典的向下采样+卷积。

class AttentionUNet(nn.Module):
    def __init__(self, in_channels=1, out_channels=1, features=[64, 128, 256, 512]):
        """
        参数：
            in_channels: 输入图像通道数，灰度图为1，RGB为3。
            out_channels: 输出通道数，即分割类别数。二分类通常为1。
            features: 编码器各层的基础特征通道数，列表长度决定网络深度。
        """
        super(AttentionUNet, self).__init__()
        self.downs = nn.ModuleList()
        self.ups = nn.ModuleList()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        # 构建编码器 (Down path)
        for feature in features:
            self.downs.append(ConvBlock(in_channels, feature))
            in_channels = feature

        # 瓶颈层 (Bottleneck)
        self.bottleneck = ConvBlock(features[-1], features[-1] * 2)

        # 构建解码器 (Up path)，注意要倒序
        for feature in reversed(features):
            # 上采样块
            self.ups.append(AttentionUpBlock(feature * 2, feature))
            # 输出卷积块（每个上采样块后已经有一个ConvBlock，这里不需要额外再加）

        # 最终输出层，使用1x1卷积将通道数映射到类别数
        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)

    def forward(self, x):
        skip_connections = []

        # 编码器前向传播
        for down in self.downs:
            x = down(x)
            skip_connections.append(x) # 保存跳跃连接
            x = self.pool(x)

        # 瓶颈层
        x = self.bottleneck(x)

        # 解码器前向传播，注意跳跃连接要倒序使用
        skip_connections = skip_connections[::-1] # 反转列表
        for idx, up in enumerate(self.ups):
            skip = skip_connections[idx]
            # 如果尺寸不匹配（由于池化舍入等），调整跳跃连接尺寸
            if x.shape != skip.shape:
                x = F.interpolate(x, size=skip.shape[2:], mode='bilinear', align_corners=False)
            x = up(x, skip)

        return torch.sigmoid(self.final_conv(x)) # 二分类用Sigmoid

这个实现里我注意了几个细节：一是用ModuleList管理子模块更规范；二是在解码器开始前反转了跳跃连接列表，逻辑更清晰；三是在forward里加了尺寸对齐的检查，避免因为输入尺寸不是2的整数次幂而导致的拼接错误，这在实际项目中很常见。

4. 实战演练：训练你的第一个Attention UNet模型

模型搭好了，怎么把它用起来呢？咱们用一个模拟的医学图像二分类任务（比如分割病灶）来走一遍流程。我会重点讲数据准备、损失函数选择和训练技巧。

4.1 准备数据：构造一个简单的数据加载器

医学数据敏感，我们这里用随机数模拟。真实项目中，你需要替换成自己的Dataset类，读取nii.gz或dicom文件。

from torch.utils.data import Dataset, DataLoader
import numpy as np

class SimMedicalDataset(Dataset):
    """模拟医学图像数据集，生成随机图像和掩码。"""
    def __init__(self, num_samples=100, img_size=256):
        self.num_samples = num_samples
        self.img_size = img_size
        # 模拟一些随机“病灶”
        self.centers = np.random.randint(50, 200, size=(num_samples, 2))
        self.radii = np.random.randint(10, 30, size=num_samples)

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        # 生成随机背景噪声图像
        image = np.random.randn(self.img_size, self.img_size).astype(np.float32) * 0.1
        mask = np.zeros((self.img_size, self.img_size), dtype=np.float32)

        # 在图像上画一个“圆形病灶”，像素值更高
        center_y, center_x = self.centers[idx]
        radius = self.radii[idx]
        y, x = np.ogrid[:self.img_size, :self.img_size]
        dist_from_center = np.sqrt((x - center_x)**2 + (y - center_y)**2)
        circle_mask = dist_from_center <= radius
        image[circle_mask] += 0.5  # 病灶区域信号更强
        mask[circle_mask] = 1.0    # 对应的掩码为1

        # 增加通道维度，模拟单通道灰度图
        image = np.expand_dims(image, axis=0)  # [1, H, W]
        mask = np.expand_dims(mask, axis=0)    # [1, H, W]

        return torch.from_numpy(image), torch.from_numpy(mask)

# 创建数据加载器
batch_size = 4
train_dataset = SimMedicalDataset(num_samples=160)
val_dataset = SimMedicalDataset(num_samples=40)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

4.2 选择损失函数与优化器

医学图像分割中，正负样本（病灶和背景）通常极不平衡。直接用BCE Loss效果可能不好。我常用的组合是 Dice Loss + BCE Loss，这个组合在实践中非常稳健。

def dice_loss(pred, target, smooth=1e-6):
    """计算Dice Loss。"""
    pred = pred.contiguous().view(-1)
    target = target.contiguous().view(-1)
    intersection = (pred * target).sum()
    dice = (2. * intersection + smooth) / (pred.sum() + target.sum() + smooth)
    return 1 - dice

def bce_dice_loss(pred, target):
    """混合损失：BCE Loss + Dice Loss。"""
    bce = F.binary_cross_entropy(pred, target)
    dice = dice_loss(pred, target)
    return bce + dice  # 可以加权重，如 0.5*bce + 0.5*dice

然后初始化模型、优化器和学习率调度器。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = AttentionUNet(in_channels=1, out_channels=1).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=5, factor=0.5, verbose=True)

4.3 编写训练与验证循环

这里给出一个简洁但完整的训练 epoch 函数。

def train_epoch(model, loader, optimizer, criterion, device):
    model.train()
    running_loss = 0.0
    for images, masks in loader:
        images, masks = images.to(device), masks.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, masks)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * images.size(0)
    return running_loss / len(loader.dataset)

def validate_epoch(model, loader, criterion, device):
    model.eval()
    running_loss = 0.0
    with torch.no_grad():
        for images, masks in loader:
            images, masks = images.to(device), masks.to(device)
            outputs = model(images)
            loss = criterion(outputs, masks)
            running_loss += loss.item() * images.size(0)
    return running_loss / len(loader.dataset)

# 开始训练
num_epochs = 50
best_val_loss = float('inf')
for epoch in range(num_epochs):
    train_loss = train_epoch(model, train_loader, optimizer, bce_dice_loss, device)
    val_loss = validate_epoch(model, val_loader, bce_dice_loss, device)
    scheduler.step(val_loss)
    print(f'Epoch {epoch+1:03d} | Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}')
    # 简单模型保存逻辑
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        torch.save(model.state_dict(), 'best_attention_unet.pth')

4.4 模型推理与可视化

训练完成后，我们看看模型预测效果。

def predict_and_visualize(model, image_tensor, device):
    """对单张图像进行预测并可视化。"""
    model.eval()
    with torch.no_grad():
        image_tensor = image_tensor.unsqueeze(0).to(device) # 增加batch维度
        pred_mask = model(image_tensor)
        pred_mask = (pred_mask > 0.5).float() # 二值化
    return pred_mask.squeeze().cpu().numpy()

# 取一个验证样本
sample_img, sample_mask = val_dataset[0]
pred_mask = predict_and_visualize(model, sample_img, device)

# 使用matplotlib进行可视化
import matplotlib.pyplot as plt
fig, axes = plt.subplots(1, 3, figsize=(12, 4))
axes[0].imshow(sample_img.squeeze(), cmap='gray')
axes[0].set_title('Input Image')
axes[0].axis('off')
axes[1].imshow(sample_mask.squeeze(), cmap='gray')
axes[1].set_title('Ground Truth')
axes[1].axis('off')
axes[2].imshow(pred_mask, cmap='gray')
axes[2].set_title('Prediction')
axes[2].axis('off')
plt.show()

5. 调优与避坑：让Attention UNet发挥真正实力

模型能跑起来只是第一步，要想在真实数据上取得好效果，还有很多细节要打磨。我结合自己项目里的经验，分享几个关键点。

第一，数据预处理是基石。 医学图像标准化至关重要。不要简单做/255.0。我常用的流程是：先对整个训练集计算图像的均值和标准差，然后对每张图做 (img - mean) / std 的标准化。对于CT值（HU值），可能还需要进行窗宽窗位调整，比如只保留[-1000, 400] HU范围内的值，并归一化到[0,1]。这能让模型训练更稳定，收敛更快。

第二，注意力机制不是银弹，通道数设置要小心。 在AttentionBlock里，中间通道数F_int是个超参数。论文和常见实现里通常设为F_l（跳跃连接通道数）的一半。但在实际中，如果F_l本身很小（比如网络浅层的64），F_int设为32可能特征压缩得太厉害。我的经验是，可以尝试设为F_l或者max(F_l//2, 32)，确保有足够的表达能力。另外，注意力图生成最后的1x1卷积后，可以不加BatchNorm，直接接Sigmoid，有时效果更直接。

第三，损失函数需要“因地制宜”。 BCE+Dice组合虽然经典，但对于极度不平衡的小目标（比如几个像素的微钙化），Dice Loss可能会不太稳定。可以尝试Focal Loss或者Tversky Loss。Tversky Loss通过调整α和β参数，可以控制对假阳和假阴的惩罚力度，在需要高召回率或高精度的场景下特别有用。我通常会在验证集上做一个小型实验，对比不同损失函数组合的效果。

第四，跳跃连接处的尺寸对齐是常见坑点。 由于最大池化、卷积步长等问题，编码器和解码器对应层的特征图尺寸可能不是严格的2倍关系，特别是输入尺寸不是2的整数次幂时。这会导致在torch.cat拼接时出错。我们的代码里虽然做了interpolate对齐，但更优雅的做法是在编码器的下采样过程中，使用卷积的padding和池化的ceil_mode参数来控制输出尺寸，或者统一在跳跃连接前使用一个适配层（如1x1卷积+插值）来调整通道和尺寸。

第五，可视化注意力图来调试模型。 想知道你的注意力机制到底有没有学到东西？一个很好的办法是把AttentionBlock中间生成的psi（注意力图）在验证时提取出来可视化。你可以把它叠加在原始图像上。一个训练良好的模型，其注意力图应该清晰地高亮目标区域。如果注意力图一片模糊或者全白，那可能是训练出了问题，或者注意力机制没有生效，需要回头检查网络结构或数据。

最后，关于训练技巧。医学数据量通常不大，强数据增强是必须的。除了常见的旋转、翻转、缩放，对于医学图像，弹性形变、随机伽马变换、对比度调整都非常有效。可以使用albumentations库。学习率方面，使用ReduceLROnPlateau调度器很有用，配合Early Stopping可以防止过拟合。如果显存不够，可以尝试使用混合精度训练（torch.cuda.amp），它能显著减少显存占用并可能加快训练速度。

把Attention UNet应用到实际项目中，它更像是一个强大的基础架构。你可以根据具体任务，把它和深度监督、多尺度输入、或者Transformer模块结合。我见过有的团队在瓶颈层加入轻量化的自注意力，或者在跳跃连接上用通道注意力（如SE Block）和空间注意力（即我们用的这种）并联，都取得了不错的效果。关键是多实验，多分析失败案例，理解你的数据特性，模型才能真正为你所用。