3D U-Net医学图像分割实战：肝脏肿瘤分割源码解析与避坑指南

1次阅读

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

医学图像分割，尤其是 3D 数据处理，存在几个新手常见痛点：

各向异性分辨率：CT/MRI 扫描通常在不同方向上分辨率不一致（如 0.5mm×0.5mm×2mm），直接输入网络会导致特征学习偏差
数据标注成本高：肝脏肿瘤标注需要放射科医生参与，导致正负样本比例可能达到 1:100 以上
GPU 显存瓶颈：3D 数据体积庞大，输入尺寸 128×128×128 时，单卡显存容易爆满

指标	2D U-Net	3D U-Net
参数量	约 8M	约 19M
DSC@LiTS	0.72	0.89
显存占用	4GB	11GB
上下文信息	单切片	三维空间

结论：当硬件允许时，3D U-Net 能更好捕捉肿瘤的空间分布特征

重采样：将所有数据统一到 1mm×1mm×1mm 各向同性分辨率
标准化：采用 CT 值截断（-200~250HU）后做 Z -score 归一化
Patch 提取：根据 GPU 显存动态调整 patch size（常用 96×96×96）

5 折交叉验证：充分利用有限标注数据
弹性形变增强：模拟器官生理变形
随机旋转：在 XY 平面进行±15°旋转
模态混合：将不同患者的正常 / 病变区域拼合成新样本

import torch
import torch.nn as nn

class DoubleConv(nn.Module):
    """(Conv3D -> BN -> ReLU) × 2"""
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.net = nn.Sequential(nn.Conv3d(in_ch, out_ch, kernel_size=3, padding=1),
            nn.BatchNorm3d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv3d(out_ch, out_ch, kernel_size=3, padding=1),
            nn.BatchNorm3d(out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.net(x)

class DownSample(nn.Module):
    """MaxPool + DoubleConv"""
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.net = nn.Sequential(nn.MaxPool3d(2),
            DoubleConv(in_ch, out_ch)
        )

    def forward(self, x):
        return self.net(x)

def focal_dice_loss(pred, target, alpha=0.7, gamma=2):
    # Focal Loss
    bce = F.binary_cross_entropy_with_logits(pred, target, reduction='none')
    pt = torch.exp(-bce)
    focal_loss = (alpha * (1-pt)**gamma * bce).mean()

    # Dice Loss
    pred = torch.sigmoid(pred)
    intersection = (pred * target).sum()
    union = pred.sum() + target.sum()
    dice_loss = 1 - (2.*intersection + 1e-5)/(union + 1e-5)

    return 0.3*focal_loss + 0.7*dice_loss

# 梯度累积训练示例
accum_steps = 4  # 累计 4 个 batch 的梯度
optimizer.zero_grad()

for i, (x, y) in enumerate(train_loader):
    pred = model(x.cuda())
    loss = criterion(pred, y.cuda()) / accum_steps
    loss.backward()

    if (i+1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

错误做法 ：直接nib.load(img_path).get_fdata() 可能得到 (x,y,z) 或(z,y,x)顺序

正确方案：

import nibabel as nib

def load_nii(path):
    img = nib.load(path)
    data = img.get_fdata()
    if img.affine[0,0] < 0:  # 判断是否需要翻转
        data = np.flip(data, axis=0)
    return np.transpose(data, (2,1,0))  # 转为 z,y,x

错误场景：同一患者的扫描出现在训练集和验证集
正确做法：按患者 ID 划分数据集，确保扫描序列完全隔离

model = nn.DataParallel(UNet3D(in_ch=1, out_ch=1),
    device_ids=[0,1]  # 使用两块 GPU
)
# 需在第一个卷积层前添加 SyncBN
model.module.conv1 = nn.SyncBatchNorm.convert_sync_batchnorm(model.module.conv1)

在 LiTS2017 测试集上的表现：