3D卷积网络UNet入门实战：从医学图像分割到模型优化

1次阅读

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

医学影像数据（如 CT、MRI）本质上是三维体数据，传统 2D 卷积只能处理单层切片，会丢失层间空间信息。举个简单例子：

2D 卷积：像看书的每一页，但不知道页与页之间的关系
3D 卷积：像真正翻书阅读，能理解整本书的上下文

实际测试显示，对脑肿瘤分割任务，3D UNet 的 Dice 系数比 2D 版本平均高出 15%-20%，尤其在细小血管识别上优势明显。

先定义编码器和解码器的基本模块，采用模块化设计方便后续调整：

import torch
import torch.nn as nn

class DoubleConv(nn.Module):
    """(卷积 => [BN] => ReLU) * 2"""
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm3d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm3d(out_channels),
            nn.ReLU(inplace=True)
        )

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

UNet 的核心特性是通过跳跃连接保留空间细节，图示说明：

编码器下采样 → [特征图] ---- 拼接 ---→ 解码器上采样
                (跳跃连接)

对应代码实现：

class UNet3D(nn.Module):
    def __init__(self, n_channels=1, n_classes=1):
        super().__init__()
        # 编码器部分
        self.enc1 = DoubleConv(n_channels, 64)
        self.pool1 = nn.MaxPool3d(2)
        # ... 中间层省略

        # 解码器部分
        self.up4 = nn.ConvTranspose3d(128, 64, kernel_size=2, stride=2)
        self.dec4 = DoubleConv(128, 64)  # 128=64+64（跳跃连接通道合并）def forward(self, x):
        # 编码过程
        enc1 = self.enc1(x)
        enc2 = self.enc2(self.pool1(enc1))
        # ...

        # 解码过程（带跳跃连接）dec4 = self.up4(enc4)
        dec4 = torch.cat([dec4, enc1], dim=1)  # 关键拼接操作
        dec4 = self.dec4(dec4)
        return dec4

医学影像常用 NIfTI 格式，推荐使用 nibabel 库读取：

import nibabel as nib

def load_nii(path):
    """
    加载 NIfTI 文件并返回数据和元信息
    返回值: (体数据数组, 体素间距, 坐标系信息)
    """
    img = nib.load(path)
    data = img.get_fdata()
    zooms = img.header.get_zooms()  # 获取体素间距(重要！)
    return data, zooms, img.affine

针对医学数据稀缺问题，推荐这些增强组合：

空间变换：
随机旋转（±15 度范围内）
弹性形变（模拟器官蠕动）
随机裁剪（确保保留 ROI 区域）
强度变换：
高斯噪声（均值 0，方差 0.1）
伽马校正（gamma∈[0.7,1.3]）

示例代码：

from scipy.ndimage import rotate

def random_rotate(volume, label, max_angle=15):
    angle = np.random.uniform(-max_angle, max_angle)
    # 对每个切片旋转相同角度
    volume = rotate(volume, angle, axes=(1,2), reshape=False)
    label = rotate(label, angle, axes=(1,2), reshape=False)
    return volume, label

通过自动混合精度 (AMP) 减少显存占用：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for inputs, labels in dataloader:
    optimizer.zero_grad()

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

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

根据任务类型选择激活函数：

二分类：Sigmoid（输出值在 0 - 1 之间）
多分类：Softmax（每个体素类别概率）
回归任务：不使用激活（直接输出数值）

3D 卷积的计算量公式：

FLOPs = C_in × C_out × K³ × H × W × D

其中 K 为卷积核大小，(H,W,D)为特征图尺寸。举例说明：

输入 128x128x128，64 通道的 3D 卷积（核 3x3x3）
计算量 ≈ 64×64×27×128³ ≈ 220G FLOPs

方法	优点	缺点
MaxPooling	保留显著特征	丢失位置细节
AvgPooling	平滑噪声	模糊边界
Strided Conv	可学习下采样	需精心调参
SpaceToDepth	无参操作	通道数爆炸