AttentionUNet医学图像分割实战：从模型原理到PyTorch实现

1次阅读

没有评论

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

医学图像分割（如 CT/MRI 中的病灶检测）面临两个核心挑战：

小目标漏检问题：肿瘤或病变区域可能只占图像的几个像素，传统 U -Net 下采样时容易丢失这些细节
边界模糊问题：器官边缘与周围组织对比度低，普通卷积难以捕捉精确边界

以肝肿瘤分割为例，传统 U -Net 的 skip connection 直接拼接高低层特征，会导致：

浅层特征中的噪声被传递到解码器
无关背景区域干扰关键部位的重建

我们对比三种典型架构在 BraTS2018 数据集上的表现：

模型	参数量(M)	推理速度(FPS)	Dice 系数
U-Net	31.4	45.2	0.781
AttentionUNet	32.1(+2%)	38.5(-15%)	0.823(+5%)
nnUNet	151.7	12.3	0.851

AttentionUNet 通过仅增加 2% 的参数，获得了显著的精度提升。其核心优势在于：

注意力门控（AG）自动抑制无关区域
保留原始 U -Net 的轻量级特性

class AttentionGate(nn.Module):
    def __init__(self, F_g, F_l, F_int):
        """
        F_g: 门控信号通道数（来自解码器）F_l: 局部特征通道数（来自编码器）F_int: 中间层通道数
        """
        super().__init__()
        self.W_g = nn.Sequential(nn.Conv2d(F_g, F_int, 1, bias=False),
            nn.BatchNorm2d(F_int)
        )
        self.W_x = nn.Sequential(nn.Conv2d(F_l, F_int, 1, bias=False),
            nn.BatchNorm2d(F_int)
        )
        self.psi = nn.Sequential(nn.Conv2d(F_int, 1, 1, bias=False),
            nn.BatchNorm2d(1),
            nn.Sigmoid())

    def forward(self, g, x):
        g1 = self.W_g(g)
        x1 = self.W_x(x)
        psi = F.relu(g1 + x1)
        psi = self.psi(psi)
        return x * psi  # 注意力加权

处理 CT 图像时必须考虑窗宽 (Window Width) 和窗位(Window Center)：

def dicom_to_tensor(dcm_path):
    """将 DICOM 文件转换为归一化后的 Tensor"""
    ds = pydicom.dcmread(dcm_path)
    img = ds.pixel_array.astype(np.float32)

    # 窗宽窗位调整（以肺窗为例）center = ds.WindowCenter if hasattr(ds, 'WindowCenter') else 40
    width = ds.WindowWidth if hasattr(ds, 'WindowWidth') else 400
    img = np.clip((img - center + 0.5*width)/width, 0, 1)

    # 处理不同位深
    if ds.BitsAllocated == 16:
        img = (img - img.min()) / (img.max() - img.min())

    return torch.from_numpy(img).unsqueeze(0)  # 增加通道维度

对于肿瘤占比不足 5% 的数据：

损失函数设计：

class MixedLoss(nn.Module):
    def __init__(self, alpha=0.5):
        super().__init__()
        self.alpha = alpha  # 平衡系数
        self.bce = nn.BCEWithLogitsLoss()

    def dice_coeff(self, pred, target):
        smooth = 1.0
        pred = torch.sigmoid(pred)
        intersection = (pred * target).sum()
        return (2. * intersection + smooth) / (pred.sum() + target.sum() + smooth)

    def forward(self, pred, target):
        return self.alpha * self.bce(pred, target) - (1-self.alpha) * torch.log(self.dice_coeff(pred, target))

数据采样策略：
对包含病灶的切片进行过采样
使用 ROI-Crop 只裁剪包含目标的区域

当使用 DataParallel 或DistributedDataParallel时：

必须将 BatchNorm 替换为SyncBatchNorm

model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
model = nn.DataParallel(model)

学习率需要按 GPU 数量线性缩放

lr = base_lr * torch.distributed.get_world_size()

在 BraTS 验证集上的典型结果：

使用 Monai 进行 3D 可视化：

from monai.visualize import plot_2d_or_3d_image

# 预测结果转伪彩色
pred_color = plt.cm.viridis(pred.cpu().numpy())
label_color = plt.cm.spring(label.cpu().numpy())
plot_2d_or_3d_image(data=[image, pred_color, label_color], 
    figsize=(12, 6),
    titles=["Input", "Pred", "GT"]
)

当前架构仍存在长程依赖建模能力不足的问题。可能的改进方向：