nnU-Net：强基线来自完整流程

Date 2026-05-11 · Category tech · Status finished · Confidence likely
论文解读, 医学影像, 图像分割

论文信息

论文题目：nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation

作者提出 nnU-Net，也就是 no-new-Net。这个名字本身就表达了论文态度：不要急着给 U-Net 添加 residual、dense、attention 等复杂模块，先把普通 U-Net 在医学图像分割中的整套流程系统地做好。

论文评测背景是 Medical Segmentation Decathlon。这个挑战要求同一个方法适配多个差异很大的医学分割任务，包括不同器官、不同模态、不同图像尺寸和不同数据量，并且不允许人为针对每个数据集手动调参。

核心动机：问题不一定在网络结构

医学图像分割里，U-Net 已经是非常强的基线。但很多论文会在 U-Net 上加新模块，然后报告性能提升。nnU-Net 作者质疑的是：这些提升到底是新结构真的更强，还是因为原始 U-Net 没有被充分优化？

作者认为，很多所谓结构创新可能是在一个没有调好的 baseline 上获得收益。比如预处理、重采样、patch size、batch size、数据增强、推理策略和后处理如果没有认真设计，U-Net 本身就会低估。此时再加入一个新模块，当然容易看起来有效。

因此，nnU-Net 的出发点不是“U-Net 不够强”，而是：

大家没有系统地、自动地把 U-Net 用好。

它把重点放在非结构因素上：预处理、网络拓扑自动配置、训练策略、推理策略、模型选择和后处理。论文想证明的是，医学图像分割性能很大程度来自完整 pipeline，而不是单个花哨模块。

nnU-Net 总体流程框架图

为什么必须自动适配

医学图像数据的差异远大于普通自然图像。不同任务之间，图像尺寸、体素间距和模态都可能完全不同。

nnU-Net 自动生成的网络拓扑

表中可以看到几个非常典型的差异：

  • Liver 的 median shape 是 482 x 512 x 512，体积非常大；
  • Hippocampus 只有 36 x 50 x 35，几乎可以被一个 3D patch 覆盖；
  • Prostate 是 20 x 320 x 319，z 轴很薄，明显各向异性；
  • Lung、Pancreas、Heart 又有自己的尺寸分布。

这意味着一个固定配置的 U-Net 不可能适合所有任务。nnU-Net 要自动决定：

  • 使用 2D U-Net、3D U-Net，还是 3D U-Net Cascade；
  • 输入 patch size 应该多大；
  • batch size 应该是多少；
  • 每个轴应该下采样几次；
  • 是否需要 low-resolution 阶段补足全局上下文。

这正是表格的意义：这些配置不是人工逐任务调出来的，而是 nnU-Net 根据数据集 fingerprint 自动生成的。

三类网络：2D、3D 和 Cascade

nnU-Net 不是单个网络，而是一组简单 U-Net：

2D U-Net
3D U-Net
3D U-Net Cascade

它没有引入复杂结构。相比原始 U-Net，只做了很小的改动，比如使用 Leaky ReLU 和 Instance Normalization。真正关键的是：不同数据集适合不同使用方式。

2D U-Net 适合强各向异性数据。比如某些 MRI 数据 x-y 平面分辨率高，但 z 轴切片很厚，强行做 3D 卷积可能会把不可靠的跨层信息混进去。

3D U-Net 适合三维结构比较可靠、图像体积不太大的数据。它可以同时利用 x、y、z 三个方向的信息，学习器官和病灶的三维形态。

3D U-Net Cascade 主要解决大图像上的显存和视野矛盾。

2D、3D 与 Cascade 的选择对比图

U-Net Cascade

Cascade 分两阶段：

  1. Stage 1：先把图像降采样，在低分辨率下训练 3D U-Net。这样网络能看到更大的全局范围，得到一个粗分割。
  2. Stage 2：把粗分割上采样回原分辨率，作为额外输入通道，再用 full-resolution 3D U-Net 精修边界和细节。

