3DUNet过拟合问题实战：从数据增强到模型正则化的综合解决方案

1次阅读

没有评论

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

在医学影像分割任务中，3DUNet 因其优秀的性能被广泛使用，但过拟合问题一直困扰着研究人员。具体表现为：

训练集 Dice 系数高达 0.95+，但验证集只有 0.7 左右
验证集损失曲线剧烈波动，无法稳定收敛
模型对未见过的数据泛化能力极差

这些问题主要源于：

医学影像数据获取困难，像 BraTS 这样的公开数据集仅有几百例样本
3DUNet 模型参数量大，极易在小数据集上过拟合
医学影像的复杂性和多样性要求模型有更强的泛化能力

我们基于 MONAI 实现了 3D 弹性变形 + 随机遮挡的组合增强：

弹性变形：模拟器官的自然形变
随机遮挡：增强模型对部分遮挡的鲁棒性
参数可调：控制形变幅度和遮挡比例

采用 Dice+CrossEntropy 混合损失：

$$L = \alpha L_{Dice} + (1-\alpha)L_{CE}$$

其中：

$\alpha$ 控制两种损失的权重
加入 Label Smoothing 减轻模型对标签的过度自信

不同于传统 Dropout，DropBlock 在 3D 空间中丢弃连续区域：

在特征图上随机选择中心点
以该点为中心丢弃一个 3D 立方体
保持重要空间信息的连续性

import monai

transforms = monai.transforms.Compose([
    monai.transforms.Rand3DElasticd(keys=["image", "label"],
        sigma_range=(5,7),
        magnitude_range=(50,100),
        prob=0.8
    ),
    monai.transforms.RandSpatialCropSamplesd(keys=["image", "label"],
        roi_size=(96,96,96),
        num_samples=4
    ),
    monai.transforms.RandCoarseDropoutd(keys=["image", "label"],
        holes=10,
        spatial_size=(16,16,16),
        max_holes=8,
        prob=0.5
    )
])

class DropBlock3D(nn.Module):
    def __init__(self, block_size=7, keep_prob=0.9):
        super(DropBlock3D, self).__init__()
        self.block_size = block_size
        self.keep_prob = keep_prob

    def forward(self, x):
        if not self.training or self.keep_prob == 1:
            return x

        # 确保特征图尺寸大于 block_size
        if x.shape[2] < self.block_size or \
           x.shape[3] < self.block_size or \
           x.shape[4] < self.block_size:
            return x

        # 计算 gamma 值
        gamma = (1 - self.keep_prob) / (self.block_size ** 3)
        mask = torch.bernoulli(torch.ones_like(x) * gamma)
        mask = F.max_pool3d(mask, 
                          kernel_size=(self.block_size, self.block_size, self.block_size),
                          stride=(1, 1, 1),
                          padding=(self.block_size // 2, self.block_size // 2, self.block_size // 2))
        x = x * mask.numel() / mask.sum() * mask
        return x

class MixedLoss(nn.Module):
    def __init__(self, alpha=0.5, label_smoothing=0.1):
        super().__init__()
        self.alpha = alpha
        self.dice = DiceLoss(sigmoid=True)
        self.ce = nn.CrossEntropyLoss(label_smoothing=label_smoothing)

    def forward(self, pred, target):
        return self.alpha * self.dice(pred, target) + (1 - self.alpha) * self.ce(pred, target)