nnU-Net 跨数据集监督预训练与微调实战：plans 对齐、权重迁移与微调流程详解

人工智能深度学习计算机视觉医疗健康【免费下载链接】nnUNet项目地址https://gitcode.com/gh_mirrors/nn/nnUNet点击查看免费下载导读在 nnU-Net v2 中不同数据集经过自动化实验规划后其 patch size、网络拓扑、batch size 乃至归一化方案往往各不相同因此不能直接把一个数据集上训练好的权重搬到另一个数据集上。本文以 documentation/pretraining_and_finetuning.md 为核心完整讲解监督预训练的可行路径先把目标数据集的 plans 转移到预训练数据集以对齐网络结构再用allfold 在预训练数据上训练最后通过-pretrained_weights参数把权重初始化到目标数据集并微调。读完本文你将掌握一套可复现的跨数据集迁移训练方案并了解其背后的源码实现原理与注意事项。需要说明的是本文讨论的是监督式的 plans/权重跨数据集迁移。如果你要找的是基于 nnssl本文不涉及该路径。术语约定整个流程涉及两个数据集原文档给出的术语定义如下术语含义pretraining dataset预训练数据集即你打算在其上执行预训练的数据集finetuning dataset微调数据集即你真正关心、最终要微调出模型的目标数据集两者的关系是先在预训练数据集上训练出一个通用的模型再把它的权重作为目标数据集的初始化从而获得更好的收敛性能或更短的训练时间。为什么不能直接复用权重plans 对齐是前提nnU-Net 最核心的特性之一是自动化数据集分析与实验规划对每个数据集它会依据数据指纹fingerprint自动决定 patch size、batch size、网络拓扑、归一化方案等训练超参数参见 nnunetv2/experiment_planning/ 下的相关实现。这意味着两个数据集的 patch size 可能不同两个数据集的网络拓扑深度、宽度、卷积核配置可能不同两个数据集的归一化方案可能不同。权重 tensor 的形状与网络结构严格绑定拓扑不一致时state_dict根本无法对齐。因此开箱即用地从一个数据集拿权重去初始化另一个数据集是不可能的。解决办法就是把目标数据集finetuning dataset的 plans 转移到预训练数据集让两边共享完全相同的网络结构、patch size、batch size 与归一化设置。原文档特别强调转移是全量的不只是网络拓扑、batch size 和 patch size归一化方案也会一并转移。因此在两个使用不同归一化方案的数据集之间做迁移可能效果不好但取决于具体方案也可能可行。对于 CT 数据还要注意CT 归一化中的 clip 值、均值mean与标准差std同样会被转移——预训练数据集的 CT 数据将按目标数据集的统计量做归一化处理。完整实操流程在预训练数据集上训练整个流程共分 5 步命令均来自原文档下面结合源码逐一展开说明。第 0 步环境准备按照 documentation/setting_up_paths.md 和 documentation/set_environment_variables.md 设置nnUNet_raw、nnUNet_preprocessed、nnUNet_results等环境变量预训练数据集与微调数据集都已按 nnU-Net 数据集格式放置于nnUNet_raw中格式规范见 documentation/dataset_format.md。第 1 步为微调数据集执行实验规划与预处理nnUNetv2_plan_and_preprocess -d FINETUNING_DATASET因为我们最终关心的是微调数据集所以需要先对它的数据进行数据指纹提取、实验规划与预处理产出它的 plans 文件。该命令的入口与实现位于 nnunetv2/experiment_planning/plan_and_preprocess_entrypoints.py核心 API 在 nnunetv2/experiment_planning/plan_and_preprocess_api.py 中。规划完成后微调数据集在nnUNet_preprocessed目录下会生成形如nnUNetPlans.json的 plans 文件。第 2 步提取预训练数据集的指纹若尚未生成nnUNetv2_extract_fingerprint -d PRETRAINING_DATASET该步骤只计算预训练数据集的指纹模态、spacing、强度分布等统计信息不做规划。从源码看它调用的是 nnunetv2/experiment_planning/plan_and_preprocess_api.py 中的extract_fingerprint_dataset实际计算由 nnunetv2/experiment_planning/dataset_fingerprint/fingerprint_extractor.py 中的DatasetFingerprintExtractor完成。指纹是后续预处理如裁剪、重采样所必需的输入。第 3 步将微调数据集的 plans 转移到预训练数据集nnUNetv2_move_plans_between_datasets -s FINETUNING_DATASET -t PRETRAINING_DATASET -sp FINETUNING_PLANS_IDENTIFIER -tp PRETRAINING_PLANS_IDENTIFIER参数含义参数说明-s源数据集此处为微调数据集的名称或 ID-t目标数据集此处为预训练数据集的名称或 ID-sp源 plans 标识符即源数据集 plans 文件名去掉.json后缀的部分通常就是nnUNetPlans除非你在nnUNetv2_plan_and_preprocess时更换过实验规划器-tp目标 plans 标识符建议设置成一个自定义名称如nnUNetPlansFrom2以免覆盖预训练数据集默认的 plans该命令的实际实现在 nnunetv2/experiment_planning/plans_for_pretraining/move_plans_between_datasets.py 中入口为entry_point_move_plans_between_datasets逻辑核心是move_plans_between_datasets见 L12-L61。从源码可以看到它在转移时做了以下几件事断言前置条件要求源数据集已存在预处理目录否则提示先运行nnUNetv2_plan_and_preprocess且源 plans 文件存在否则提示先执行实验规划见 L23-L29改写dataset_name把 plans 中的dataset_name改成目标预训练数据集名称L32保证后续训练/推理能正确找到目标数据集改写data_identifier若-tp与-sp不同会把每个 configuration 里的data_identifier前缀替换为目标标识符L35-L43这样预处理产出的数据会落在以新标识符命名的子目录下避免与预训练数据集原有预处理数据冲突改写 reader/writer根据目标数据集的dataset.json与示例图像重新确定image_reader_writerL45-L56保证读取器与目标数据集的图像格式匹配写入新 plans 文件最终保存为nnUNet_preprocessed/预训练数据集/PRETRAINING_PLANS_IDENTIFIER.jsonL60-L61。第 4 步用转移后的 plans 预处理预训练数据集nnUNetv2_preprocess -d PRETRAINING_DATASET -plans_name PRETRAINING_PLANS_IDENTIFIER注意这里必须显式指定-plans_name PRETRAINING_PLANS_IDENTIFIER即第 3 步中你自定义的目标 plans 标识符。由于该 plans 中记录的是微调数据集的 patch size、spacing 重采样目标和归一化参数预训练数据集的预处理数据将完全按微调数据集的口径生成从而保证后续训练时两边的数据形态一致。第 5 步在预训练数据集上训练使用 all foldnnUNetv2_train PRETRAINING_DATASET CONFIG all -p PRETRAINING_PLANS_IDENTIFIERPRETRAINING_DATASET预训练数据集名称或 IDCONFIG配置名如2d、3d_fullres等在转移的 plans 中存在的 configurationall使用全部数据训练。原文档明确指出预训练场景下没有必要划分数据做交叉验证所以这里用allfold 而非0/1/2/3/4之一-p PRETRAINING_PLANS_IDENTIFIER指定使用转移过来的自定义 plans。训练完成后或当你以其他方式获得了兼容权重后预训练数据集的 checkpoint如checkpoint_final.pth即作为可复用的初始化权重。使用预训练权重初始化目标数据集模型基本命令nnUNetv2_train FINETUNING_DATASET CONFIG FOLD -pretrained_weights PATH_TO_CHECKPOINT其中PATH_TO_CHECKPOINT指向预训练阶段产出的 checkpoint 文件例如.../checkpoint_final.pth。训练正常结束后你得到的就是在目标数据集上微调好的模型可用标准的 documentation/run_inference_with_pretrained_models.md 流程进行推理。源码层面的实现细节-pretrained_weights的处理逻辑位于 nnunetv2/run/run_training.py与-c继续训练互斥maybe_load_checkpoint中明确断言——continue_training与pretrained_weights不能同时使用预训练权重只能在训练开始时加载见 L100-L104加载时机当既不是继续训练也不是纯验证时且指定了pretrained_weights_file会先初始化 trainerinitialize()再调用load_pretrained_weights(nnunet_trainer.network, pretrained_weights_file, verboseTrue)将权重灌入网络见 L120-L125继续训练时的 checkpoint 查找顺序checkpoint_final.pth→checkpoint_latest.pth→checkpoint_best.pth见 L106-L111这与微调后继续训练的场景相关。权重加载函数位于 nnunetv2/run/load_pretrained_weights.py其机制值得展开按 key 与 shape 匹配只把 key 相同且 shape 相同的参数从预训练state_dict复制到当前网络L62-L65并在加载前逐一断言非分割层参数必须存在且形状一致L43-L51跳过分割层key 中包含.seg_layers.的参数即最终输出分割图的 1x1(x1) 卷积层不会被转移L30-L32。这正是原文档所说加载预训练权重时除分割层外的所有层都会被使用的源码依据——分割头保持随机初始化以适应目标数据集的类别数兼容 DDP 与 torch.compile 前缀如果预训练权重来自 nnU-Net 之外的训练如 DDP 会加module.前缀、torch.compile 会加_orig_mod前缀需要先修正 key而 nnU-Net 自己保存的 checkpoint 由nnUNetTrainer.save_checkpoint处理过无需担心L15-L18加载时也会自动剥离这些前缀L34-L40。微调阶段的训练器选择与自定义原文档明确说明目前 nnU-Net 没有专门针对微调场景定制的 trainer当前推荐做法就是直接使用标准nnUNetTrainer并配合-pretrained_weights参数。同时你完全可以根据需要编写自己的 trainer 子类例如学习率 warmup/ramp up在微调初期使用较低或线性上升的学习率避免破坏预训练特征分割头专项微调先冻结编码器只训练分割头再解冻全网络缩短训练轮数因为模型已有良好的初始化可以显著减少 epochs。关于自定义 trainer仓库中已有大量可参考的变体实现如 nnunetv2/training/nnUNetTrainer/variants/ 下的各类 Trainer自定义 trainer 需要继承nnUNetTrainer并置于nnunetv2.training.nnUNetTrainer命名空间下之后通过nnUNetv2_train的-tr参数指定该查找逻辑可参考 nnunetv2/utilities/find_class_by_name.py 的recursive_find_python_class。需要特别区分的是本仓库 nnunetv2/training/nnUNetTrainer/pretraining/pretrainedTrainer.py 中的PretrainedTrainer以及PretrainedTrainer_Primus、PretrainedTrainer_150ep等系列是自监督 nnssl 扩展路径专用的训练器——它从 plans 中的pretrain_info读取 checkpoint 路径与 key 映射负责把 nnssl 预训练编码器适配进分割网络包括按输入通道重复 stem 权重、按 patch size 对可学习位置编码做三线性插值等并配合 warmup→train 两阶段学习率调度。该训练器通过nnUNetv2_train_pretrained入口nnunetv2/run/run_training_from_pretrained.py启动与本文讨论的监督式 plans 转移 -pretrained_weights初始化是两条相互独立的流程详见 documentation/finetuning_from_nnssl_checkpoints.md。实操注意事项汇总结合原文档与源码以下几点在实践中尤其重要plans 必须严格对齐预训练与微调必须使用同一份结构相同的plans否则load_pretrained_weights中非分割层 key 与 shape 必须匹配的断言会直接报错-tp使用自定义标识符防止覆盖预训练数据集原本的nnUNetPlans同时后续nnUNetv2_preprocess与nnUNetv2_train都要记得带-plans_name/-p指向同一标识符归一化方案会整体转移包括 CT 的 clip 值、mean、std。若预训练数据集与微调数据集的归一化差异过大迁移效果可能不佳分割层不迁移目标数据集的类别数不必与预训练数据集一致分割头会从头训练预训练用allfold预训练阶段不划分交叉验证 fold把全部数据用于训练-pretrained_weights与-c互斥无法一边继续训练一边加载预训练权重预训练权重只在训练初始阶段生效checkpoint 路径要写对PATH_TO_CHECKPOINT应指向真实存在的.pth文件如checkpoint_final.pth源码会直接torch.load读取其network_weights字段。总结nnU-Net v2 的监督预训练方案本质上是一个plans 搬运 权重初始化的流程通过nnUNetv2_move_plans_between_datasets把微调数据集的完整训练配置网络拓扑、patch size、batch size、归一化参数等迁移给预训练数据集在预训练数据集上用allfold 完成预训练最后用-pretrained_weights把预训练权重初始化到微调数据集分割层除外并正常训练。整套流程的所有关键环节都有对应的源码实现可供核对任何一步出错如 plans 未对齐、checkpoint 路径不存在、归一化不匹配都能在报错信息或代码断言中找到明确提示。赞分享人工智能深度学习计算机视觉医疗健康【免费下载链接】nnUNet项目地址https://gitcode.com/gh_mirrors/nn/nnUNet点击查看免费下载相关推荐SimCLR 对比学习实战TensorFlow Model Garden 中的无监督预训练与半监督微调全流程SimCLR 对比学习实战TensorFlow Model Garden 中的无监督预训练与半监督微调全流程 TensorFlow Model Garden人工智能深度学习计算机视觉NLP语音Fairseq 中 wav2vec 2.0 全流程实战数据清单、自监督预训练、CTC 微调与跨语言扩展Fairseq 中 wav2vec 2.0 全流程实战数据清单、自监督预训练、CTC 微调与跨语言扩展 本文以 examples/wav2vec/README人工智能深度学习预训练NLP语音Ant Design DatePicker 的 needConfirm掌控选择确认按钮的提交时机Ant Design DatePicker 的 needConfirm掌控选择确认按钮的提交时机 本篇文章聚焦 Ant Design DatePicker 组大模型模型微调上一篇Komodo用户行为分析优化部署工具的使用体验下一篇Satellizer逆向工程从编译代码还原认证流程设计创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考