直观理解是：

第一阶段解决“看得远”
第二阶段解决“看得清”

这对 Liver、Lung、Pancreas 这类大体积任务很重要。普通 3D U-Net 受显存限制只能看局部 patch，容易缺少全局解剖上下文。

预处理：医学图像分割的关键工程

nnU-Net 的预处理包括三件事：cropping、resampling 和 normalization。

Cropping 会裁掉 nonzero region 之外的空白区域。对 skull-stripped brain MRI 这类图像尤其有效，因为脑外大量区域可能都是 0。裁剪可以减少无效背景，降低计算负担。

Resampling 会把同一个数据集内的病例重采样到 median voxel spacing。医学图像里的 voxel spacing 有真实物理意义，CNN 本身却不知道每个 voxel 对应多少毫米。如果不同病例 spacing 不统一，网络看到的同样数组大小可能对应不同真实尺度。

重采样时：

  • 图像使用三阶样条插值；
  • 标签使用最近邻插值，避免产生无意义的小数标签；
  • 如果数据集太大，nnU-Net 会自动判断是否启用 cascade；
  • 如果数据各向异性，先处理高分辨率轴，再整体降采样。

Normalization 会区分 CT 和 MRI。CT 的 HU 值有物理含义，所以使用数据集级别统计信息归一化；MRI 灰度没有统一物理尺度，所以通常按每个 patient 单独做 z-score normalization。

这一节的重点是：预处理不是附属细节，而是决定模型是否公平、稳定、可泛化的关键部分。

训练流程：稳健规则优先

nnU-Net 所有模型都从零开始训练，并使用五折交叉验证。训练损失是 Dice loss 和 Cross Entropy loss 的组合：

nnU-Net 训练损失公式

用公式表示就是：

L_total = L_dice + L_CE

Cross Entropy 负责逐 voxel 分类稳定性，Dice loss 负责区域重叠，对前景很小、背景很多的医学分割尤其重要。

训练中的几个关键策略：

  • 使用 Adam，初始学习率为 3e-4；
  • 一个 epoch 定义为 250 个 training batches；
  • loss 长时间不改善时自动降低学习率；
  • 验证 loss 长时间不改善并且学习率足够小时自动停止训练；
  • 使用随机旋转、缩放、弹性形变、gamma correction、镜像等数据增强；
  • 对各向异性 3D patch，改用 slice-wise 的 2D 数据增强；
  • cascade 第二阶段会随机扰动第一阶段的 segmentation，避免第二阶段过度依赖完美粗分割；
  • patch sampling 会保证 batch 中至少一部分样本包含前景类别。

这些策略都不是特别花哨，但组合起来非常实用。nnU-Net 的强度来自“每个环节都不掉链子”。

nnU-Net v2 的安装

nnU-Net v2 实践路线图

现在常用的是 nnU-Net v2。安装时建议先装 PyTorch，再装 nnU-Net，因为 PyTorch 必须匹配本机 CUDA 和 GPU 环境。

conda create -n nnunet python=3.10 -y
conda activate nnunet

根据自己的 CUDA 版本安装 PyTorch。例如 CUDA 12.1 环境可能使用：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

然后安装 nnU-Net v2：

pip install nnunetv2

如果要改源码或做研究开发，可以源码安装：

git clone https://github.com/MIC-DKFZ/nnUNet.git
cd nnUNet
pip install -e .

配置工作目录

nnU-Net v2 需要三个目录：

nnUNet_raw            原始数据
nnUNet_preprocessed   预处理后的数据
nnUNet_results        训练结果和模型权重

Linux 或 WSL 中可以这样配置：

export nnUNet_raw="/path/to/nnUNet_raw"
export nnUNet_preprocessed="/path/to/nnUNet_preprocessed"
export nnUNet_results="/path/to/nnUNet_results"

Windows PowerShell 中可以这样配置：

