nnUnet源码解析

type

status

date

slug

summary

系统框架

实验规划：分析数据集属性并生成管道配置

预处理：准备训练数据集（规范化、重采样等）

模型训练：使用配置的设置训练模型

模型评估：计算训练模型的性能指标

最佳配置选择：确定最佳模型或集成

推理：将选定的模型应用于新数据

安装

硬件要求

CPU

现代多核处理器

内存

最低：16GB
建议：32GB 或更多（特别是对于包含大图像的 3D 数据集）

图形处理器

用于训练：NVIDIA GPU，至少具有 11GB VRAM（RTX 2080 Ti、3090、4090、A5000 或更高版本）
仅用于推理：配备 8GB+ VRAM 的 NVIDIA GPU，或仅使用 CPU（速度明显较慢）

存储空间

至少 100GB 可用空间（根据数据集大小而变化）

软件要求

Python

3.10 或更高版本

操作系统

Linux（推荐，尤其是 Ubuntu）
Windows 10/11
macOS（通过 MPS 或仅 CPU 提供有限的 GPU 支持）

CUDA 和 cuDNN

GPU 加速所需（兼容 PyTorch 2.1.2+）

安装方法

通过pip安装

安装 nnU-Net v2 最简单的方法是使用 pip：

这将自动安装 nnU-Net v2 及其所有依赖项，如pyproject.toml32-55文件。

从 GitHub 仓库安装

对于最新的开发版本或者如果您想为代码库做出贡献：

核心模块代码解读

1.训练数据预处理

nnU-Net v2 中的预处理系统将原始医学影像数据转换为适用于神经网络训练和推理的标准化输入。

预处理系统的核心是DefaultPreprocessor类，它协调所有预处理操作。

代码文件为：preprocessing/preprocessors/default_preprocessor.py

a.图像加载与转置第一步使用计划中指定的读取器加载图像并应用轴转置以确保方向一致：

b.图像裁减

裁剪通过删除没有相关信息的背景区域来减少内存需求,nnU-Net 记录用于裁剪的边界框，以便在推理过程中进行逆转：

c.正则化

归一化使图像间的强度值标准化，使网络训练更加稳定。在重采样之前应用归一化，以确保插值的准确性：

d.重采样（各向异性处理）

重采样将体素之间的间距调整为配置中指定的目标间距：

重采样过程包括：

根据原始和目标间距计算新形状

应用适当的插值（图像和分割不同）

e. 前景采样以提高训练效率

对于分割任务，nnU-Net 对前景位置进行采样，以通过平衡的块采样实现高效的训练：

2.推理预处理

在推理过程中，预处理是即时执行的，流程如下：

3.用户单个case的图像预处理

实验自动化设置

设置的理念

nnUNet “无配置”理念指根据数据集属性自动设计和配置网络架构，平衡性能和硬件限制。该部分负责根据数据集特征自动配置神经网络架构和训练参数。实验规划系统会分析数据集属性和硬件约束，从而生成预处理、训练和推理的最佳设置。

实验自动化分析数据集特征以确定：

重采样的目标间距

网络架构和拓扑（池化操作、内核大小）

内存高效的patch和batch size的大小

适当的数据增强和预处理策略

代码文件位置：

文件中ExperimentPlanner几个主要职责：

读取数据集属性——分析数据集指纹以了解图像特征，如形状、间距和强度分布。

网络架构配置——根据数据集属性选择适当的网络深度、内核大小和特征图。

硬件感知优化——它估计 GPU 内存需求并调整补丁和批次大小以适应可用资源。

多配置规划它为 2D、3D 全分辨率和 3D 低分辨率方法创建配置计划。

设置的类型

配置	维度	解决	用例
2d	2D	全分辨率	训练速度快，适合高度各向异性的数据
3d_全分辨率	3D	全分辨率	适合中等大小的 3D 体积
3d_lowres	3D	分辨率降低	对于非常大的 3D 体积
3d_cascade_fullres	3D	全分辨率	3d_lowres 之后的第二阶段，用于大型数据集

系统还会根据GPU的现存自动化设计ResEncUnet的参数大小

系统规划	目标 GPU 内存	用例
nnUNetPlannerResEncM	8 GB	RTX 2080Ti、1080Ti等
nnUNetPlannerResEncL	24 GB	RTX 3090、RTX 4090、A5000
nnUNetPlannerResEncXL	40 GB	A100 40GB、A6000等

模型的训练

训练框架

训练模块是 nnU-Net 的核心组件，负责模型训练，协调从数据加载到模型验证的整个过程

代码文件位置：

训练系统的核心是nnUNetTrainer类，它协调整个训练过程。它提供了一个灵活的框架，可以扩展以适应不同的训练策略。在

主要职责：

初始化和管理网络架构

配置优化器和损失函数

管理数据加载器和增强管道

执行训练和验证循环

实施检查点和日志记录

和模型训练相关的文件都在nnUNetTrainer文件夹下面，比如优化器，网络结构定义，训练流程等。

模型结构

网络架构根据计划和配置动态构建。默认架构类似 U-Net，但可以自定义，该方法build_network_architecture负责按照计划构建适当的网络架构。

数据增强

nnU-Net 采用大量数据增强来提高模型泛化能力：

增强类型	示例
空间	旋转、缩放、弹性变形、镜像
强度	亮度、对比度、伽马校正
噪音	高斯噪声、高斯模糊
其他	模拟低分辨率、随机二进制算子

增强管道是根据数据集属性配置的：

损失函数

nnU-Net 结合使用了 Dice 损失和交叉熵（或基于区域的分割的 BCE），损失函数是根据任务类型（基于标签或基于区域的分割）以及是否启用深度监督来构建的

