nnU-Net训练测试数据集全流程:从目录构建到推理避坑指南
发布时间:2026/10/2 5:51:56来源:尧图网络
简介这套数据集源自马萨诸塞道路遥感数据集已按 nnU-Net 常用目录结构整理包含训练图像、训练标签、测试图像与测试标签非常适合用于道路提取、遥感图像分割等任务的模型训练与效果验证。数据集中图像为高分辨率航空影像道路标注来自 OpenStreetMap 栅格化处理覆盖城市、郊区、乡村等多种场景道路与周边地物对比清晰非常适用于训练和评估深度学习分割模型。压缩包共 59 个文件其中 58 张 PNG 同时包含原始影像与对应标签掩膜另含 1 个 dataset.json 配置文件用于描述数据集基本信息、模态和类别设置整体大小 142.02 MB。整体目录划分明确可直接对应 nnU-Net 预处理流程节省数据整理与格式转换时间适合需要开展遥感道路分割实验或学习 nnU-Net 数据组织方式的研究者。目前已有 1772 人学习浏览与下载。1. nnU-Net 训练测试数据集从目录结构到推理的一整套闭环用 nnU-Net 跑过一次图像分割训练的人大概率都有这种体验标注画完了模型结构也选好了结果卡在最不性感的一步——数据集整理。目录结构摆错、label 编号和训练代码对不上、预处理跑到一半 OOM随便一个都能耗掉大半天。这篇笔记把 nnU-Net 训练测试数据集的完整流程拆开讲透从任务目录怎么建、dataset.json 怎么写到 plan_and_preprocess 的参数含义、五折交叉验证和 nnUNet_predict 推理命令最后列几条我实际翻过车的坑。适合正在做 MRI/CT 分割、需要快速跑通 nnU-Net baseline 的研究生和算法工程师也适合想把 nnU-Net 自动规划黑匣子看明白的从业者。内容以 3D 医学数据为例2D 场景同样适用。2. 数据集目录与 dataset.json输入规矩立不住后面全是坑nnU-Net 对数据组织的要求非常死板好处是只要目录结构对了后续的预处理、训练、验证全自动。先说版本差异老版本用 Task501_XXX 命名任务新版本统一用 Dataset501_XXX下面以 Dataset 前缀为例说明命令里的-t 501取的是数字编号。2.1 任务目录结构三级目录与命名规范顶层nnUNet_raw_data_base下有三个核心目录nnUNet_raw放原始数据nnUNet_preprocessed放预处理产物nnUNet_trained_models放训练权重。原始数据集放在nnUNet_raw/Dataset501_BrainTumor下内部固定包含imagesTr训练图像、labelsTr训练标签和可选的imagesTs测试图像。nnUNet_raw_data_base/ ├── nnUNet_raw/ │ └── Dataset501_BrainTumor/ │ ├── imagesTr/ │ │ ├── case_000_0000.nii.gz │ │ ├── case_001_0000.nii.gz │ │ └── ... │ ├── labelsTr/ │ │ ├── case_000.nii.gz │ │ └── ... │ └── imagesTs/ │ ├── case_test_0000.nii.gz │ └── ... ├── nnUNet_preprocessed/ └── nnUNet_trained_models/目录层级是 nnU-Net 的数据输入硬约束路径错一层训练直接找不到数据。逻辑说明imagesTr里的文件命名带_0000后缀这是模态序号多模态时按_0000、_0001顺序递增labelsTr里的文件与图像同名但不带模态后缀。参数说明前缀case_000可以任意取名但同一个数据集内必须唯一_0000模态序号必须从 0 开始连续否则 nnU-Net 加载数据时会把多个模态当成同一卷。训练前先跑一个校验脚本把缺失和多余的 label 一次性找出来import os, glob root /path/to/nnUNet_raw/Dataset501_BrainTumor imgs sorted(glob.glob(os.path.join(root, imagesTr, *.nii.gz))) lbls sorted(glob.glob(os.path.join(root, labelsTr, *.nii.gz))) img_ids {os.path.basename(p).replace(_0000.nii.gz, ) for p in imgs} lbl_ids {os.path.basename(p).replace(.nii.gz, ) for p in lbls} print(missing:, img_ids - lbl_ids) print(extra:, lbl_ids - img_ids) print(samples:, len(img_ids))逻辑说明把图像和标签文件名都转成集合求差集就能定位所有对不上的案例。参数说明这段脚本只适配单模态数据多模态时要把_0000、_0001逐个去除或者用正则re.sub(r_\d{4}\.nii\.gz$, , name)一步搞定。这个检查我在多个项目里都跑过filesystem 层面的错位大多能提前暴露。2.2 dataset.jsonlabels 编号决定训练生死dataset.json放在 Dataset 根目录下nnU-Net 读取它才知道类别数量、模态类型和标签映射关系。与 YOLOv8 那套 txt 标注体系完全不同nnU-Net 不接受任何矢量标注只认 nii.gz 逐像素掩码。这里最关键的字段是labels背景必须固定在 0其余类别从 1 开始连续递增中间不能跳号。import json labels {background: 0, edema: 1, tumor_core: 2, enhancing: 3} modalities {0: MR_T1, 1: MR_T2} # 两个模态 info { name: Dataset501_BrainTumor, description: example for nnU-Net training, tensorImageSize: 3d, labels: dict(sorted(labels.items(), keylambda x: x[1])), numTraining: 0, modalities: modalities, file_ending: .nii.gz } with open(dataset.json, w) as f: json.dump(info, f, indent4)逻辑说明sorted按 value 排序确保 background 永远排在前面numTraining写 0 也没关系nnU-Net 会按imagesTr实际文件数在预处理时覆盖。参数说明labels的 value 必须是没有空洞的连续整数这是训练阶段 Dice 计算的硬前提modalities的 key 是字符串形式的0、1value 建议写CT、MR_T1这类可读名nnU-Net 内部会据此决定归一化策略。提示掩码里只有 0 和 2、没有 1训练不会报错但验证 Dice 会一直在 0 附近。这种坑极难察觉我在这里浪费过两天时间。写 dataset.json 前务必先打印 label 文件的 unique 值核对。3. 预处理与自动配置nnUNet_plan_and_preprocess 到底在做什么数据集整理完下一步是nnUNet_plan_and_preprocess。这步很多人当一键预处理跑完就交差实际上它做了指纹提取、方案规划、数据重采样三件事输出直接决定训练时的显存占用和精度上限。对训练标注来说nnU-Net 对掩码格式的刚性要求远超分类网络预处理阶段就是把这些刚性要求落到数据上。3.1 命令参数与三个阶段先设置三个环境变量再执行规划命令。export nnUNet_raw_data_base/data/nnunet export nnUNet_preprocessed/data/nnunet/nnUNet_preprocessed export RESULTS_FOLDER/data/nnunet/results nnUNet_plan_and_preprocess -t 501 -pl3d ExperimentPlanner3D_v21 -tl 2 -tf 2参数说明-t 501是 Dataset 编号对应Dataset501_BrainTumor-pl3d是 3D 规划器当前固定用 v21 版本-tl 2和-tf 2分别控制预处理和指纹提取线程数内存紧张时调到 2线程开太高容易在读取大体积 nii.gz 时把内存打满。整个流程三个阶段先扫描全部训练样本的 spacing、shape、强度分布生成 dataset_fingerprint.json再根据指纹决定 target spacing、patch size、batch size写出 plans.json最后按 plans.json 对每个样本做 crop、重采样、归一化输出到nnUNet_preprocessed。预处理产物看一眼目录结构ls /data/nnunet/nnUNet_preprocessed/Dataset501_BrainTumor/ # dataset_fingerprint.json plans.json splits_final.json # folder_0/ folder_1/ ... folder_4/逻辑说明splits_final.json是五折划分结果训练时完全按它加载数据如果你对划分不满意比如同一个病人的多期扫描散在不同折里必须在训练前手动替换这个文件。folder_0到folder_4里存放的是重采样后的图像和标签训练不再直接读原始数据。3.2 plans.json 关键参数与三种配置plans.json 是预处理和训练的总纲。对多数使用者来说只需要关注四类字段patch_size、batch_size、target_spacing、median_image_size_in_voxels。nnU-Net 默认生成三种配置各有分工配置典型 patch size典型 batch size适用场景2d512×51212单张切片独立分析的 2D 任务3d_fullres128×128×1282常见 CT/MRI 体积分割3d_lowres32×192×1608各向异性强、低分辨率轴维度大的数据三种配置都会在 plan 阶段生成训练时不一定全跑。资源充足时优先跑3d_fullres五折资源紧张时先跑3d_fullres一折看曲线。低分辨率版主要用来配合 fullres 做模型集成单独用精度通常差一截。显存和 patch size 的对应关系里3d_fullres的 128³ patch 配 batch_size 2 大约需要 11GB 显存24GB 卡上比较从容8GB 卡就得按第 6 章的方法手动调参。3.3 被自动决定的归一化与重采样策略nnU-Net 的自动不只是选 patch size。它还会根据模态名决定归一化策略CT 数据按窗宽窗位裁剪后归一化到 01MR 数据按 z-score 标准化。对 spacing 不均的数据集它会统一重采样到中位数 spacing预处理后整体 shape 趋向一致。黑匣子不是没有代价的plan 阶段自动把数据缩放到中位数 spacing如果数据集里某几个样本 spacing 特别离谱中位数被带偏所有样本都会被 resample 到错误间距。常见做法是在 plan 之前自己统计所有样本的 spacing 分布异常值提前剔掉或重采样。标注时注意同一批病例的采集参数尽量统一否则预处理阶段会替你做很多额外工作。注意不想让 nnU-Net 替你决定时可以手动写入 plans.json 的target_spacing。但改完必须重跑预处理只想调训练超参时别动这个字段。4. 五折交叉验证与推理nnUNet_train / predict 的命令闭环预处理完成splits_final.json把数据切成五折。nnU-Net 默认跑 5 折交叉验证每一折训练一个模型五个模型汇总的指标才可信最终推理也可以五模型 ensemble。这一章把训练、汇总、推理三个命令串成一个完整闭环。4.1 训练命令与 checkpoint 管理训练命令的四个位置参数分别是配置名、trainer 类名、任务号、fold 编号。export RESULTS_FOLDER/data/nnunet/results nnUNet_train 3d_fullres nnUNetTrainerV2 501 0参数说明第一个位置参数3d_fullres是配置名可选2d、3d_fullres、3d_lowres第二个nnUNetTrainerV2是 trainer 类名第三个 501 是任务号第四个 0 是 fold 编号范围 04。想跑全部五折可以串行写0,1,2,3,4循环执行也可以直接传5nnU-Net 会把 5 解释成 all folds。训练过程每个 epoch 结束都会写盘checkpoint 有两个关键文件model_final_checkpoint.model是训练结束时的状态model_best.model是验证 loss 最低点的状态。推理时 nnU-Net 默认优先用 best如果训练过程验证 loss 曲线持续下降没有拐点两者差别不大。训练中断不会丢失已保存的 epoch重新执行同样的命令并加-r参数可以从最后一个 checkpoint 恢复不加-r就会从头开始跑这个细节很多人栽过。长时间训练建议用 tmux 保活tmux new -s nnunet nnUNet_train 3d_fullres nnUNetTrainerV2 501 0 # Ctrl-b d 退出当前会话 tmux attach -t nnunet逻辑说明tmux 把训练进程挂到后台关掉终端也不影响训练。参数说明-s nnunet是会话名后续 attach 和 kill 都靠这个名字定位。4.2 五折指标汇总与最优配置选择五折全部跑完后用汇总命令选出最优配置。nnUNet_find_best_configuration 501 -c 3d_fullres 3d_lowres 2d这个命令读取RESULTS_FOLDER下所有已完成的 fold 结果汇总每折验证集 Dice输出最优配置名和对应平均 Dice。-c参数列出你想对比哪些配置只跑了3d_fullres就只传3d_fullres传了缺失的配置它会一直等待。想看每类的详细指标到对应 fold 目录打开validation.json里面有逐病例的 per-case 结果summary.json则汇总了各类别的平均 Dice 和标准差。五折的验证指标波动能反映数据集本身的稳定性——如果某一折明显偏低优先检查该折的数据划分是否混入了异常样本。4.3 nnUNet_predict 推理与后处理训练验证完毕对imagesTs里的测试图像做推理。测试图像同样需要_0000模态后缀且模态顺序必须和训练时一致。nnUNet_predict \ -i /data/imagesTs \ -o /data/pred \ -t 501 \ -m 3d_fullres \ -f 0 1 2 3 4 \ --step_size 0.8参数说明-i是测试图像目录-o是输出目录-t是任务号-m是配置名。-f 0 1 2 3 4表示五折模型一起 ensemble单卡显存紧张时可以只传-f 0先出单模型结果。--step_size 0.8是滑动窗口步长比例默认 0.5步长越大窗口重叠越少、推理越快但边界精度略降。输出的是未后处理的原始预测概率图。推理完成后还要做后处理和训练阶段的验证指标对齐nnUNet_apply_postprocessing \ -i /data/pred \ -o /data/post \ -pp_pkl_file /data/nnunet/results/nnUNet/3d_fullres/Dataset501_BrainTumor/.../postprocessing.pkl \ -plans_json /data/nnunet/nnUNet_preprocessed/Dataset501_BrainTumor/plans.json逻辑说明后处理主要移除预测结果中过小的连通域比如一个只有几千体素的小块在整个 mask 里通常不是目标器官。参数说明-pp_pkl_file指向训练阶段生成的 postprocessing.pkl路径比较深直接用find定位最稳妥。提示五模型 ensemble 是 nnU-Net 的标准用法但不是唯一用法。显存紧张就少传几个 fold先跑-f 0出结果后面有时间再补其余 fold。5. nnU-Net 避坑清单五条翻车记录与排查路径以下五条全部来自我或同事实际翻车的记录按「现象 → 原因 → 解决」写训练前对着过一遍能省至少三天。每一条都配了排查路径和验证方法。5.1 标签编号有空洞Dice 全程为零现象训练能正常走loss 也在下降但验证集 Dice 始终在 0 附近后处理也没有任何改善。最先怀疑模型有问题换网络、调学习率都没效果。原因掩码像素值只有 0 和 2没有 1。nnU-Net 的损失函数把类别索引当成类别编号映射全部错位Dice 自然算不准。这种空洞在标注工具导出时很常见尤其是只标注了部分类别的半成品数据集。解决写脚本重映射 label确保从 0 开始的连续整数。import nibabel as nib import numpy as np path label.nii.gz img nib.load(path) data img.get_fdata().astype(np.int16) print(unique before:, np.unique(data)) # 先确认空洞 remap {5: 1, 10: 2, 15: 3} # 原始值 - 目标值 for src, dst in remap.items(): data[data src] dst nib.save(nib.Nifti1Image(data, img.affine, img.header), label_fixed.nii.gz)逻辑说明重映射前先打印 label 的 unique 值确认和 dataset.json 的映射一致背景 0 保持不变。参数说明remap字典按实际数据调整覆盖写入前先备份原始文件这条我在一个 300 例的数据集上吃过亏重映射逻辑写反把背景覆盖了。5.2 imagesTr 和 labelsTr 文件名错位现象splits_final.json生成后训练日志显示的样本数比imagesTr实际文件数少一半或者某个 fold 的验证集为空导致该折指标是 NaN。原因label 文件缺失或者文件名多了一个空格、后缀nnU-Net 在 plan 阶段按文件名配对时悄悄跳过。解决用第 2 章的集合差脚本过一遍补上缺失文件并检查是否有隐藏副本比如 macOS 的._开头文件。nii.gz 文件大小如果全是几 KB 的大概率是坏文件用 nibabel 逐个读一次无法正常加载的提前剔除。5.3 3D 数据被误判成 2D验证效果离奇现象plans.json 里 2d 配置的 patch size 变成 512×512×13d_fullres 的 patch size 的 z 维度特别小整个训练输出看起来像在跑 2D。原因数据里 z 轴 spacing 异常大比如几百毫米一层。dataset fingerprint 统计出的中位数间距把低分辨率轴标记成切片方向于是 3D 任务被规划成 2D 任务。解决查看dataset_fingerprint.json里的spacings字段用 nibabel 读出所有样本的 spacing把异常样本重采样或直接从训练集剔除。常见做法是把 z 轴 spacing 控制在 x/y 的 25 倍以内超出就先重采样再进预处理。5.4 训练阶段 CUDA OOM现象第一个 epoch 刚开始就报CUDA out of memory或者训练到一半崩掉重启后又在同一位置崩。这个和预处理 OOM 还不一样是纯显存问题。原因显存本身不够plans.json 里 batch_size 对当前数据偏大或者有其它程序在占显存。Windows 下显存碎片化更严重同样配置更容易崩。解决先nvidia-smi看显存占用再改 plans.json 里 3d_fullres 的batch_size从 2 降到 1patch_size往下调 8 的倍数128 改成 112 或 96。改完必须重新执行nnUNet_plan_and_preprocess否则训练仍按旧 plan 加载数据形状对不上。5.5 改 plans.json 没生效白等三天现象改完 plans.json 里的 patch size 和 batch size重新发起训练日志里显示的 patch size 还是旧值训练跑到第三天出了 NaN。原因训练读的是nnUNet_preprocessed目录下已经生成的 plans.json 副本你在原始 Dataset 目录里改的那份不是它真正读的文件。解决把nnUNet_preprocessed/Dataset501_BrainTumor/plans.json删掉重新预处理或者在 preprocessed 目录里直接改、改完保存并重跑预处理。我一般直接在 preprocessed 目录里改改完看一眼文件 mtime确认新 plans 生成时间晚于训练启动时间才放心。6. 手动覆盖 plans 与自定义推理把默认配置改造成自己的工作流nnU-Net 的默认配置适合多数场景但遇到显存只有 8GB、或者需要快速出中间结果的情况还是得手动干预。我常用的两板斧是「改 batch size 硬塞进小显存」和「调滑动窗口步长提速」。先看小显存方案。打开 preprocessed 目录的 plans.json找到 3d_fullres 下的batch_size把 2 改成 1patch_size也按 8 的倍数往下调。注意 2d 和 3d_lowres 配置也同步调整否则后面集成时模型资源占用不一致。改完保存并重跑预处理训练命令一行不用变。8GB 卡跑 3d_fullresbatch size 1 加 patch 112×112×96 是验证过能稳的组合再低就建议直接跑 2d 配置别硬撑体积分割。再看推理阶段。显存紧张时把--step_size从默认 0.5 提到 0.8推理时间几乎减半边界 Dice 通常只掉 0.1 到 0.3 个百分点。数据分布稳定时开启测试时增强nnUNet_predict -i /data/imagesTs -o /data/pred -t 501 -m 3d_fullres \ -f 0 1 2 3 4 --step_size 0.8 --tta加了--tta后推理时间约翻倍但分割边界更规整适合小目标任务。五折 ensemble 与单折相比Dice 平均提升 0.51 个百分点是常态折叠数越全结果越接近真实泛化水平。有个项目我图省事改完 plans.json 没重跑预处理训练到第三天出了 NaN一查是数据形状和旧 plan 不匹配三天的卡白烧。从那以后我每次改 plans.json 都强制检查 preprocessed 目录的 mtime确认新的 plans 文件生成时间晚于训练启动时间才开跑。希望帮到你。本文还有配套的精品资源点击获取
网站建设高端定制企业官网