《nnU-Net 0基础入门(11):修改 data augmentation 与采样策略,让 nnU-Net 适应你的数据》

503611908 发布于 5 小时前 8 次阅读


本篇学习目标

这是《nnU-Net 0基础入门》系列的第 11 篇。上一篇我们修改了 loss。本文继续修改训练流程中的另一个关键环节:data augmentation 和采样策略。

读完本文,你应该能够:

  1. 理解 data augmentation 在 nnU-Net v2 训练流程中的位置。
  2. 知道默认增强大致包括哪些类型。
  3. 理解医学图像增强为什么不能无脑加。
  4. 通过自定义 Trainer 调整 mirror 和 foreground oversampling。

1. data augmentation 是什么

data augmentation,中文常译为数据增强,指训练时对图像和标签做随机变换,让模型看到更多变化形式。常见增强包括旋转、缩放、翻转、加噪声、模糊、亮度变化、对比度变化等。

在医学图像分割中,增强的目标不是把数据“变花哨”,而是让模型对合理变化更鲁棒。例如不同扫描仪、不同病人姿态、不同图像噪声,都可能造成输入变化。

flowchart TD
    A[预处理后的 case] --> B[nnUNetDataLoader 采样 patch]
    B --> C[get_training_transforms]
    C --> D[空间增强 / 强度增强 / mirror / cascade transforms]
    D --> E[训练 batch]
    E --> F[network forward + loss]

注意,augmentation 发生在训练时,不是 preprocessing 阶段永久写入磁盘。每个 epoch、每个 batch 都可能看到不同随机变换。

2. nnU-Net v2 默认增强包括什么

根据当前官方 nnUNetTrainer.py 中的 get_training_transforms,默认训练增强大致包括:

增强类型 源码中的 transform 直观作用
空间变换 SpatialTransform 旋转、缩放、裁剪到 patch
高斯噪声 GaussianNoiseTransform 模拟图像噪声
高斯模糊 GaussianBlurTransform 模拟分辨率或清晰度变化
亮度变化 MultiplicativeBrightnessTransform 模拟整体强度变化
对比度变化 ContrastTransform 模拟组织对比差异
低分辨率模拟 SimulateLowResolutionTransform 模拟较低采样质量
gamma 变换 GammaTransform 改变强度分布曲线
镜像翻转 MirrorTransform 沿允许轴随机翻转

如果是 cascade 训练,还会有和上一阶段分割结果相关的 transform。如果启用 deep supervision,还会加入 DownsampleSegForDSTransform,为多尺度输出准备标签。

3. 医学图像增强的风险

医学图像增强不能简单照搬自然图像。原因是医学图像里的方向、强度和空间关系可能有临床含义。

增强 可能风险 什么时候要谨慎
左右翻转 可能改变左右解剖语义 左右器官、侧别诊断、脑半球任务
大角度旋转 产生不真实体位 方向严格标准化的数据
强度变化 破坏定量意义 CT HU 值、定量 MRI、显微定量染色
低分辨率模拟 模糊小病灶 小目标分割或边界敏感任务

所以修改 augmentation 前,先问一个问题:这个变换在你的医学任务中是否仍然符合真实世界?如果不符合,它可能不是增强,而是在制造错误样本。

4. 最小修改一:关闭 mirror

如果你的任务对左右方向敏感,可以先尝试关闭 mirror。当前官方 Trainer 中,mirror 轴由 configure_rotation_dummyDA_mirroring_and_inital_patch_size 返回,并保存为 inference_allowed_mirroring_axes

可以写一个自定义 Trainer:

from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer


class nnUNetTrainerNoMirror(nnUNetTrainer):
    def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self):
        (
            rotation_for_DA,
            do_dummy_2d_data_aug,
            initial_patch_size,
            mirror_axes,
        ) = super().configure_rotation_dummyDA_mirroring_and_inital_patch_size()

        mirror_axes = None
        self.inference_allowed_mirroring_axes = mirror_axes
        return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes

训练命令:

nnUNetv2_train 1 3d_fullres 0 -tr nnUNetTrainerNoMirror --npz

这个修改比重写整个 get_training_transforms 更小,也更适合初学者验证一个具体假设:mirror 是否对当前任务有害。

5. 最小修改二:调整 foreground oversampling

foreground oversampling 指训练 patch 采样时,提高包含前景类别的 patch 比例。医学图像分割中,目标结构可能很小。如果随机采样,大量 patch 可能只有背景,模型很难学到小目标。

当前官方 Trainer 中有一个超参数:

self.oversample_foreground_percent = 0.33

可以创建一个更偏向前景采样的 Trainer:

from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer


class nnUNetTrainerMoreForeground(nnUNetTrainer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.oversample_foreground_percent = 0.5

训练命令:

nnUNetv2_train 1 3d_fullres 0 -tr nnUNetTrainerMoreForeground --npz

这并不等于“一定更好”。前景采样太多,可能让模型低估背景分布,或者对真实推理场景泛化变差。它适合在小目标任务中作为一个受控实验。

6. 不建议一开始重写整个 get_training_transforms

get_training_transforms 很强大,但也很容易改坏。它不仅包含图像增强,还包含:

  • dummy 2D augmentation。
  • cascade 相关 transform。
  • region-based training 转换。
  • ignore label 处理。
  • deep supervision 标签下采样。

如果你直接复制整段官方函数再改几行,未来 nnU-Net 更新后很容易和新版本脱节。更好的做法是先做小范围覆盖,例如关闭 mirror、调整采样比例或改一个概率参数。只有在你明确知道每个 transform 的输入输出时,再重写完整 pipeline。

7. 实验对照设计

augmentation 和采样策略的实验必须做对照。推荐至少比较:

实验 Trainer 目的
默认基线 nnUNetTrainer 确认官方默认表现
关闭 mirror nnUNetTrainerNoMirror 验证翻转是否伤害左右语义
更多前景采样 nnUNetTrainerMoreForeground 验证小目标召回是否改善

比较时不要只看平均 Dice。还应检查每个类别的 Dice、小目标召回、误报区域和可视化结果。增强策略有时会改善整体平均值,却让某个临床关键类别变差。

8. 调试建议

  • 先用短训练变体检查自定义 Trainer 能否跑起来。
  • 每次只改一个增强或采样参数。
  • 改 mirror 时,同步关注 inference mirroring,因为训练和推理时的 mirroring 语义应一致。
  • 改采样比例后,观察训练 loss、验证 Dice 和小目标类别指标。
  • 可视化增强后的 patch,确认没有产生明显不合理图像。

9. 官方资料入口

本文主要参考:

本篇总结

nnU-Net v2 默认 augmentation 已经覆盖空间变换、强度变化、噪声、模糊、低分辨率模拟、gamma 和 mirror 等多种情况。修改增强策略时,应从医学任务合理性出发,而不是盲目增加变换。初学者最适合从小改动开始,例如关闭 mirror 或调整 foreground oversampling,并用严格对照实验判断是否真的有效。

下一篇预告

下一篇是本系列最后一篇:修改 network architecture 与 plans。我们会讲 quick-and-dirty 的 Trainer 覆盖路线、proper planner 路线、ResEnc preset,以及为什么网络替换必须处理 deep supervision、patch size 和显存约束。

此作者没有提供个人介绍。
最后更新于 2026-05-14