- 人工智能
- 深度学习
- 计算机视觉
- 医疗健康
【免费下载链接】nnUNet
导读
本文基于 documentation/competitions/AortaSeg24.md(DKFZ 团队在 AortaSeg24 挑战赛的参赛方案说明),完整还原其从数据预处理、模型训练到推理部署的全流程技术方案。该方案的核心是:以 nnU-Net v2 的 ResEnc L 架构为骨干,关闭左右镜像、改用重型 DA5 数据增强,并通过"全量模型 + 五折交叉验证集成"的组合完成推理。读完本文,你将掌握如何复用 nnUNet 仓库内现成的 nnUNetTrainer_onlyMirror01_DA5 训练器复现该方案,并理解每个命令与参数背后的源码实现原理。
AortaSeg24 方案总览:从方案设计到源码落点
该参赛方案的技术要点可以浓缩为以下几条(均可在仓库源码中得到印证):
- 骨干网络:nnU-Net ResEnc L(Residual Encoder U-Net 大配置),由
nnUNetPlannerResEncL规划器生成对应 plans; - 数据增强策略:禁用左/右镜像,仅保留轴 0/1 镜像(即
nnUNetTrainer_onlyMirror01),叠加 DA5 重型增强(即nnUNetTrainerDA5); - 训练硬件:A100 40GB GPU,从零训练(from scratch);
- 提交策略:1 个"全量数据单模型" + 1 个"五折交叉验证集成模型";
- 推理:直接使用 nnU-Net 官方推理工具
nnUNetv2_predict_from_modelfolder。
这些要素在仓库中有明确的代码实现:组合训练器类定义在 nnunetv2/training/nnUNetTrainer/variants/competitions/aortaseg24.py(一个仅 5 行的多继承类),DA5 增强管线定义在 nnunetv2/training/nnUNetTrainer/variants/data_augmentation/nnUNetTrainerDA5.py,镜像控制定义在 nnunetv2/training/nnUNetTrainer/variants/data_augmentation/nnUNetTrainerNoMirroring.py。
第一步:数据集转换与预处理
数据格式要求
参赛数据需要先转换为 nnU-Net 数据集格式。AortaSeg24 官方提供的数据以.mha文件为主,文档给出了两条等价路径:
- 保留
.mha并直接重命名:利用 nnU-Net v2 的多格式支持,NibabelIO与SimpleITKIO均原生支持.mha扩展名(见 nnunetv2/imageio/readme.md); - 转换为
.nii.gz:若希望与其他工具链统一,可先转成 NIfTI 格式。
无论选择哪种格式,必须遵守 nnU-Net 数据集目录规范:imagesTr(训练图像,文件名形如{CASE}_{XXXX}.{ext},其中 XXXX 为四位数通道标识)、labelsTr(分割标签)、以及描述元数据的dataset.json。注意:nnU-Net 仅支持无损压缩格式(不能使用.jpg),且整个数据集必须统一使用同一种文件格式。
运行规划与预处理命令
数据就绪后,运行预处理命令:
nnUNetv2_plan_and_preprocess -d 610 -c 3d_fullres -pl nnUNetPlannerResEncL -np 16各参数含义:
| 参数 | 取值 | 作用 |
|---|---|---|
-d | 610 | 数据集 ID(AortaSeg24 在本方案中被编号为 610,对应Dataset610_XXX) |
-c | 3d_fullres | 配置为 3D 全分辨率 |
-pl | nnUNetPlannerResEncL | 指定 ResEnc L 规划器,生成nnUNetResEncUNetLPlans |
-np | 16 | 并行预处理进程数 |
从源码看,nnUNetPlannerResEncL是ResEncUNetPlanner的子类,定义于 nnunetv2/experiment_planning/experiment_planners/residual_unets/residual_encoder_unet_planners.py#L224-L246。其默认gpu_memory_target_in_gb=24(面向 RTX 4090 / Titan RTX 级别显卡),生成的名称为nnUNetResEncUNetLPlans,并使用ResidualEncoderUNet作为骨干网络类。这解释了为何后文训练命令中 plans 名写为nnUNetResEncUNetLPlans。
第二步:模型训练
训练命令与参数拆解
nnUNetv2_train 610 3d_fullres all -p nnUNetResEncUNetLPlans -tr nnUNetTrainer_onlyMirror01_DA5参数说明:
| 参数 | 取值 | 作用 |
|---|---|---|
610 | 数据集 ID | 与预处理阶段一致 |
3d_fullres | 配置名 | 3D 全分辨率配置 |
all | fold 标识 | 使用全部数据训练单模型;若改为0~4则训练单个交叉验证折 |
-p | nnUNetResEncUNetLPlans | 显式指定 ResEnc L 的 plans 文件 |
-tr | nnUNetTrainer_onlyMirror01_DA5 | 定制训练器:仅 0/1 轴镜像 + DA5 增强 |
该命令针对"全量单模型"训练;五折交叉验证则需分别以0 1 2 3 4替换all各跑一次(即五个 fold 的模型,用于集成)。
定制训练器的源码原理
nnUNetTrainer_onlyMirror01_DA5的定义极其简洁——它通过多重继承组合了两个训练器:
class nnUNetTrainer_onlyMirror01_DA5(nnUNetTrainer_onlyMirror01, nnUNetTrainerDA5): pass(见 nnunetv2/training/nnUNetTrainer/variants/competitions/aortaseg24.py#L4-L5)
两个父类的职责如下:
1.nnUNetTrainer_onlyMirror01(镜像限制)
定义于 nnunetv2/training/nnUNetTrainer/variants/data_augmentation/nnUNetTrainerNoMirroring.py#L38-L52。它重写configure_rotation_dummyDA_mirroring_and_inital_patch_size:2D 数据只允许沿轴 0 镜像,3D 数据只允许沿轴 0 和 1 镜像,并同步将inference_allowed_mirroring_axes设置为(0,)或(0, 1)。
为什么禁掉左右镜像?AortaSeg24 的主动脉分割任务涉及解剖学左右语义——镜像是全局的左右翻转(关于 sagittal 平面,即轴 2),会破坏主动脉及其分支的左右解剖对应关系。保留轴 0/1(上下、前后)镜像则不会破坏这种语义。同样的思路也体现在nnUNetTrainerNoMirroring中(完全禁用所有镜像,nnunetv2/training/nnUNetTrainer/variants/data_augmentation/nnUNetTrainerNoMirroring.py#L29-L35),该方案按需选择镜像轴而非一刀切。
2.nnUNetTrainerDA5(重型增强管线)
定义于 nnunetv2/training/nnUNetTrainer/variants/data_augmentation/nnUNetTrainerDA5.py,是 nnU-Net v2 提供的"DA5 系列"增强训练器。其get_training_transforms构建的增强链包括:
- 空间变换:旋转(3D 各轴
p_rotation=0.4、每轴独立p_rot_per_axis=0.5)、缩放(p_scaling=0.2,比例范围(0.7, 1.43)); - 几何变换:90° 旋转(
Rot90Transform,概率 0.5)、轴转置(TransposeAxesTransform,概率 0.5); - 噪声/模糊:中值滤波与高斯模糊二选一(
OneOfTransform,各概率 0.2)、高斯噪声(概率 0.1); - 强度扰动:亮度加性变换(概率 0.1)、对比度变换(两档
BGContrast((0.5, 2)),各概率 0.2)、Gamma 变换(两个实例,各概率 0.1); - 分辨率模拟:低分辨率模拟
SimulateLowResolutionTransform(概率 0.15,尺度(0.25, 1)); - 局部伪影:空白矩形块
BlankRectangleTransform(概率 0.4)、亮度梯度BrightnessGradientAdditiveTransform(概率 0.3)、局部 GammaLocalGammaTransform(概率 0.3)、锐化SharpeningTransform(概率 0.2)。
镜像变换(MirrorTransform)在链中的位置受mirror_axes控制(nnunetv2/training/nnUNetTrainer/variants/data_augmentation/nnUNetTrainerDA5.py#L195-L196),因此与nnUNetTrainer_onlyMirror01组合后,仅轴 0/1 会执行镜像。该文件中还顺带提供了nnUNetTrainerDA5ord0(图像与分割均用最近邻插值)、nnUNetTrainerDA5Segord0(仅分割用最近邻插值)、nnUNetTrainerDA5_10epochs(10 个 epoch 的快速验证版)等变体,可作为进一步调参的参考。
训练资源建议:数据增强进程数
文档特别强调:DA5 属于重型增强管线,训练时容易触发CPU 瓶颈(数据增强的预处理速度跟不上 GPU 消费速度)。解决方案是提高数据增强进程数:
export nnUNet_n_proc_DA=32该环境变量的底层读取逻辑位于 nnunetv2/utilities/default_n_proc_DA.py#L20-L21:一旦设置了nnUNet_n_proc_DA,它就会覆盖脚本中基于主机名查表得到的默认值,并且最终会取min(设置值, os.cpu_count())(nnunetv2/utilities/default_n_proc_DA.py#L49)。也就是说,设置值不能超过物理 CPU 核数,32只是推荐起点,具体取值要结合你的机器核数与内存余量("if your system permits")。
第三步:推理与集成
单模型推理(全量训练模型)
训练完成后,使用nnUNetv2_predict_from_modelfolder直接从模型文件夹推理:
nnUNetv2_predict_from_modelfolder -i INPUT_FOLDER -o OUTPUT_FOLDER -m MODEL_FOLDER -f all-i INPUT_FOLDER:输入图像文件夹(格式与训练数据一致,注意推理输入图像的通道与命名必须与训练完全对应);-o OUTPUT_FOLDER:预测结果输出文件夹;-m MODEL_FOLDER:模型所在文件夹(包含 plans、dataset.json 与训练权重);-f all:必须显式指定,表示仅使用"全量数据训练的那个模型"进行预测。
如果不加-f all,该命令默认行为是集成已训练的 5 个 fold(详见 documentation/how-to/run-inference.md),而"全量单模型"只有一个模型,因此这里要明确-f all。
五折集成推理
五折交叉验证模型的集成推理则保持默认即可:
nnUNetv2_predict_from_modelfolder -i INPUT_FOLDER -o OUTPUT_FOLDER -m MODEL_FOLDER此时 5 个 fold 的模型各自输出概率图并求平均集成,这正是 nnU-Net 官方推荐的推理模式(各 fold 权重在 nnunetv2/inference/predict_from_raw_data.py 对应的推理流程中自动加载并平均)。
两个模型(全量单模型 + 五折集成)的预测结果都可用于提交,官方提交通常取两者中的更优者,或者按挑战赛规则分别提交。
复现该方案的完整流水线清单
将上述步骤串成端到端流程(假设已安装 nnU-Net v2 并设置好nnUNet_raw、nnUNet_preprocessed、nnUNet_results三个环境变量):
- 数据转换:将 AortaSeg24 原始
.mha按 nnU-Net 格式整理到nnUNet_raw/Dataset610_AortaSeg24/(imagesTr+labelsTr+dataset.json); - 预处理:
nnUNetv2_plan_and_preprocess -d 610 -c 3d_fullres -pl nnUNetPlannerResEncL -np 16; - 训练全量模型:
nnUNetv2_train 610 3d_fullres all -p nnUNetResEncUNetLPlans -tr nnUNetTrainer_onlyMirror01_DA5; - 训练五折模型:将上一步
all依次替换为0~4各执行一次(建议先export nnUNet_n_proc_DA=32缓解 CPU 瓶颈); - 推理:按上文分别运行两条
nnUNetv2_predict_from_modelfolder命令,产出两份预测结果用于提交。
总结与扩展阅读
AortaSeg24 参赛方案展示了 nnU-Net v2 高度可组合的架构:通过"训练器多继承"(nnUNetTrainer_onlyMirror01+nnUNetTrainerDA5)即可在不修改任何框架代码的情况下定制出符合任务语义的数据增强策略,配合nnUNetPlannerResEncL规划出的 ResEnc L 网络,构成了一个开箱即用、可直接复现的高性能分割流水线。方案中所有组件(规划器、训练器、增强变换、推理入口)均以源码形式存在于当前仓库,是研究 nnU-Net 定制化训练的绝佳实例。
若想深入了解相关机制,可继续阅读仓库内以下资料:
- 数据集格式规范:documentation/dataset_format.md 与 documentation/reference/dataset-format.md;
- 推理通用指南:documentation/how-to/run-inference.md;
- DA5 训练器完整实现:nnunetv2/training/nnUNetTrainer/variants/data_augmentation/nnUNetTrainerDA5.py;
- 镜像控制训练器:nnunetv2/training/nnUNetTrainer/variants/data_augmentation/nnUNetTrainerNoMirroring.py;
- ResEnc 规划器系列:nnunetv2/experiment_planning/experiment_planners/residual_unets/residual_encoder_unet_planners.py;
- 数据增强进程数控制:nnunetv2/utilities/default_n_proc_DA.py。
- 人工智能
- 深度学习
- 计算机视觉
- 医疗健康
【免费下载链接】nnUNet
相关推荐
MONAI与nnU-Net集成实战指南:医学图像分割的强强联合
MONAI与nnU Net集成实战指南:医学图像分割的强强联合 前言 在医学影像分析领域,nnU Net和MONAI都是备受推崇的开源框架。本文将深入探讨这两个
示例工程教程OpenCore Legacy Patcher完整教程:4步让老Mac重获新生的终极指南
OpenCore Legacy Patcher完整教程:4步让老Mac重获新生的终极指南 还在为你的老Mac无法升级最新macOS而烦恼吗?看着2012年的Ma
人工智能深度学习计算机视觉医疗健康MMDetection 中 InstaBoost 数据增强实战:基于概率图引导 Copy-Pasting 的实例分割提升方案
MMDetection 中 InstaBoost 数据增强实战:基于概率图引导 Copy Pasting 的实例分割提升方案 InstaBoost 是一种基于概
人工智能计算机视觉深度学习模型评测
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考