尧图网络科技YAOTU DIGITAL 获取报价
获取报价
首页 / 资讯中心 / 文章详情

nnU-Net ResEnc L 在 ToothFairy2 挑战赛中的完整复现指南:数据集转换、超大 Patch 训练与限时推理优化

发布时间:2026/9/25 4:07:42

资讯中心
01
ARTICLE

nnU-Net ResEnc L 在 ToothFairy2 挑战赛中的完整复现指南:数据集转换、超大 Patch 训练与限时推理优化

nnU-Net ResEnc L 在 ToothFairy2 挑战赛中的完整复现指南:数据集转换、超大 Patch 训练与限时推理优化
人工智能深度学习计算机视觉医疗健康【免费下载链接】nnUNet项目地址https://gitcode.com/gh_mirrors/nn/nnUNet点击查看免费下载导读本文基于 nnU-Net 开源仓库中 DKFZ 团队对 Toothfairy2 挑战赛口腔 CBCT 全景牙/颌分割的官方提交方案文档完整还原其技术路线以 ResEnc L 残差编码器 U-Net 为骨干将 Patch 尺寸放大至 160×320×320、禁用左右镜像并延长训练至 1500 epoch同时讲解针对 grand-challenge 平台单例 10 分钟推理时限所做的规划器torch 重采样、双模型集成与体积阈值后处理等工程优化。读者将掌握从 mha 数据集转换、plans 文件手工扩写、定制 Trainer 训练到双模型集成推理与后处理的端到端复现能力并理解每个环节背后的源码级原理。一、方案概览为 Toothfairy2 挑战赛定制的 nnU-NetToothfairy2 挑战赛要求算法在口腔 CBCT 影像上分割牙齿与颌骨结构。本仓库 Toothfairy2 提交文档作者为 DKFZ 医学图像计算部门记录的官方方案其核心并非重新设计网络而是对 nnU-Net v2 的标准流程做四处关键改造放大 Patch将 ResEnc L 配置的 Patch 尺寸上采样到160×320×320体素让网络看到更大的解剖上下文加深网络架构比标准 ResEnc L 多一个 stage7 个 stage多一次池化与残差块调整训练策略禁用左右镜像left/right mirroring训练轮数从默认 1000 提升到 1500面向推理时限优化用 torch 实现的重采样替代默认重采样更快但略欠精确训练两个模型并以 cross-validation 集成策略合并以满足 grand-challenge 单例 10 分钟的时间上限。训练硬件为 2×A100 40GB 或单块 GH200 96GB模型从零开始训练不加载预训练权重。二、数据集转换mha 转 NIfTI 与标签重映射2.1 转换脚本与标签语义官方在 Dataset119_ToothFairy2_All.py 中完成数据集转换。该脚本的作用正如文档所述只是将 mha 文件转换为 nifti更小的文件体积并移除未使用的标签 id。从源码看脚本同时完成了三类工作格式转换image_to_nifi使用 SimpleITK 将*.mha图像直接写出为*.nii.gz标签重映射label_mapping通过查表把原始标签 id 映射为连续 id并压缩为uint8np.zeros_like(label_np, dtypenp.uint8)丢弃不存在的NA 类生成元数据重写dataset.jsonfile_ending改为.nii.gz、更新labels字典并额外生成 70:30 划分的splits_final.json随机种子固定为 42。2.2 标签映射表mapping_DS119转换脚本中mapping_DS119()定义了核心映射规则移除所有 NA 类并让类 id 连续def mapping_DS119() - Dict[int, int]: Remove all NA Classes and make Class IDs continuous mapping {} mapping.update({i: i for i in range(1, 19)}) # [1-10]-[1-10] | [11-18]-[11-18] mapping.update({i: i - 2 for i in range(21, 29)}) # [21-28]-[19-26] mapping.update({i: i - 4 for i in range(31, 39)}) # [31-38]-[27-34] mapping.update({i: i - 6 for i in range(41, 49)}) # [41-48]-[35-42] return mapping即原始标签 1–18 保持不变21–28 映射到 19–2631–38 映射到 27–3441–48 映射到 35–42最终得到 42 个前景类别标签 1–42。脚本中还保留了mapping_DS120仅保留牙与颌类与mapping_DS121仅保留牙类两种变体默认注释掉说明同源数据可衍生出多个 nnU-Net 数据集。转换入口需按实际路径修改rootroot /media/l727r/data/Teeth_Data/ToothFairy2_Dataset process_ds(root, Dataset112_ToothFairy2, Dataset119_ToothFairy2_All, mapping_DS119(), None)运行后将得到符合 nnU-Net v2 数据集规范的Dataset119_ToothFairy2_All含imagesTr/、labelsTr/、dataset.json。三、实验规划与预处理指纹提取、torch 重采样规划与 plans 文件扩写3.1 提取数据集指纹nnUNetv2_extract_fingerprint -d 119 -np 48该命令并行-np 48分析 Dataset 119 的体素间距、强度分布、形状统计等信息生成数据指纹是规划阶段的前提。3.2 使用 torch 重采样规划器进行规划nnUNetv2_plan_experiment -d 119 -pl nnUNetPlannerResEncL_torchres这里的核心是自定义规划器nnUNetPlannerResEncL_torchres定义于 resample_with_torch.py。它继承自nnUNetPlannerResEncL见 residual_encoder_unet_planners.py目标显存约 24GBUNet_reference_val_3d 2100000000并做了两点改动替换默认重采样方案determine_resampling()返回resample_torch_fornnunet实现在 resample_torch.py基于torch.nn.functional.interpolate而非默认的基于 scipy 的方案——更快但精度略低修改 plans 标识plans_name默认为nnUNetResEncUNetLPlans_torchresgenerate_data_identifier()返回plans_identifier _ configuration_name保证不同 plans 文件产生的预处理数据互不混淆。如文档所强调由于挑战赛训练/测试图像本就统一为0.3×0.3×0.3 mm间距理论上无需重采样改用快速方案纯粹是安全冗余真正的动机是推理端的速度——grand-challenge 平台对每个病例限时 10 分钟。3.3 编辑 plans 文件手工追加超大 Patch 配置规划完成后需在生成的 plans 文件nnUNetResEncUNetLPlans_torchres中追加以下配置块。除修改 Patch 尺寸外该配置还将网络加深了一个 stage多一次池化 残差块从而让网络更充分地利用更大输入3d_fullres_torchres_ps160x320x320_bs2: { inherits_from: 3d_fullres, data_identifier: nnUNetPlans_3d_fullres_torchres_ctnorm, patch_size: [ 160, 320, 320 ], normalization_schemes: [ CTNormalization ], architecture: { network_class_name: dynamic_network_architectures.architectures.unet.ResidualEncoderUNet, arch_kwargs: { n_stages: 7, features_per_stage: [ 32, 64, 128, 256, 320, 320, 320 ], conv_op: torch.nn.modules.conv.Conv3d, kernel_sizes: [ [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3] ], strides: [ [1, 1, 1], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [1, 2, 2] ], n_blocks_per_stage: [ 1, 3, 4, 6, 6, 6, 6 ], n_conv_per_stage_decoder: [ 1, 1, 1, 1, 1, 1 ], conv_bias: true, norm_op: torch.nn.modules.instancenorm.InstanceNorm3d, norm_op_kwargs: { eps: 1e-05, affine: true }, dropout_op: null, dropout_op_kwargs: null, nonlin: torch.nn.LeakyReLU, nonlin_kwargs: { inplace: true } }, _kw_requires_import: [ conv_op, norm_op, dropout_op, nonlin ] } }关键点逐一解读字段值含义inherits_from3d_fullres继承标准 3D 全分辨率配置的其余属性patch_size[160, 320, 320]相比 ResEnc L 默认约 128×128×128 量级大幅放大利用各向同性 0.3mm 间距数据n_stages7比标准 ResEnc L 的 6 个 stage 多一级下采样features_per_stage(32,64,128,256,320,320,320)末三级特征数封顶在 320strides末项[1, 2, 2]最后一个 stage 仅在 Y/Z 轴池化第 0 轴不池化适配 160 这一较小轴normalization_schemes[CTNormalization]与data_identifier中的ctnorm对应CT 强度归一化3.4 预处理nnUNetv2_preprocess -d 119 -c 3d_fullres_torchres_ps160x320x320_bs2 -plans_name nnUNetResEncUNetLPlans_torchres -np 48注意-plans_name必须与规划阶段一致nnUNetResEncUNetLPlans_torchres预处理产物按nnUNetPlans_3d_fullres_torchres_ctnorm数据标识落盘供训练时按配置名读取。四、训练双模型、1500 epoch 与仅 0/1 轴镜像4.1 训练命令官方在两个模型上训练全部训练病例allfoldnnUNetv2_train 119 3d_fullres_torchres_ps160x320x320_bs2 all -p nnUNetResEncUNetLPlans_torchres -tr nnUNetTrainer_onlyMirror01_1500ep nnUNet_results${nnUNet_results}_2 nnUNetv2_train 119 3d_fullres_torchres_ps160x320x320_bs2 all -p nnUNetResEncUNetLPlans_torchres -tr nnUNetTrainer_onlyMirror01_1500ep两行命令的唯一差别在于第二行通过环境变量覆盖nnUNet_results追加_2后缀从而在不互相覆盖结果的前提下把同一模型训练两遍——这是后续双模型集成的基础。模型均从零开始训练。4.2 定制 Trainer 的源码原理nnUNetTrainer_onlyMirror01_1500ep定义于 nnUNetTrainerNoMirroring.py继承链为nnUNetTrainer_onlyMirror01_1500ep → nnUNetTrainer_onlyMirror01 → nnUNetTrainerclass nnUNetTrainer_onlyMirror01(nnUNetTrainer): Only mirrors along spatial axes 0 and 1 for 3D and 0 for 2D def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes \ super().configure_rotation_dummyDA_mirroring_and_inital_patch_size() patch_size self.configuration_manager.patch_size dim len(patch_size) if dim 2: mirror_axes (0, ) else: mirror_axes (0, 1) self.inference_allowed_mirroring_axes mirror_axes return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes class nnUNetTrainer_onlyMirror01_1500ep(nnUNetTrainer_onlyMirror01): def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, device: torch.device torch.device(cuda)): super().__init__(plans, configuration, fold, dataset_json, device) self.num_epochs 1500镜像策略mirror_axes (0, 1)表示训练时只沿 3D 空间的前两个轴对应 sagittal/coronal 方向做镜像增强禁用了左右轴向axis 2镜像同时inference_allowed_mirroring_axes也被设为(0, 1)保证推理时测试时增强TTA的镜像轴与训练一致。这一点对牙齿这类左右高度对称的结构很有意义——避免网络被镜像引入的左右位置先验干扰。训练长度self.num_epochs 1500覆盖默认的 1000 epoch。同文件还提供了nnUNetTrainerNoMirroring完全禁用镜像与nnUNetTrainer_onlyMirror01_DA5、nnUNetTrainer_onlyMirror01_DASegOrd0等变体可在其他项目按需复用。4.3 数据增强的 CPU 瓶颈与nnUNet_n_proc_DA文档特别提醒建议提高数据增强进程数否则容易遭遇 CPU 瓶颈export nnUNet_n_proc_DA32系统允许时可继续调高。原因在于训练所用增强管线旋转、缩放、弹性形变、噪声、低分辨率模拟、镜像等实现在上述 Trainer 的get_training_transforms中均需 CPU 逐 batch 执行而 160×320×320 的超大 Patch 使单 batch 的增强计算量显著增大若 DA 进程数不足GPU 将被迫空转等待。五、推理双模型集成的限时优化5.1 集成策略把双模型伪装成交叉验证 foldgrand-challenge 平台对推理限时 10 分钟/病例因此官方没有采用跑两次推理再平均的朴素集成而是复用 nnU-Net 内置的交叉验证集成机制技术上把两个模型的fold_all目录复制到同一个训练输出目录下分别重命名为fold_0与fold_1从而启用 nnU-Net 的 cross-validation ensembling 策略计算效率更高满足平台时限。这样推理时两个模型共享同一份预处理后的滑窗高斯权重、同一批 patch 采样逐个 fold 加载权重并把 softmax 概率累加避免重复搬运大张量。5.2 推理脚本整体流程官方推理脚本为 inference_script_semseg_only_customInf2.py其main流程如下if __name__ __main__: os.environ[nnUNet_compile] f parser argparse.ArgumentParser() parser.add_argument(-i, --input_folder, typePath, default/input/images/cbct/) parser.add_argument(-o, --output_folder, typePath, default/output/images/oral-pharyngeal-segmentation/) parser.add_argument(-sem_mod, --semseg_trained_model, typestr, default/opt/app/_trained_model/semseg_trained_model) parser.add_argument(--semseg_folds, typestr, nargs, default[0, 1]) args parser.parse_args() args.output_folder.mkdir(exist_okTrue, parentsTrue) semseg_folds [i if i all else int(i) for i in args.semseg_folds] semseg_trained_model args.semseg_trained_model rw SimpleITKIO() input_files list(args.input_folder.glob(*.nii.gz)) list(args.input_folder.glob(*.mha)) for input_fname in input_files: output_fname args.output_folder / input_fname.name # load test image im, prop rw.read_images([input_fname]) with torch.no_grad(): semseg_pred predict_semseg(im, prop, semseg_trained_model, semseg_folds) torch.cuda.empty_cache() gc.collect() # now postprocess semseg_pred postprocess(semseg_pred, np.prod(prop[spacing]), True) semseg_pred map_labels_to_toothfairy(semseg_pred) # now save rw.write_seg(semseg_pred, output_fname, prop)默认路径已按 grand-challenge 的容器规范设定输入为/input/images/cbct/输出为/output/images/oral-pharyngeal-segmentation/脚本同时支持*.nii.gz与*.mha输入。单例流程为读取 → 集成预测 → 体积后处理 → 标签反向映射 → 写盘。5.3 CustomPredictor面向显存与速度的定制predict_semseg使用CustomPredictor(nnUNetPredictor)源码同文件内定义初始化参数为tile_step_size0.5, use_mirroringTrue, use_gaussianTrue。相比基类它重写了三个关键方法initialize_from_trained_model_folder从model_training_output_dir读取dataset.json、plans.json与checkpoint_final.pth用trainer_class.build_network_architecture(...)恢复网络enable_deep_supervisionFalse并支持nnUNet_compile环境变量触发torch.compilepredict_preprocessed_image这是最核心的提速/省显存手段——每个 fold 的权重按需加载self.network.load_state_dict(torch.load(p, ...)[network_weights])而不是一次性加载全部参数每个参数集在测试集上只用一次运行时间几乎不变但显存占用大幅下降滑窗预测在torch.autocast下以 float16 进行高斯权重compute_gaussian(tuple(patch_size), sigma_scale1./8, value_scaling_factor10)见 sliding_window_prediction.py同样以 float16 计算且数据与 logits 驻留 CPU、仅计算时搬运到cuda:0每个 patch 的预测先做pred / (pred.max() / 100)归一化再乘高斯权重累加保证不同 fold 贡献可比较convert_predicted_logits_to_segmentation_with_correct_shape用configuration_manager.resampling_fn_probabilities把 logits 重采样回裁剪前原始间距注意这里同样使用 torch 重采样argmax取类别、回填到裁剪 bbox最后按transpose_backward还原轴序。单线程相关细节预测与后处理中分别torch.set_num_threads(7)在受限容器环境中平衡 CPU 吞吐。5.4 标签反向映射由于训练标签是连续 id1–42提交结果必须映射回 ToothFairy2 官方标签体系。map_labels_to_toothfairy完成该逆映射mapping数组长度为 43仅对 19–42 重新映射def map_labels_to_toothfairy(predicted_seg: np.ndarray) - np.ndarray: max_label 42 mapping np.arange(max_label 1) remapping {19: 21, 20: 22, 21: 23, 22: 24, 23: 25, 24: 26, 25: 27, 26: 28, 27: 31, 28: 32, 29: 33, 30: 34, 31: 35, 32: 36, 33: 37, 34: 38, 35: 41, 36: 42, 37: 43, 38: 44, 39: 45, 40: 46, 41: 47, 42: 48} for k, v in remapping.items(): mapping[k] v return mapping[predicted_seg]这与数据集转换时的mapping_DS119恰好互逆确保提交结果符合挑战赛的标签定义。六、后处理基于五折交叉验证的体积截断文档明确后处理规则若某个类别的预测连通体积小于对应 cutoff则将其移除替换为背景。6.1 Cutoff 的确定方法cutoff 值在 Toothfairy2 训练数据上通过五折交叉验证优化得到分别针对HD95与Dice两个指标优化出两组 cutoff每个类别的最终 cutoff 取两组中的较小值更保守、宁可多删因此对假阳性小碎块有较强的抑制作用。6.2 推理脚本中的postprocess实现官方把最终体积 cutoff 直接写死在推理脚本的postprocess函数中按类 id 1–42 组织def postprocess(prediction_npy, vol_per_voxel, verbose: bool False): cutoffs {1: 0.0, 2: 78411.5, 3: 0.0, 4: 0.0, 5: 2800.0, 6: 1216.5, 7: 0.0, 8: 6222.0, 9: 1573.0, 10: 946.0, 11: 0.0, 12: 6783.5, 13: 9469.5, 14: 0.0, 15: 2260.0, 16: 3566.0, 17: 6321.0, 18: 4221.5, 19: 5829.0, 20: 0.0, 21: 0.0, 22: 468.0, 23: 1555.0, 24: 1291.5, 25: 2834.5, 26: 584.5, 27: 0.0, 28: 0.0, 29: 0.0, 30: 0.0, 31: 1935.5, 32: 0.0, 33: 0.0, 34: 6140.0, 35: 0.0, 36: 0.0, 37: 0.0, 38: 2710.0, 39: 0.0, 40: 0.0, 41: 0.0, 42: 970.0} vol_per_voxel_cutoffs 0.3 * 0.3 * 0.3 for c in cutoffs.keys(): co cutoffs[c] if co 0: mask prediction_npy c pred_vol np.sum(mask) * vol_per_voxel if 0 pred_vol (co * vol_per_voxel_cutoffs): prediction_npy[mask] 0 if verbose: print( fremoved label {c} because predicted volume of {pred_vol} is less than the cutoff {co * vol_per_voxel_cutoffs}) return prediction_npy实现要点传入的vol_per_voxel np.prod(prop[spacing])是该病例真实体素体积mm³而 cutoff 以0.3mm 间距下的体素数为单位比较前乘以vol_per_voxel_cutoffs 0.3×0.3×0.3换算为 mm³从而对实际间距与标称 0.3mm 的微小偏差保持稳健cutoff 为 0 的类别如 1、3、4、7、11 等直接跳过不做删除该后处理是类别级的按prediction_npy c统计整个类别的总体积而非连通域级因此实现极快几乎不占用推理时限预算。七、端到端复现路线图与关键文件索引把整条流水线汇总为可执行的步骤清单阶段命令/动作关键产物1. 数据转换运行 Dataset119_ToothFairy2_All.pyDataset119_ToothFairy2_AllNIfTI 连续标签2. 指纹提取nnUNetv2_extract_fingerprint -d 119 -np 48dataset_fingerprint3. 规划nnUNetv2_plan_experiment -d 119 -pl nnUNetPlannerResEncL_torchresnnUNetResEncUNetLPlans_torchres4. 编辑 plans追加3d_fullres_torchres_ps160x320x320_bs2配置见 3.3含超大 Patch 与 7-stage 架构5. 预处理nnUNetv2_preprocess -d 119 -c 3d_fullres_torchres_ps160x320x320_bs2 -plans_name nnUNetResEncUNetLPlans_torchres -np 48预处理后的训练数据6. 训练 ×2两条nnUNetv2_train ... -tr nnUNetTrainer_onlyMirror01_1500ep第二条改nnUNet_results两个fold_all模型7. 集成准备复制两个fold_all→fold_0/fold_1单目录双 fold 模型8. 推理 后处理运行 inference_script_semseg_only_customInf2.py提交用分割结果复现所需的源码锚点定制 TrainernnUNetTrainer_onlyMirror01与nnUNetTrainer_onlyMirror01_1500ep定义于 nnUNetTrainerNoMirroring.py镜像轴(0,1)、num_epochs1500torch 重采样规划器nnUNetPlannerResEncL_torchres定义于 resample_with_torch.py底层重采样实现见 resample_torch.pyResEnc L 家族规划器显存目标与参考值见 residual_encoder_unet_planners.py滑窗高斯权重compute_gaussiansigma_scale1/8、value_scaling_factor10见 sliding_window_prediction.py官方提交文档原文Toothfairy2/readme.md。八、结语一套标准方案 限时优化的范式Toothfairy2 的这套提交方案的价值不仅在于挑战赛成绩更在于它示范了 nnU-Net v2 如何在标准自动化流程与平台级工程约束之间取得平衡通过手工扩写 plans 文件把 ResEnc L 的 Patch 放大一个量级、增加一个 stage通过定制 Trainer 微调镜像与训练长度通过 torch 重采样与双模型集成把推理时间压进 10 分钟窗口最后用五折交叉验证训出的体积 cutoff 清洗假阳性。对任何需要在有限显存与推理时限内把 nnU-Net 推到极限的任务这套从 数据集转换脚本 到 推理脚本 的完整链路都是一份可直接迁移参考的实战范本。赞分享人工智能深度学习计算机视觉医疗健康【免费下载链接】nnUNet项目地址https://gitcode.com/gh_mirrors/nn/nnUNet点击查看免费下载相关推荐OpenCore Legacy Patcher完整教程4步让老Mac重获新生的终极指南OpenCore Legacy Patcher完整教程4步让老Mac重获新生的终极指南 还在为你的老Mac无法升级最新macOS而烦恼吗看着2012年的Ma人工智能深度学习计算机视觉医疗健康nnU-Net v2 入门指南从安装、数据准备到训练与推理的完整工作流nnU Net v2 入门指南从安装、数据准备到训练与推理的完整工作流 本指南是面向首次使用 nnU Net v2 用户的完整入门路线图它把官方文档中分散的人工智能深度学习计算机视觉医疗健康eldarion-ajax实战构建现代化CRUD应用的最佳实践eldarion ajax实战构建现代化CRUD应用的最佳实践 在当今Web开发领域构建响应式、用户友好的CRUD创建、读取、更新、删除应用已成为标配需人工智能深度学习计算机视觉医疗健康上一篇Hotkey Detective让Windows热键冲突无所遁形的智能诊断工具下一篇深度解析Inferact/Kimi-K3-DSpark基于vLLM的MLA原生草稿模型核心原理与架构设计创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
02
RELATED NEWS

相关资讯

更多网站建设与数字化升级内容

03
WHY YAOTU

想打造同款高转化官网?

懂行业、懂生意,从建站到增长一站式陪跑

◈

场景化定制

不做模板站,围绕你的业务场景量身设计,小众不撞款。

◐

营销型架构

以转化目标组织内容与路径,让官网真正带来询盘。

▲

全周期服务

设计、开发、运营、运维一体,上线只是开始。

免费获取你的建站方案

留下需求,专属顾问 24 小时内为你输出方案建议。