ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀,每一条都经过真实交付验证。

nnU-Net训练测试数据集全流程指南:从数据组织到推理调优

2026/10/2 8:55:12 拓冰建站 浏览量
nnU-Net训练测试数据集全流程指南:从数据组织到推理调优 简介一套面向深度学习道路提取任务的nnU-Net格式训练与测试数据集源自马萨诸塞道路遥感数据集Massachusetts Roads Dataset适合使用nnU-Net框架进行遥感图像分割的研究者与工程师。数据包共59个文件含58张PNG图像和1份dataset.json配置压缩后约142.02MB目录严格按nnU-Net规范组织imagesTr/imagesTs存放训练与测试影像labelsTr/labelsTs存放道路标注掩膜训练与测试集划分清晰默认命名符合nnU-Net预处理要求可直接被读取使用。影像覆盖城市、郊区和农村等多种地形道路标签由OpenStreetMap道路中心线栅格化生成适合检验模型在树木遮挡、阴影等复杂场景下的提取能力。目前已有1772人浏览学习拿到后无需格式转换即可用于道路分割实验也可作为遥感语义分割基准数据集对比不同算法。1. nnunet训练测试数据集为什么这个框架能“开箱即用”“nnunet训练测试数据集”这几个词放在一起容易让人误以为这是某个现成项目的操作手册。实际经历过的人都知道nnU-Net最反直觉的地方在于它不靠发明新网络结构而是靠一套自动配置流程把预处理、训练策略和推理参数按数据集特征给你定好。你只需要按它的规矩把数据集摆放整齐、写好dataset.json剩下的plan、train、predict三步命令就能跑起来。对做医学影像分割、又不想在参数调优上耗太多时间的团队来说这套流程省掉的不是几分钟是几周。我按自己惯用的落地路径来写数据集怎么组织、预处理怎么跑、训练怎么训、测试推理怎么出结果以及最容易返工的那几个坑。2. 把数据集整理成nnU-Net要的样子任务目录与JSON配置2.1 从原始标注到nnU-Net的数据组织labelsTr、imagesTr与dataset.jsonnnU-Net v2的数据组织是按“任务Task”管理的。每个任务是一个独立的文件夹放在nnUNet_raw_data目录底下命名规则是Task序号_任务名称比如Task001_liver或Task002_brain。序号就是后面plan和train命令里用的任务编号名称是给你自己看的两个部分用下划线连接。文件夹内部固定是imagesTr、labelsTr、imagesTs三个目录分别放训练图像、训练标签和测试图像命名中的Tr是training的缩写Ts是testing的缩写。nnU-Net只吃NIfTI格式.nii.gz其他格式需要先转这个没得商量。文件名要按“病例名_模态标识符_病例名”来组织。比如一个病例叫case_001的CT图像在imagesTr里就叫case_001_0000.nii.gz其中_0000代表模态序号0。标签文件在labelsTr里就叫case_001.nii.gz。这里最容易翻车的是图像文件和标签文件名必须完全对应但标签不能带模态后缀。你有一张case_001_0000.nii.gz就必须有一张case_001.nii.gz在labelsTr里否则训练时nnU-Net报“找不到配对”的错误。另一个容易忽略的规则是case_001这个名称本身不能再包含下划线因为nnU-Net解析文件名时靠下划线切分字段case名字里再带下划线会让解析逻辑直接混乱。下面的小脚本可以快速把散落的图像和标签整理成这个结构import shutil from pathlib import Path raw_root Path(raw_data) # 原始图像目录 label_root Path(raw_label) # 原始标签目录 target Path(nnUNet_raw/nnUNet_raw_data/Task001_demo) (target / imagesTr).mkdir(parentsTrue, exist_okTrue) (target / labelsTr).mkdir(parentsTrue, exist_okTrue) (target / imagesTs).mkdir(parentsTrue, exist_okTrue) for i, img_path in enumerate(sorted(raw_root.glob(*.nii.gz))): case_id img_path.stem.split(.)[0] # 取文件名当作病例名 tgt_img target / fimagesTr/{case_id}_0000.nii.gz shutil.copy(img_path, tgt_img) # 找到同名标签 lab_path label_root / f{case_id}.nii.gz if lab_path.exists(): shutil.copy(lab_path, target / flabelsTr/{case_id}.nii.gz) else: print(f缺少标签: {case_id})这段脚本做的事情遍历原始图像目录把每个.nii.gz文件复制到imagesTr并自动加上_0000后缀同时检查同名标签是否存在存在的复制到labelsTr不存在的单独打印出来提醒你补数据。我用它处理过一次几百例的肝脏分割数据唯一要注意的是case_id不能包含下划线以外的特殊字符否则训练阶段解析文件名时会异常。2.2 写dataset.json的字段channel_names、labels与file_ending每个Task目录下都要有一个dataset.jsonnnU-Net从plan阶段开始就读它。最核心的三个字段是channel_names、labels和file_ending。channel_names描述每个模态通道是什么单模态CT就是[CT]多模态比如MRI的T1和T2就写[T1,T2]顺序要和文件名里的_0000、_0001一一对应。labels是字典从0开始0必须是background后面的1、2就是你要分割的解剖结构。file_ending固定写.nii.gz。{ channel_names: { 0: CT }, labels: { 0: background, 1: liver, 2: tumor }, numTraining: 120, file_ending: .nii.gz }这里的numTraining要和imagesTr里实际图像数量一致不一致时plan阶段会警告。labels的键必须是字符串值是类别名建议用英文方便在日志里识别。多模态场景常见做法是把MRI的多个序列按固定顺序排列比如T1加权、T2加权、FLAIR分别作为_0000、_0001、_0002同一病例的所有模态都必须存在缺失一个模态时nnU-Net不会替你补全训练会以报错告终。如果你处理的是带子区域的分割任务比如肝脏和肝脏肿瘤分别标注还可以在dataset.json里额外写regions字段来描述子区域归属但这个字段是可选的新手阶段先不用管。2.3 检查图像与标签形状跑plan之前先做自检数据放好、json写完之后我不会立刻跑plan_and_preprocess而是先做一个自检把每个病例的图像和标签都load进来比对shape和spacing。因为nnU-Net在训练阶段会自动把图像裁剪到非零区域但如果你图像的尺寸和标签尺寸不一致等到预处理重采样之后才发现再回去改数据就晚了。下面这个脚本遍历一遍import nibabel as nib from pathlib import Path img_dir Path(nnUNet_raw/nnUNet_raw_data/Task001_demo/imagesTr) lab_dir Path(nnUNet_raw/nnUNet_raw_data/Task001_demo/labelsTr) for img_path in sorted(img_dir.glob(*.nii.gz)): case_id img_path.name.replace(_0000.nii.gz, ) lab_path lab_dir / f{case_id}.nii.gz if not lab_path.exists(): print(f标签缺失: {case_id}) continue img nib.load(img_path) lab nib.load(lab_path) if img.shape ! lab.shape: print(fshape不一致: {case_id} img{img.shape} lab{lab.shape}) else: print(fOK: {case_id} {img.shape})这一步花五分钟能避免后面plan阶段报一堆形状不匹配的错。shape不一致通常是因为原始标注软件把标签裁剪过或重采样过常见做法是把标签重新resample到图像的空间建议用SimpleITK的Resample插值方式选最近邻避免引入新的标签值。多模态数据还要把每个模态的shape也互相核对一遍我遇到过T1和T2序列采集范围不同导致shape不一致的情况这种病例要么裁到共同区域要么干脆剔除。3. 跑通预处理与训练从nnUNet_plan_and_preprocess到nnUNet_train3.1 安装与环境变量nnUNet_raw、nnUNet_preprocessed、nnUNet_results三件套nnU-Net v2用pip安装命令是pip install nnunetv2它的命令行工具前缀是nnUNet_。要跑起来必须先设置三个环境变量。nnUNet_raw对应上面放原始数据的目录nnUNet_preprocessed是预处理输出目录plan阶段生成的plans.json和预处理后的数组都落在那里nnUNet_results是训练结果目录权重和日志都在里面。这三个目录建议都放到同一块剩余空间足够的盘上因为预处理后的数据往往比原始数据大好几倍。export nnUNet_raw/data/nnUNet/nnUNet_raw export nnUNet_preprocessed/data/nnUNet/nnUNet_preprocessed export nnUNet_results/data/nnUNet/nnUNet_results这三个环境变量在每个终端都要export一遍如果忘记设置命令会直接提示找不到路径。我一般把它们写进~/.bashrc省得每次重开终端都要敲一遍。另外要注意nnunetv1和v2的命令不兼容网上一搜一大把旧教程里nnUNet_plan_and_preprocess的用法是v1风格v2里已经统一成nnUNet_plan_and_preprocess -t 任务ID这种写法不需要再手动指定数据集名称。3.2 nnUNet_plan_and_preprocess数据指纹与自动配置预处理阶段做的事情是扫描所有数据的统计信息形状、spacing、强度分布、前景占比然后在3d_fullres、3d_lowres、2d这些配置里选出适合当前数据的方案同时完成重采样、裁剪、z-score归一化。这一步的输出是预处理后的数据、plans.json和dataset.json副本。命令很简单nnUNet_plan_and_preprocess -t 1 --verify_dataset_integrity-t 1对应Task001那个序号--verify_dataset_integrity是额外做一次数据完整性校验会检查每个病例的图像和标签是否一一对应、尺寸是否一致比2.3节那个自检脚本更严格。跑完之后看nnUNet_preprocessed/Task001_demo目录里面应该出现一个plans.json和一堆预处理后的.npz、.npy文件。预处理时间取决于数据量和分辨率几十例CT大概十几分钟几百例高分辨率MRI可能要跑几个小时这个速度是正常的别看到进度条不动就去CtrlC。plan阶段不是所有配置都会跑。如果数据量和前景区域都不满足3d_lowres的条件nnU-Net会自动跳过低分辨率配置只跑2d和3d_fullres。这一步看起来像个黑匣子但它其实把传统调参过程里“图像应该重采样成什么大小”“loss怎么配”“要不要用级联”这些问题都自动回答了。对新手来说不需要手动干预对想干预的熟手plans.json里的preprocessor_name、resampling_scheme这些字段都可以改但改完要能说得清为什么否则不如不改。3.3 nnUNet_train开始训练配置、Trainer与fold参数预处理完成后就能训练了。nnUNet_train命令格式是nnUNet_train 3d_fullres nnUNetTrainerV2 1 0第一个参数是配置名3d_fullres是完整分辨率的三维模型数据少或显存小时用2d第二个参数是Trainer类一般就写nnUNetTrainerV2想换loss或改学习率就继承这个类写自己的Trainer第三个参数是任务ID第四个是fold。fold从0到4代表5折交叉验证的哪一折如果只想快速验证流程通不通可以用0这一个fold训完后续要出稳定结果再训全5个fold。模型权重默认写在nnUNet_results/Task001_demo/nnUNetTrainerV2__3d_fullres/fold_0下面文件名是model_best和model_final_checkpoint。训练日志在同一目录下的training_log_*.txt里每个epoch会打印loss、Dice等指标。显存不够时可以在命令里加--disable_deep_supervision减少显存占用但会牺牲一点精度也可以手动把batch size调小训练脚本默认会按数据大小自动选batch size手动改的话一般从2开始逐步往下试同时观察GPU利用率别掉到30%以下。3.4 用--val跑验证集先拿到训练阶段的客观指标训练完fold_0之后我想先看一眼这个fold在验证集上的表现再决定要不要把5个fold都训练完。做法是重新执行一次nnUNet_train但加上--val参数nnUNet_train 3d_fullres nnUNetTrainerV2 1 0 --val这个命令不会重新训练而是用刚才训练好的权重在fold_0的验证集就是训练时留出来的那20%上做推理然后输出一个validation文件夹里面是预测结果同时打印每个类别的Dice系数。这个验证指标的意义在于它是nnU-Net自己按数据划分算出来的和后面nnUNet_predict阶段你自己拿测试集去推的结果在数据分布上是一致的可以用来判断训练是否收敛、需不需要加大epochs或者换Trainer。如果你的数据类别极度不均衡这里还要关注一下每个类别的Dice而不是只看平均Dice某个类别Dice特别低通常是该类别样本太少需要回去检查标注质量或考虑用更大的学习率。4. 用nnUNet_predict做测试推理输入输出与后处理4.1 nnUNet_predict的标准命令与关键参数测试阶段的核心命令是nnUNet_predict。它读取一个输入文件夹里所有图像输出对应的分割结果。一个完整的调用是nnUNet_predict -i /data/test_images -o /data/test_pred \ -t 1 -m 3d_fullres -f 0 \ --save_npz参数含义-i是测试图像目录里面放和训练集同名的.nii.gz文件模态后缀_0000也是必需的-o是输出目录-t是任务ID-m是配置名-f是指定用哪几个fold的模型多个fold时写-f 0 1 2 3 4nnU-Net会把多个fold的结果取平均再argmax这是它比单模型更稳的重要原因5折ensemble的预测结果通常比单折高1到2个点Dice--save_npz会把softmax概率保存下来占空间但后续做集成或分析方便。如果训练时用的自定义Trainer还要加-tr nnUNetTrainerV2指定如果plan阶段生成了多个配置要选对你想用的那个。这些参数不一致时最常见的报错是“checkpoint not found”——因为你指定的fold没有训练过或者Trainer名对不上。排查时先看nnUNet_results/Task001_demo目录下有哪些配置和fold再回头对照命令里的参数。4.2 单张与批量推理输入目录结构决定成败nnUNet_predict按目录批量处理输入目录里的每个文件都会被当成一个病例。所以如果你只测一张图也得放在一个目录里不能直接给文件路径。输出目录不需要提前创建程序会自动建。推理完成之后输出目录里是case名加.nii.gz后缀的分割结果标签值对应dataset.json里定义的1、2等类别。结果读进Python里很简单import nibabel as nib import numpy as np pred nib.load(/data/test_pred/case_001.nii.gz) data pred.get_fdata().astype(np.int16) print(np.unique(data)) # 应该出现0、1、(2)这些标签这里一个容易踩的坑是输出的是label还是probabilities。默认输出的是标签图只有加了--save_npz才会额外生成npz概率文件。如果你拿到结果发现值全是0和1而训练时有两个目标类别多半是label id写错或者模型本身没训好。批量推理时建议分批拷贝输入文件比如一次放50例跑完一批再放下批这样万一某个文件有问题不会整个任务失败重来。4.3 TTA开不开、step_size怎么设nnU-Net推理时默认会启用test time augmentationTTA即对输入做几种空间变换后分别预测再取平均这个机制常被称为测试时训练。TTA通常能涨0.5到1个点Dice代价是推理时间翻倍。显存够、时间不急时我不关但如果是做线上服务或处理几百例的大批量数据加--disable_tta能省一半时间。这里我一般先开着TTA跑一遍测试集把结果存好再关掉TTA跑一遍对比指标差异如果差异不到0.3个点后续就关掉TTA换速度。step_size是滑窗推理的步长比例默认0.5表示窗口每次移动一半。对特别大的图像step_size调到0.8能加快推理但边界区域的分割会略差调小到0.3会变慢但更稳。这个参数是推理阶段少数值得手动调的因为它直接控制滑窗重叠程度对分割稳定性影响明显。调的时候观察一下输出mask的连通性如果目标区域出现很多细碎的碎片多半是step_size太大导致窗口间预测不一致。4.4 模型checkpoint的“后悔药”中断续训与权重选择训练中途断了不用从头跑。nnUNet_train会记录断点重新执行同样的训练命令会从最近的checkpoint继续。这里要分清两个文件model_best是验证集上表现最好的权重model_final_checkpoint是训练结束时的权重两个文件在同名目录下。如果训练到一半被kill掉只有model_final_checkpoint存在可以用它继续如果要推理优先用model_best还是model_final_checkpoint可以看验证指标常见做法是哪个在验证集上高用哪个。我在实际项目中习惯每个fold训完后把model_best单独复制一份到另一个目录避免后续误操作把结果覆盖。权重文件是几百MB级别不占空间但这个习惯救过我两次一次是重跑训练时把fold_0的结果覆盖了另一次是清理目录时误删了训练输出。nnU-Net的自动化程度再高也架不住手动操作失误备份永远是成本最低的后悔药。5. nnunet训练测试数据集的5个常见问题现象、原因与解决5.1 标签是RGB索引而不是单通道标签loss怎么都不降现象训练几十个epochloss在0.6左右不动Dice一直是0或极低。原因很多人从标注工具比如Labelme或某些Web标注平台导出的标签是RGB彩色图比如背景是(0,0,0)目标是(255,0,0)直接存成PNG转成NIfTI后每个像素是一个三维向量不是单一类别ID。nnU-Net读进去会把每个向量当成多通道标签来处理模型无法学习这种离散映射loss自然降不下去。解决把RGB标签映射成索引标签。常见做法是按颜色构建一个查找表把每个像素的RGB向量映射成整数类别import numpy as np from PIL import Image # 颜色到类别ID的映射按你的实际标注颜色填写 color_map { (0, 0, 0): 0, # background (255, 0, 0): 1, # liver (0, 0, 255): 2, # tumor } label_img np.array(Image.open(label.png)) # 形状 (H, W, 3) out np.zeros((label_img.shape[0], label_img.shape[1]), dtypenp.int16) for rgb, cls_id in color_map.items(): mask (label_img[..., 0] rgb[0]) (label_img[..., 1] rgb[1]) (label_img[..., 2] rgb[2]) out[mask] cls_id nib.save(nib.Nifti1Image(out, affine), label_mapped.nii.gz)关键是映射后要检查np.unique(out)只能出现0、1、2这些连续整数中间不能有缺口否则模型训练时类别数量和dataset.json定义对不上又产生新的报错。我遇到过标注工具把两种颜色混在一起的情况比如两个类别都含有(255,0,0)分量这时候要回到标注软件里重新修正靠脚本补不干净。5.2 plan_and_preprocess在计算数据指纹时中断现象nnUNet_plan_and_preprocess跑到某个病例时直接报错或Killed日志里没有明确报错信息。原因最常见的是图像非零区域dense比如整幅图都没有全零切片导致裁剪时申请的内存过大或者图像spacing异常某个维度spacing为0导致重采样计算溢出再就是内存不够被系统OOM kill。解决先看plans.json是否已经生成如果生成了一部分说明是在统计阶段之后就挂的。然后单跑一遍数据完整性校验nnUNet_plan_and_preprocess -t 1 --verify_dataset_integrity如果这步能过再把内存上限放开或者分批预处理。nnU-Net支持--mode参数把plan和preprocess拆开执行先plan再preprocess这样能定位到底是哪一步挂的。数据本身有问题的用2.3节的自检脚本把报错case找出来单独修掉再重跑。这一步最忌讳反复直接重跑完整命令日志一大片根本看不出哪个case有问题。5.3 重采样后目标器官变形严重spacing不一致先做归一化现象预处理后的图像和标签里目标器官的形状看起来和原始图像不一致细小的结构扭曲、断裂。原因nnU-Net的plan阶段会计算所有病例的median spacing把所有图像重采样到目标spacing。如果一部分病例是1x1x1mm体素另一部分是5x5x5mm体素重采样到中间值后后者被放大的倍数过大细节就糊了。解决在数据准备阶段统一spacing。常见做法是先用SimpleITK把所有图像重采样到同一个spacing范围比如CT统一到1x1x1mm然后再交给nnU-Net。注意标签重采样必须用最近邻插值用线性或三次样条会在边缘产生中间值破坏标签语义。重采样后还要再跑一次自检确认图像和标签的shape仍然一致。另外不同设备采集的同一器官数据spacing差异太大时不要硬塞进同一个数据集训练先按spacing分桶处理看各自的分割效果再决定是否合并。5.4 训练完找不到model_final_checkpointfold和命令没对应现象训练日志正常打印但nnUNet_results对应目录里是空的或者只有partial_checkpoint。原因训练进程被提前终止或者store_checkpoint的路径没写对。nnU-Net训练过程中会周期性保存partial_checkpoint只有正常跑完一个fold才会落model_best和model_final_checkpoint。如果你只是快速验证模型能跑通训练被CtrlC中断自然找不到完整权重。解决确认训练日志里出现了类似“Epoch done”的完整周期输出。如果只是验证流程建议把epochs调小、data少放点但最好还是让它完整跑完一个fold不然测试阶段nnUNet_predict找不到checkpoint白忙一场。也可以训练过程中手动把partial_checkpoint复制出来续训但我不建议新手这么做续训的路径和参数比较容易搞错我见过有人把partial_checkpoint当成final去推理结果输出全黑。5.5 推理结果全黑或全背景label IDs和JSON定义对不上现象测试集推理完成了结果图打开全黑np.unique输出只有0。原因测试图和训练图强度分布差太远比如训练数据是增强CT测试数据是平扫CT归一化之后目标区域的响应完全不同更常见的是labels定义和标注数据不一致比如标注文件里类别ID是1到5dataset.json里只定义了0和1两个类别模型从未见过那些标签。解决先跑验证集看看是否正常。如果验证集正常、测试集全黑优先怀疑数据分布差异把测试图像的强度分布和训练集对比一下看是不是范围差太多。如果验证集也全黑检查标注数据里到底有哪些值把labels字典补全。不要先怀疑模型这一步排查顺序能省很多时间。还要顺带检查一下是不是测试图像命名里的模态序号和训练集不一致_0000写成了_0001模型拿到的通道根本不是它认识的那个。6. 一个让nnU-Net结果更稳的习惯先读预测图再谈指标我见过不少同事拿到Dice 0.9就急着交付结果临床医生打开一看分割结果比原始图像小了一圈边缘完全对不上。Dice是整体重叠率对边界偏差不敏感但对医学分割来说边界才是最能看出问题的。所以我现在有一个习惯无论训练还是测试推理出来的预测图一定先叠加到原图上肉眼看一遍再去看指标。用SimpleITK或nibabel都能做这个叠加最快速的方法是先把预测转成numpy然后做伪彩色覆盖import numpy as np import nibabel as nib import matplotlib.pyplot as plt img nib.load(case_001_0000.nii.gz).get_fdata() pred nib.load(case_001_pred.nii.gz).get_fdata() plt.figure(figsize(12, 6)) plt.subplot(1, 2, 1) plt.imshow(img[:, :, img.shape[2] // 2], cmapgray) plt.subplot(1, 2, 2) plt.imshow(img[:, :, img.shape[2] // 2], cmapgray) plt.imshow(np.ma.masked_where(pred[:, :, img.shape[2] // 2] 0, pred[:, :, img.shape[2] // 2]), alpha0.5, cmapautumn) plt.savefig(overlay_check.png)这个脚本每次推理后我都会跑一遍挑几个典型切片存成图。如果要系统验证一个测试集我会从中随机抽10个病例做叠加图肉眼检查有没有分割区域跑到图像外面去、目标器官之间的边界粘连、小结构被漏掉。这三个问题指标上不一定看得出来但叠加图一眼就能发现。做nnU-Net这两年我最大的教训就是不要迷信框架的自动化自动化的plan和train只是把常规参数调好了数据本身的问题它管不了。每一次拿到新数据集先把数据质量关卡住再谈训练和推理。希望这个习惯对你也有用。本文还有配套的精品资源点击获取