本篇学习目标
这是《nnU-Net 0基础入门》系列的第 8 篇。从这一篇开始,我们不再只停留在“会用命令”,而是进入 nnU-Net v2 的内部框架。理解这些模块,是后面修改 Trainer、loss、augmentation 和 network architecture 的前提。
读完本文,你应该能够:
- 说清楚 fingerprint、plans、preprocessing、Trainer、Predictor 分别负责什么。
- 理解
nnUNetPlans.json为什么是连接数据分析和训练配置的关键文件。 - 知道哪些修改通常需要重跑 preprocessing,哪些通常不需要。
- 建立一张源码级的 nnU-Net v2 框架地图。
1. 先看总图:nnU-Net v2 内部数据流
nnU-Net v2 的内部流程可以理解为“数据特征 → 实验计划 → 预处理数据 → 训练器 → 推理器”。
flowchart TD
A[原始数据 DatasetXXX] --> B[Dataset Fingerprint Extractor]
B --> C[dataset_fingerprint.json]
C --> D[Experiment Planner]
D --> E[nnUNetPlans.json]
E --> F[Preprocessor]
F --> G[预处理后的训练数据]
E --> H[nnUNetTrainer]
G --> H
H --> I[checkpoints / validation / logs]
E --> J[nnUNetPredictor]
I --> J
J --> K[预测结果]
这张图说明一个核心事实:nnU-Net v2 不是直接把原图丢给 Trainer。它会先分析数据,再生成 plans,然后用 plans 指导预处理、训练和推理。
2. fingerprint:数据集的“体检报告”
dataset fingerprint 可以理解为数据集的体检报告。它不是训练结果,而是对训练数据本身的统计和描述,例如:
- 图像尺寸分布。
- 体素 spacing 分布。
- 强度统计信息。
- 通道和标签相关信息。
这些信息会保存到 dataset_fingerprint.json。后续 planner 会读取它,并根据数据特征决定预处理和训练配置。
你可以把 fingerprint 想成医生检查病人的基础指标:不先量身高体重、血压、血常规,就直接开训练方案,风险很大。
3. plans:nnU-Net 的实验计划书
nnUNetPlans.json 是第 8 篇最重要的文件。官方 reference 说明,plans 文件定义了 nnU-Net 如何预处理数据,以及每个 configuration 如何训练。
它包含两类信息:
- global dataset-level settings:数据集级别设置,例如 reader/writer、label manager、转置规则、数据集名称。
- configuration-level settings:每个配置自己的设置,例如 spacing、patch size、batch size、preprocessor、normalization、network architecture。
| plans 字段 | 含义 | 影响 |
|---|---|---|
spacing |
目标 spacing | 影响重采样和图像尺度 |
patch_size |
训练 patch 大小 | 影响上下文范围和显存 |
batch_size |
batch 大小 | 影响训练稳定性和显存 |
preprocessor_name |
预处理器名称 | 决定预处理逻辑 |
normalization_schemes |
各通道归一化策略 | 影响输入强度分布 |
network_arch_class_name |
网络架构类 | 影响模型结构 |
所以 plans 不是普通日志,而是后续训练和推理共享的配置来源。随意手改 plans 可能导致训练、推理、预处理不一致。
4. preprocessing:把原始数据变成训练数据
preprocessing 是预处理。它负责把原始医学图像转换成 nnU-Net 训练时真正读取的数据。常见工作包括:
- 按 plans 中的 spacing 重采样。
- 按通道归一化强度。
- 裁剪无效背景区域。
- 保存可快速读取的中间数据。
预处理结果通常放在 nnUNet_preprocessed/DatasetXXX_Name 下。训练时 Trainer 读取的不是 nnUNet_raw 里的原图,而是这里的预处理数据。
这也解释了一个常见问题:如果你改了 spacing、normalization 或 preprocessor,但没有重跑 preprocessing,训练可能仍然读旧数据。
5. Trainer:训练行为的核心入口
Trainer 是控制训练流程的类。nnU-Net v2 中最核心的是 nnUNetTrainer。它负责把数据加载、网络构建、loss、optimizer、训练循环、验证、checkpoint 保存等步骤组织起来。
从学习路径看,Trainer 是最重要的扩展入口。后面我们要修改 loss、augmentation、optimizer 或训练流程时,通常不是直接改 nnU-Net 核心文件,而是继承 nnUNetTrainer 写一个自己的 Trainer 变体。
| 想修改什么 | 通常从哪里入手 |
|---|---|
| loss | 自定义 Trainer 中覆盖 loss 构建逻辑 |
| optimizer / scheduler | 自定义 Trainer 中覆盖优化器相关方法 |
| augmentation | 自定义 Trainer 中修改 transform 构建 |
| network architecture | 覆盖或调整网络构建逻辑,并保证 plans 兼容 |
6. Predictor:把训练结果用于新数据
Predictor 负责推理流程。官方源码中,推理相关逻辑集中在 nnunetv2/inference/predict_from_raw_data.py。它会读取训练产生的模型、plans 和 dataset 信息,对输入图像执行与训练匹配的预处理和预测。
这就是为什么推理时不能只拿一个 checkpoint 随便写 PyTorch 代码加载。nnU-Net 推理不仅需要权重,还需要知道训练时的 plans、normalization、spacing、patch size、滑窗预测策略等信息。
7. 哪些修改需要重跑 preprocessing
官方 plans reference 给出一个重要判断:如果修改影响已经准备好的数据,就通常需要重跑 preprocessing。
| 修改内容 | 通常是否需要重跑 preprocessing | 原因 |
|---|---|---|
spacing |
需要 | 重采样结果会变 |
preprocessor_name |
需要 | 预处理逻辑会变 |
normalization_schemes |
需要 | 输入强度分布会变 |
| resampling functions | 需要 | 重采样算法会变 |
batch_size |
通常不需要 | 只影响训练加载批量,不改变预处理数据 |
| 部分 architecture-only 设置 | 通常不需要 | 若仍复用同一份预处理数据,则不必重跑 |
实践中,如果你不确定某个改动是否改变预处理数据,保守做法是新建一个 configuration 和新的 data_identifier,避免旧数据和新实验混在一起。
8. 框架地图:从命令到模块
把前面用过的命令和内部模块对应起来,可以得到这张表:
| 命令 | 主要模块 | 关键产物 |
|---|---|---|
nnUNetv2_plan_and_preprocess |
fingerprint extractor、planner、preprocessor | dataset_fingerprint.json、nnUNetPlans.json、预处理数据 |
nnUNetv2_train |
nnUNetTrainer |
checkpoint、日志、validation 结果 |
nnUNetv2_find_best_configuration |
模型选择、ensemble、postprocessing 评估 | inference_instructions.txt |
nnUNetv2_predict |
nnUNetPredictor |
预测分割文件、可选 probability files |
9. 官方资料入口
本文主要参考:
- How nnU-Net Works
- Plans and Configuration Reference
- Extending nnU-Net
- nnUNetTrainer.py 官方源码
- predict_from_raw_data.py 官方源码
本篇总结
nnU-Net v2 的内部框架可以按 fingerprint、plans、preprocessing、Trainer、Predictor 来理解。fingerprint 描述数据,plans 连接数据分析和训练配置,preprocessing 生成训练可读的数据,Trainer 控制训练过程,Predictor 使用模型和 plans 对新数据推理。理解这条链路后,后面修改 loss、augmentation 和 network architecture 才不会变成盲改。
下一篇预告
下一篇我们正式开始修改 nnU-Net v2:从继承 nnUNetTrainer 写一个自定义 Trainer 开始。我们会讲为什么不要直接改核心源码,以及如何让 nnUNetv2_train 调用你的 Trainer。
Comments NOTHING