3D医学图像分割技术解析：从算法原理到PyTorch实战

1次阅读

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

医学影像分析中，CT/MRI 数据本质是三维体数据（voxel），但传统处理方式常切片成 2D 图像分析，这会丢失空间上下文信息。实际落地时面临几个核心挑战：

各向异性分辨率：Z 轴分辨率常比 XY 轴低 5 -10 倍（如 1mm×1mm×5mm），导致 3D 卷积核难以均衡捕捉特征
标注成本高：专业医生标注单个 3D 病例需 4 - 6 小时，是 2D 标注的 20 倍工作量
显存爆炸：512×512×512 的 CT 扫描，float32 格式下原始数据就占用 1GB 显存

方法	优势	医疗影像缺陷
2D CNN	显存占用低，训练快	丢失层间关联，肿瘤边界识别差
3D CNN	保持空间关系，分割精度高	计算量 O(n³)级增长
Transformer	长程依赖建模能力强	需要超大规模数据，收敛慢

UNet3D 的杀手锏：

编码器 - 解码器结构缓解梯度消失
跳跃连接 (skip connection) 融合多尺度特征
可扩展性强，支持深度监督训练

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):
    """下采样层含 skip connection"""
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = DoubleConv(in_ch, out_ch)
        self.pool = nn.MaxPool3d(2)

    def forward(self, x):
        skipped = self.conv(x)
        down = self.pool(skipped)
        return down, skipped  # 返回下采样结果和跳跃特征

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    output = model(input)
    loss = criterion(output, target)

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

方法	显存占用(MB)	速度(iter/s)
原始模型	12000	1.2
+ 梯度检查点	7200	0.9
+ 混合精度	3800	2.1

梯度检查点实现：

from torch.utils.checkpoint import checkpoint

# 在 forward 时启用
x = checkpoint(block, x)

样本不均衡：Dice Loss 中加 squared 项抑制背景主导

def dice_loss(pred, target, smooth=1e-5):
    intersection = (pred * target).sum()
    return 1 - (2.*intersection + smooth)/(pred.sum() + target.sum() + smooth)