本篇学习目标
这是《nnU-Net 0基础入门》系列的第 5 篇。上一篇我们跑通了 plan、preprocess、train、predict 的主流程。本文把训练阶段拆开讲清楚:fold、configuration、checkpoint 和日志怎么看。
读完本文,你应该能够:
- 理解
nnUNetv2_train DATASET CONFIGURATION FOLD三个核心参数。 - 知道 5-fold cross-validation 为什么重要。
- 看懂训练输出目录、checkpoint、validation 结果和
progress.png。 - 知道
--npz、--val --npz、-device和多 GPU 的基础用法。
1. 训练命令再看一遍
官方训练命令的基本形式是:
nnUNetv2_train DATASET_NAME_OR_ID CONFIGURATION FOLD
例如:
nnUNetv2_train 1 3d_fullres 0
这三个位置参数分别是:
| 参数 | 例子 | 含义 |
|---|---|---|
DATASET_NAME_OR_ID |
1 或 Dataset001_Liver |
要训练的数据集编号或名称 |
CONFIGURATION |
2d、3d_fullres |
训练配置,由 plan 阶段生成 |
FOLD |
0、1、2、3、4 |
交叉验证中的 fold 编号 |
初学者最常见的误解是:以为训练一次 fold 0 就完成了整个 nnU-Net 实验。实际上,nnU-Net 的标准流程通常依赖 5-fold cross-validation 来比较配置、做 ensemble 和确定后处理。
2. fold 是什么:为什么不是只训练一次
cross-validation 中文常译为交叉验证。它的核心思想是:把训练集分成几份,多次训练和验证,让每个样本都有机会作为验证样本。
nnU-Net 通常使用 5-fold。可以把它理解成这样:
flowchart TD
A[训练集] --> B[fold 0: 第0份做验证,其余做训练]
A --> C[fold 1: 第1份做验证,其余做训练]
A --> D[fold 2: 第2份做验证,其余做训练]
A --> E[fold 3: 第3份做验证,其余做训练]
A --> F[fold 4: 第4份做验证,其余做训练]
B --> G[汇总验证表现]
C --> G
D --> G
E --> G
F --> G
这样做的好处是:模型选择不只依赖一次随机划分,更稳定,也能让后续的 configuration 比较和 ensemble 更有依据。
训练 5 个 fold 的命令形式如下:
nnUNetv2_train 1 3d_fullres 0 --npz
nnUNetv2_train 1 3d_fullres 1 --npz
nnUNetv2_train 1 3d_fullres 2 --npz
nnUNetv2_train 1 3d_fullres 3 --npz
nnUNetv2_train 1 3d_fullres 4 --npz
--npz 会保存验证集的 softmax 概率文件。softmax 概率指模型对每个像素或体素属于各个类别的概率。官方文档说明,如果后续要用 nnUNetv2_find_best_configuration 做自动配置比较和 ensemble,应在训练时加 --npz。
3. configuration 怎么理解
configuration 是 nnU-Net 根据数据自动生成的训练配置。常见有:
2d:适合一些 2D 数据,或者 3D 数据中层间 spacing 差异很大的情况。3d_fullres:全分辨率 3D U-Net,是 3D 医学图像中常见配置。3d_lowres:低分辨率 3D U-Net,用更低分辨率获得更大上下文。3d_cascade_fullres:级联全分辨率模型,通常依赖低分辨率模型先完成粗分割。
官方训练文档提醒:不是每个数据集都会有每一种 configuration。对于 cascade,必须先训练 3d_lowres,再训练 3d_cascade_fullres。
# 如果数据集生成了 cascade,并且你要训练 cascade:
nnUNetv2_train 1 3d_lowres 0 --npz
nnUNetv2_train 1 3d_cascade_fullres 0 --npz
不要在没有生成 3d_lowres 的情况下强行训练 cascade。configuration 是否存在,以 nnUNet_preprocessed 中的 plans 和实际预处理结果为准。
4. 输出目录怎么看
官方文档说明,训练输出写在:
nnUNet_results/DatasetXXX_Name/TRAINER__PLANS__CONFIGURATION/fold_X
例如:
nnUNet_results/
└── Dataset001_Liver/
└── nnUNetTrainer__nnUNetPlans__3d_fullres/
├── fold_0/
├── fold_1/
├── fold_2/
├── fold_3/
└── fold_4/
每个 fold 目录里常见文件包括:
| 文件或目录 | 作用 | 怎么看 |
|---|---|---|
checkpoint_final.pth |
训练结束时保存的模型权重 | 说明该 fold 正常训练到结束 |
checkpoint_best.pth |
验证表现较好的模型权重 | 推理或模型选择时常会用到 |
progress.png |
训练曲线图 | 观察 loss、Dice、学习率趋势 |
validation/summary.json |
验证结果摘要 | 查看验证集指标 |
validation/*.npz |
验证概率文件 | 只有启用 --npz 时保存,用于后续模型选择和 ensemble |
5. progress.png 和日志里有哪些指标
官方 logging reference 说明,nnU-Net v2 的本地 logger 是训练曲线和 progress.png 的主要来源。默认每个 epoch 会记录:
mean_fg_dice:前景类别的平均 Dice。ema_fg_dice:指数滑动平均后的前景 Dice。dice_per_class_or_region:每个类别或 region 的 Dice。train_losses和val_losses:训练和验证 loss。lrs:学习率。- epoch 开始和结束时间。
Dice 是医学图像分割里最常见的重叠指标之一,用来衡量预测区域和真实标签区域的重合程度。Dice 越高通常表示分割越接近标签,但不同任务之间不能简单横向比较。
看 progress.png 时,可以先关注三件事:
- 训练 loss 是否整体下降。
- 验证 Dice 是否整体上升或至少稳定。
- 训练和验证表现是否明显背离,例如训练越来越好但验证越来越差。
6. 如果忘了加 --npz 怎么办
如果你已经训练完某个 fold,但忘了加 --npz,官方训练文档给出补救方式:重新运行 validation 并保存 npz。
nnUNetv2_train 1 3d_fullres 0 --val --npz
这不会从头训练,而是运行验证流程并保存后续模型选择所需的概率文件。建议从一开始就加 --npz,避免后面补跑。
7. 设备选择和多 GPU 训练
官方文档说明,可以用 -device 选择设备:
nnUNetv2_train 1 3d_fullres 0 --npz -device cuda
nnUNetv2_train 1 2d 0 --npz -device cpu
nnUNetv2_train 1 2d 0 --npz -device mps
实际项目中,医学图像分割训练通常应使用 CUDA GPU。CPU 更适合做命令演示或极小规模测试。
多 GPU 机器上,官方推荐的常见策略是“一张 GPU 跑一个训练任务”。例如:
CUDA_VISIBLE_DEVICES=0 nnUNetv2_train 1 3d_fullres 0 --npz
CUDA_VISIBLE_DEVICES=1 nnUNetv2_train 1 3d_fullres 1 --npz
也可以使用分布式训练:
nnUNetv2_train 1 3d_fullres 0 --npz -num_gpus 2
对初学者来说,先掌握单 GPU 和一 GPU 一任务的方式即可。分布式训练涉及更多环境和同步问题,后面有实际需求再深入。
8. 常见训练判断
| 现象 | 可能含义 | 优先检查 |
|---|---|---|
没有 checkpoint_final.pth |
训练可能未正常结束 | 查看终端输出、日志和是否中断 |
没有 validation/*.npz |
训练时可能没加 --npz |
用 --val --npz 补跑验证 |
| Dice 长期接近 0 | 标签、类别、数据方向或任务本身可能有问题 | 先可视化原图和标签,不要只调参数 |
| 训练 loss 下降,验证 Dice 不动 | 可能过拟合、验证集困难或标签质量问题 | 检查数据划分、标签质量、类别分布 |
| 显存不足 | 配置和硬件不匹配 | 尝试其他 configuration,或使用更大显存 GPU |
9. 官方资料入口
本文主要参考:
本篇总结
nnU-Net v2 的训练不是只跑一次命令。你需要明确数据集、configuration 和 fold 的关系;理解 5-fold cross-validation 对模型选择的重要性;训练时尽量加 --npz;训练后要检查 checkpoint、validation 结果和 progress.png。只有把这些产物看懂,后续的最佳配置选择、ensemble 和 postprocessing 才有依据。
下一篇预告
下一篇我们会进入模型选择:如何用 nnUNetv2_find_best_configuration 比较 2D、3D、cascade 等配置,什么时候需要 ensemble,postprocessing 又是在解决什么问题。
Comments NOTHING