AttentionUNet在医学图像分割中的原理与实践：从模型架构到性能优化

1次阅读

没有评论

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

医学图像分割是计算机辅助诊断的关键步骤，但在实际应用中常常遇到两个主要问题：

器官边缘模糊：由于医学成像设备的物理限制和生物组织的特性，不同组织之间的边界往往不够清晰
小病灶漏检：对于微小病变或细小结构，传统分割方法容易产生漏检

我们对比了 UNet 和 AttentionUNet 在胰腺分割数据集上的表现：

标准 UNet 的 Dice 系数为 0.78
AttentionUNet 将 Dice 提升到 0.85
特别是在小血管分割任务上，召回率提高了 12%

AttentionUNet 的核心创新是在跳跃连接 (skip connection) 中加入了注意力门控机制。这个机制可以自动学习并突出对分割任务重要的区域特征。

输入特征：来自编码器的低级特征 x 和来自解码器的高级特征 g
线性变换 ：通过两个 1×1 卷积(W_x 和 W_g) 将特征映射到相同维度
激活处理：ReLU 激活后接 Sigmoid 得到注意力权重
特征重加权：原始特征 x 与注意力权重逐点相乘

数学表达式为：

F_att = σ(ψ^T(ReLU(W_x*x + W_g*g + b))) ⊙ x

其中⊙表示逐元素乘法，σ 是 Sigmoid 函数。

import torch
import torch.nn as nn

class AttentionGate(nn.Module):
    def __init__(self, F_g, F_l, F_int):
        super(AttentionGate, self).__init__()
        self.W_g = nn.Sequential(nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(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)
        )

        self.psi = nn.Sequential(nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(1),
            nn.Sigmoid())

        self.relu = nn.ReLU(inplace=True)

    def forward(self, g, x):
        # 维度对齐
        g1 = self.W_g(g)
        x1 = self.W_x(x)
        # 特征融合
        psi = self.relu(g1 + x1)
        psi = self.psi(psi)
        # 特征重加权
        return x * psi

import SimpleITK as sitk
import numpy as np

def n4_bias_correction(image):
    """MRI N4 偏场校正"""
    input_image = sitk.GetImageFromArray(image)
    mask_image = sitk.OtsuThreshold(input_image, 0, 1, 200)
    corrector = sitk.N4BiasFieldCorrectionImageFilter()
    output_image = corrector.Execute(input_image, mask_image)
    return sitk.GetArrayFromImage(output_image)

def apply_window(image, window_center, window_width):
    """窗宽窗位调整"""
    min_val = window_center - window_width/2
    max_val = window_center + window_width/2
    image = np.clip(image, min_val, max_val)
    image = (image - min_val) / (max_val - min_val)
    return image

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast():
    outputs = model(inputs)
    loss = criterion(outputs, labels)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

针对类别不平衡问题，我们使用带权重的 Dice Loss：

def dice_loss(pred, target, smooth=1e-5):
    """
    pred: (N, C, H, W)
    target: (N, H, W)
    """
    target = F.one_hot(target, num_classes=pred.shape[1]).permute(0,3,1,2)
    intersection = (pred * target).sum(dim=(2,3))
    union = pred.sum(dim=(2,3)) + target.sum(dim=(2,3))

    # 类别权重
    weights = 1.0 / (target.sum(dim=(0,2,3))**2 + 1e-5)
    weights = weights / weights.sum()

    dice = (2. * intersection + smooth) / (union + smooth)
    dice_loss = 1 - dice
    weighted_loss = (dice_loss * weights).sum()
    return weighted_loss

解决方案：