超参设置

默认情况下，nnU-Net 使用带有 Nesterov 动量的 SGD 进行优化和多项式学习率衰减：

可通过训练器变体获得替代优化器：

nnUNetTrainerAdam：使用 Adam 或 AdamW 优化器

nnUNetTrainerAdan：使用Adan优化器（需要安装adan-pytorch）

多卡训练

nnU-Net 支持使用 PyTorch 的 DistributedDataParallel (DDP) 进行分布式训练，当使用多个 GPU 时，批量大小分布在各个工作器上，并且系统处理跨设备的梯度同步。

多种训练方法

nnU-Net 提供了基础训练器的几种变体，以支持不同的用例：

训练变体	目的
nnUNetTrainer无深度监督	缺乏深度监督的训练
nnUNetTrainerAdan	使用 Adan 优化器
nnUNetTrainerAdam	使用 Adam 优化器
nnUNetTrainer_Xepochs	训练指定数量的 epoch
nnUNetTrainerBenchmark_5epochs	用于基准性能

变体系统允许轻松定制，而无需修改核心训练器，比如下面的变体，直接用True或者False设置即可：

开始训练

训练过程通常通过命令行启动：

关键参数：

DATASET_NAME_OR_ID：用于训练的数据集

CONFIGURATION：要使用的配置（例如，“2d”、“3d_fullres”）

FOLD：交叉验证倍数（0-4 或“全部”）

tr：自定义训练器类（默认值：'nnUNetTrainer'）

p：计划标识符（默认值：'nnUNetPlans'）

num_gpus：用于训练的 GPU 数量

模型的推理

nnU-Net 推理系统负责应用已训练的模型对新的医学图像进行预测，将原始输入数据转换为精确的分割图。

推理系统的核心类是nnUNetPredictor，它协调整个预测过程。它管理：

模型初始化——加载网络架构和权重

预处理协调——确保正确的图像准备

预测执行——运行滑动窗口算法

结果处理——将逻辑转换为最终分割

初始化配置

使用前必须nnUNetPredictor进行初始化，一般使用以下参数：

范围	默认	描述
`tile_step_size`	0.5	移动滑动窗口的量（0.5 = 50％重叠）
`use_gaussian`	True	是否应用高斯加权进行窗口混合
`use_mirroring`	True	是否通过镜像使用测试时间增强
`perform_everything_on_device`	True	处理期间是否将数据保留在 GPU 上
`device`	CUDA	计算设备（推荐使用 CUDA）

初始化预测器后，必须使用经过训练的模型对其进行配置：

推理方法

代码文件位置:nnunetv2/inference/predict_from_raw_data.py

推理系统提供了多种预测方法，每种方法适用于不同的用例：

方法	用例	优势	缺点
`predict_from_files()`	基于多个文件的图像的批量预测	最佳内存效率，并行处理	需要磁盘上的文件
`predict_from_list_of_npy_arrays()`	多张图片已作为数组加载	无需文件 I/O	更高的内存使用率
`predict_single_npy_array()`	单幅图像预测	最简单的 API	最慢，无并行化
`predict_from_data_iterator()`	自定义数据加载方案	最大的灵活性	更复杂的实现

推理系统的一个关键组件是滑动窗口预测机制，它可以处理可能无法一次性放入 GPU 内存的大型医学图像。

多折交叉预测

推理系统通过集成平均支持多个网络（通常来自不同的交叉验证折叠）的预测：

使用多重折叠进行初始化：use_folds=(0, 1, 2, 3, 4)

对于每个输入：

循环遍历所有模型权重

从每个模型生成预测

平均预测

这通常比使用单折叠产生更稳健的结果，但代价是预测时间更长。集成平均发生在logits级别（softmax/argmax之前），从数学上讲，这比平均分割更合理。

模型评估与最佳选择

nnU-Net v2 中的评估和模型选择系统提供了一个强大的框架，用于评估已训练分割模型的性能、选择最佳配置，并通过后处理和集成来改进结果。

评估指标

nnU-Net 的评估系统主要使用 Dice 系数作为主要性能指标，但也会计算：

Dice系数：测量预测和地面实况之间的空间重叠

交并比（IoU）：替代重叠度量

真正例（TP）假正例（FP）假负例（FN）真负例（TN）

模型最佳选择

nnU-Net 为每个数据集训练多个模型配置，并自动选择性能最佳的模型。选择过程包括评估单个模型和模型集成，以确定哪个模型能获得最高的 Dice 分数。

默认情况下，nnU-Net 会考虑以下配置进行评估：

配置	描述
2d	二维U-Net
3d_全分辨率	全分辨率 3D U-Net
3d_lowres	低分辨率的 3D U-Net
3d_cascade_fullres	3D U-Net 级联（低分辨率 → 全分辨率）

交叉验证结果收集

在模型选择之前，收集并合并所有交叉验证的结果：

将各个折叠的所有验证预测复制到统一文件夹中

根据事实评估收集到的预测

生成交叉验证性能的摘要

寻找最佳配置

代码文件位置：nnunetv2/evaluation/find_best_configuration.py

该find_best_configuration函数比较所有模型变体：

评估单个模型配置（2D、3D、级联等）

创建和评估模型集成（如果允许）

选择 Dice 分数最高的配置

确定最佳配置的最佳后处理

生成并保存推理指令

💡

本文由LaputaK算法团队撰写，如有任何疑问或发现疏漏之处，欢迎批评指正。同时，我们也诚挚邀请您与我们联系，共同探讨交流！