$Env:nnUNet_raw = "D:/nnUNet_raw"
$Env:nnUNet_preprocessed = "D:/nnUNet_preprocessed"
$Env:nnUNet_results = "D:/nnUNet_results"

实际训练时更推荐 Linux 或 WSL2 + NVIDIA GPU。Windows 原生也能跑，但路径、依赖和 CUDA 组合更容易出问题。

数据集格式

nnU-Net v2 的数据集通常放在 nnUNet_raw 下：

nnUNet_raw/
└── Dataset001_MyDataset/
    ├── dataset.json
    ├── imagesTr/
    ├── labelsTr/
    └── imagesTs/

单模态 CT 示例：

imagesTr/case_001_0000.nii.gz
labelsTr/case_001.nii.gz

多模态 MRI 示例：

imagesTr/case_001_0000.nii.gz
imagesTr/case_001_0001.nii.gz
labelsTr/case_001.nii.gz

注意：图像文件带 _0000、_0001 这样的通道编号，标签文件不带通道编号。

dataset.json 示例：

{
  "channel_names": {
    "0": "CT"
  },
  "labels": {
    "background": 0,
    "liver": 1,
    "tumor": 2
  },
  "numTraining": 100,
  "file_ending": ".nii.gz"
}

标签值一般是：

0 = background
1 = class 1
2 = class 2
...

最容易出错的是命名、标签值、通道编号和 dataset.json 不一致。

规划、预处理、训练和推理

数据准备好后，先运行规划和预处理：

nnUNetv2_plan_and_preprocess -d 1 --verify_dataset_integrity

其中 -d 1 对应 Dataset001_MyDataset。第一次跑建议加 --verify_dataset_integrity，这样能提前发现数据格式问题。

训练 3d_fullres：

nnUNetv2_train 1 3d_fullres 0

含义是：

数据集 ID = 1
配置 = 3d_fullres
fold = 0

完整五折训练通常是：

nnUNetv2_train 1 3d_fullres 0
nnUNetv2_train 1 3d_fullres 1
nnUNetv2_train 1 3d_fullres 2
nnUNetv2_train 1 3d_fullres 3
nnUNetv2_train 1 3d_fullres 4

如果后续需要自动选择配置或做 ensemble，训练时可以加 --npz：

nnUNetv2_train 1 3d_fullres 0 --npz

推理时，输入文件夹里的图像仍然要带通道编号：

input/
├── case_101_0000.nii.gz
├── case_102_0000.nii.gz

预测命令：

nnUNetv2_predict -i input -o output -d 1 -c 3d_fullres

如果只训练了某一个 fold，可以指定：

nnUNetv2_predict -i input -o output -d 1 -c 3d_fullres -f 0

实践建议

刚开始用 nnU-Net 时，建议先跑最稳的 baseline：

nnUNetv2_plan_and_preprocess -d 1 --verify_dataset_integrity
nnUNetv2_train 1 3d_fullres 0
nnUNetv2_predict -i input -o output -d 1 -c 3d_fullres -f 0

如果数据很大，nnU-Net 可能会生成 3d_lowres 或 cascade 相关配置。完整使用 cascade 时，通常先训练 low-resolution 阶段，再训练 cascade full-resolution 阶段。

目前 nnU-Net v2 也提供 residual encoder presets，通常能作为更强 baseline，但显存需求更高。普通实验建议先用 3d_fullres 建立可靠基线，再考虑更大的 preset。

总结

nnU-Net 的价值不在于提出一个复杂网络，而在于把医学图像分割中大量容易被忽视的工程决策系统化：

  • 根据数据集自动决定网络配置；
  • 根据体素间距和图像尺寸自动预处理；
  • 针对 2D、3D、cascade 使用稳健训练策略；
  • 使用统一的交叉验证、推理和模型选择流程；
  • 让 U-Net 成为一个强而可靠的自动化 baseline。

如果只记住一句话，可以是：

nnU-Net 不是一个新网络，而是一套把 U-Net 用到位的自动化医学图像分割框架。

参考

See also