人工智能深度学习计算机视觉医疗健康【免费下载链接】nnUNet项目地址https://gitcode.com/gh_mirrors/nn/nnUNet点击查看免费下载Primus 是一系列以 Transformer 为核心的 3D 医学图像分割架构旨在回答一个关键问题在不把大部分工作隐藏进 CNN 骨干的情况下Transformer 能否真正在体素级医学分割中具备竞争力本文基于仓库中的 documentation/primus.md 与其在 nnU-Net 框架内的实现primus_trainers.py系统讲解 Primus / PrimusV2 / PrimusV3 的设计动机、架构细节、各尺度配置、公开测试集结果以及如何在 nnU-Net 中使用对应 Trainer 完成训练与推理。背景为什么需要 Transformer-centric 的分割网络在医学影像领域许多自称 基于 Transformer 的分割网络实际上是 CNN–Transformer 混合体。分析表明这类架构在 Transformer 模块之外分配了大量容量其中不少模型在把 Transformer 块替换为恒等映射identity后性能几乎不下降——也就是说在很多流行的混合架构中CNN 组件基本可以独立解决任务注意力机制并没有真正发挥作用。Primus 的思路是让架构回归 Vision Transformer 范式几乎所有的参数与 FLOPs 都集中在 Transformer 骨干中使用高分辨率 3D token并搭配一个轻量解码器。PrimusV2 在此基础上加入更强的迭代式 tokenizeriterative tokenizer在提升视觉 token 质量的同时保持模型对 Transformer 的可度量依赖PrimusV3 则进一步针对 patch embedding 的信道瓶颈做了激进改进详见下文 Primus V3 小节。核心设计思想Primus 系列架构的关键设计可以归纳为六点Transformer-centric 设计最小化非 Transformer 容量确保表示学习主要由注意力机制驱动避免 CNN 骨干 抢走 Transformer 的职责。高分辨率 3D token与常见的16 × 16 × 16分块不同Primus 采用8 × 8 × 8的 token 化方式以保留更精细的局部解剖细节。3D 旋转位置编码3D axial RoPE在基于 patch 的 3D 训练下提升空间感知能力。现代 Transformer 块采用 SwiGLU MLP、LayerScale、post-attention 归一化、DropPath、AdamW 优化器以及强权重衰减weight decay等现代设计元素。nnU-Net 集成Primus 直接实现于 nnU-Net 框架内可以复用其预处理、规划planning、数据增强与训练基础设施。PrimusV2 tokenizerPrimusV2 添加了极简的残差迭代 tokenizer显著改善小目标与精细结构的分割性能同时保持 Transformer 在功能上的重要性。架构总览Primus 的处理流水线为将 3D 输入体素化tokenize为高分辨率视觉 token → 送入采用 3D axial RoPE 与现代 Transformer 块的骨干网络 → 将 token 序列解码回密集分割图。PrimusV2 沿用同一总体设计但用迭代残差 tokenizer 取代单一跨步投影strided projection以产生更强的视觉 token。在仓库实现中这一流水线由 nnunetv2/training/nnUNetTrainer/primus/primus_trainers.py 中的AbstractPrimus.build_network_architecture组装它从dynamic_network_architectures.architectures.primus导入Primus、PrimusV2{S,B,M,L}与PrimusV3{S,B,M,L}V3 系列以 try/except 方式导入若所装dynamic_network_architectures版本不支持则置为None。每个 Trainer 将网络实例化所需的超参数嵌入维度、层数、注意力头数、drop_path_rate、scale_attn_inner、init_values固化在类内而 patch size 则取自configuration_manager.patch_size。为什么选择 PrimusTransformer 天然以 token 序列为操作对象这使其非常适合 3D 医学影像中的可扩展自监督学习、掩码图像建模masked image modeling以及未来多模态系统——例如同时融合图像、文本、报告、测量值与其他临床数据的系统。但这些优势的前提是Transformer 必须真正学到高质量的 3D 视觉 token。Primus 提供了这样的骨干它弥合了基于 Transformer 的分割模型与强 CNN 基线之间的差距PrimusV2 更进一步在多个公开 3D 医学分割数据集上与 ResEnc-L、MedNeXt 等先进卷积架构达到同等水平具体数值见下节测试集结果。架构家族与各尺度配置Primus 提供多个尺度的变体适用于不同算力与任务规模模型LayersHeadsEmbedding dim.主要用途Primus-S126396轻量实验与小数据集Primus-B1212792均衡基线Primus-M1612864推荐默认配置Primus-L24161056大规模实验PrimusV2 遵循相同的缩放哲学但将单一跨步 patch 投影替换为迭代残差 tokenizer。从源码看各尺度的具体实例化参数与文档表格完全对应nnUNet_Primus_S_Trainer使用Primus(num_input_channels, 396, (8,8,8), num_output_channels, 12, 6, patch_size, ...)B/M/L 分别对应(792, 12, 12)、(864, 16, 12)、(1056, 24, 16)token 尺寸统一为(8, 8, 8)。所有 Trainer 的代码注释都明确标注this architecture will crash if the patch size is not divisible by 8!——即 patch size 必须能被 8 整除这是使用 Primus 系列时最重要的前置约束之一。仓库还提供若干实用变体nnUNet_Primus_M_Trainer_BS8/nnUNet_Primus_M_Trainer_BS8_2e4将 batch size 调整为 8后者另将initial_lr设为2e-4适用于显存更充足或需要更大批次稳定性的场景_Primus_S/B/M_96_BS1、_Primus_L_48_BS1直接在 plans 中将 patch size 强制设为(96,96,96)L 为(48,48,48)并将 batch size 设为 1对应论文原始实验设定。测试集结果下表汇总了 Primus、PrimusV2、PrimusV3 与所选基线在测试集上的 Dice 相似系数DSC结果。数值为对应交叉验证折上的平均 DSC数据集ACDC、AMOS22、KiTS23、LiTS、SST3、MAMA、SBM、Atlas22、WORD。方法Params. (M)ACDCAMOS22KiTS23LiTSSST3MAMASBMAtlas22WORD平均nnU-Net31.291.3488.6185.9979.2990.2778.3266.5263.1183.1180.73ResEnc-L102.492.5489.3988.0681.2090.2879.0064.0063.1285.7981.48MedNeXt-L61.892.5589.5888.2081.5789.9379.4265.8563.0385.3781.72CoTr41.990.5087.9384.6378.4489.6076.9559.9662.1483.1179.25Primus-M146.692.2688.1886.3879.5288.3176.3957.6360.1082.9879.08PrimusV2-M147.292.2789.3588.0981.7388.2679.4066.3663.2384.1581.43PrimusV3-S70.692.7089.2988.6581.2989.6680.6863.9564.1084.9981.70观察要点Primus-M146.6M 参数在参数规模显著大于 CNN 基线的前提下平均 DSC 为 79.08已超越 CoTr 但略低于 nnU-Net 基线——这正是文档所述 closes the gap 的体现PrimusV2-M147.2M平均 DSC 达 81.43与 ResEnc-L81.48、MedNeXt-L81.72基本持平实现了与强卷积架构的 parity尤其在 LiTS、SBM、Atlas22 上明显反超 Primus-M印证迭代 tokenizer 对细粒度结构的作用PrimusV3-S70.6M以不足一半的参数达到平均 81.70接近 MedNeXt-L 的水平。需要强调的是上述数值为文档与论文报告的结果实际复现时受数据预处理、硬件与训练细节影响应以本仓库代码在自有数据集上的 5 折交叉验证结果为准。在 nnU-Net 中使用 Primus TrainerPrimus 以标准 nnU-Net Trainer 的形式接入框架训练入口与默认 nnU-Net 完全一致仅需用-tr指定 Trainer 名称。可用 Trainer 清单原始 Primus立即 8×8×8 token 化nnUNet_Primus_S_Trainer nnUNet_Primus_B_Trainer nnUNet_Primus_M_Trainer nnUNet_Primus_L_TrainerPrimusV2迭代残差 tokenizer对应 TMLR 发表版本的最佳模型nnUNet_PrimusV2S_Trainer nnUNet_PrimusV2B_Trainer nnUNet_PrimusV2M_Trainer nnUNet_PrimusV2L_TrainerPrimusV3更重的卷积型 patch embedding初步版本nnUNet_PrimusV3S_Trainer # (推荐) nnUNet_PrimusV3B_Trainer nnUNet_PrimusV3M_Trainer nnUNet_PrimusV3L_Trainer官方当前推荐使用nnUNet_PrimusV3S_Trainer。训练示例需先按 nnU-Net 标准流程完成数据准备与nnUNetv2_plan_and_preprocessnnUNetv2_train DATASET_NAME_OR_ID 3d_fullres FOLD -tr nnUNet_PrimusV3S_Trainer推理同样使用标准命令Trainer 名称保持一致并配合-f all等参数加载各折模型nnUNetv2_predict -i INPUT_FOLDER -o OUTPUT_FOLDER \ -d DATASET_NAME_OR_ID -c 3d_fullres \ -tr nnUNet_PrimusV3S_Trainer训练超参数与优化策略源码级说明所有 Primus Trainer 统一继承自AbstractPrimus其父类为nnUNetTrainer_warmup关键的训练配置如下见 primus_trainers.pyself.initial_lr 3e-4self.weight_decay 5e-2强权重衰减与文档所述 strong weight decay 一致self.enable_deep_supervision False默认关闭深度监督这是与 nnU-Net 默认 Trainer 的重要差异两阶段优化器调度继承自 nnUNetTrainer_warmup。前warmup_duration_whole_net 50个 epoch 使用Lin_incr_LRScheduler对全网络做线性升温warmup随后切换到PolyLRScheduler_offset幂 0.9 的多项式衰减直至num_epochs 1000。on_train_epoch_start会在 epoch 0 与 epoch 50 自动切换两个阶段断点续训时load_checkpoint也会依据当前 epoch 恢复正确的调度器优化器torch.optim.AdamW(params, lr3e-4, weight_decay5e-2, amsgradFalse, betas(0.9, 0.98), fusedTrue)。Primus 系列使用 AdamW 而非 nnU-Net 默认的 SGD-NesterovfusedTrue在支持时启用融合内核以加速梯度裁剪自定义的train_step在反向后执行torch.nn.utils.clip_grad_norm_(self.network.parameters(), 1)并在 CUDA 上启用torch.autocast混合精度配合grad_scaler非 CUDA 设备则走dummy_context()set_deep_supervision_enabled为空实现pass进一步确认训练过程不使用深度监督。这些超参数与文档描述 modern Transformer blocksSwiGLU、LayerScale、post-attention norm、DropPath配合 AdamW 与强权重衰减 一一对应。各 Trainer 实例化网络时统一使用drop_path_rate0.2, scale_attn_innerTrue, init_values0.1。使用预训练权重微调如需从 nnSSL 预训练 checkpoint 微调 Primus 模型可参考 documentation/finetuning_from_nnssl_checkpoints.md对于 PrimusTransformercheckpoint训练命令使用-tr PretrainedTrainer_Primusx。其实现位于 pretrainedTrainer.pyPretrainedTrainer_Primus使用预训练权重时initial_lr 1e-4从零训练时自动切换为3e-4weight_decay 5e-2同样关闭深度监督warmup_duration_whole_net 50PretrainedTrainer_Primus_150ep将 warmup 缩至 15 个 epoch、总 epoch 数设为 150适合快速微调另外get_network_via_name.py 支持按名称实例化PrimusS / PrimusM / PrimusL / PrimusB架构同样需要安装相应版本的dynamic_network_architectures包供自定义实验调用。Primus V3修复 patch embedding 的 瓶颈 问题在 PrimusV2 投稿 TMLR 之后作者继续探索设计空间重点放在 patch embedding 上。UK-Essen 的 Luc Bouteille 指出PrimusV2 的迭代 patch embedding 在进入 Transformer 之前就形成了一个信息瓶颈。问题的根源在于信道增长与空间降采样之间的失衡。PrimusV2 的 stem 路径为32 (stem) → 32 (stride 2) → 64 (stride 4) → 128 (stride 8)沿此路径空间分辨率每个轴缩小 8 倍即体积缩小到原来的 1/8³ 1/512而信道只增长了 4 倍32 → 128。空间体积收缩×1/512远快于信道增长×4意味着表示在到达 Transformer 之前已被严重压缩。PrimusV3 的 patch embedding 正是针对这一点设计它把信道扩张大幅加快32 → 64 (stride 2) → 256 (stride 4) → 1024 (stride 8)使信道增长更好地抵消空间下采样带来的信息损失。此外patch embedding 现在带有一条直接跳跃连接direct skip connection每个空间分辨率都被投影到最终嵌入维度并加到 token 序列上从而生成由 8×8×8 核与迭代 token 尺寸共同构成的复合 tokencomposite tokens。下图直观对比了 PrimusV2-S 与 PrimusV3-S 的 stem 结构差异PrimusV2-S PrimusV3-S input 1 64³ input 1 64³ | | stem (k3,s1) stem (k3,s1) 32 64³ 32 64³ ───────────────┐ skip (k8,s8) | | | stage0 (s2) stage0 (s2) | 32 32³ 64 32³ ──────────┐ | skip (k4,s4) | | | | stage1 (s2) stage1 (s2) | | skip (k2,s2) 64 16³ 256 16³ ────┐ | | | | | | | stage2 (s2) stage2 (s2) | | | 128 8³ 1024 8³ | | | | | | | | final 1×1 final 1×1 | | | 396 8³ 396 8³ | | | | | | | | | () ◄───────┴────┴────┘ Σ scaleᵢ·proj_i | | (learnable scales, init ≈ 1e-5) tokens 396 8³ tokens 396 8³ | | EVA encoder (no spatial change) EVA encoder (no spatial change) 396 8³ 396 8³注意图中 PrimusV3 的复合 token 求和使用了可学习缩放系数learnable scales初始化 ≈ 1e-5即Σ scaleᵢ · proj_i让网络在训练初期近似于纯 8×8×8 token再逐步放大各分辨率跳跃连接的贡献。代价与验证这种设计把大量参数集中到了 patch embedding约 60MUNet Index ≈ 2。但即便如此仅靠卷积 stem 依然不足以取得好性能——因为缺少全局上下文移除 Transformer 后LiTS、ACDC、KiTS、AMOS 四数据集 5 折平均 DSC 从87.98 掉到 84.48。这从实验层面再次印证了 Primus 系列 Transformer-centric 的立场Transformer 承担了不可替代的角色而非可被 CNN 骨干替换的装饰件。使用建议与注意事项patch size 必须能被 8 整除所有 Primus/PrimusV2/PrimusV3 Trainer 在 patch 不可整除时都会直接崩溃源码注释明示规划时建议优先采用 64、96、128 等尺寸显存规划Primus 参数规模较大M 级约 146–147M 参数默认 batch size 可能较高可选用_BS8/_96_BS1等变体或自行调整 plans 中的batch_size深度监督默认关闭若自定义实验需要深度监督需自行覆写enable_deep_supervision与set_deep_supervision_enabled依赖要求Primus 网络定义位于外部包dynamic_network_architectures中使用前须确保安装包含architectures.primus模块的版本V3 系列的导入失败会被静默降级为None建议装齐后再训练结果复现测试集表格中的 DSC 以 5 折交叉验证平均报告建议按 documentation/how-to/find-best-configuration.md 的流程在自己的数据上做交叉验证后再下结论。延伸阅读架构说明原文documentation/primus.mdTrainer 实现全部 Primus/PrimusV2/V3 变体nnunetv2/training/nnUNetTrainer/primus/primus_trainers.pyWarmup PolyLR 两阶段调度nnunetv2/training/nnUNetTrainer/variants/lr_schedule/nnUNetTrainer_warmup.py、nnunetv2/training/lr_scheduler/warmup.pyPrimus 预训练微调PretrainedTrainer_Primusnnunetv2/training/nnUNetTrainer/pretraining/pretrainedTrainer.py使用说明见 documentation/finetuning_from_nnssl_checkpoints.md架构按名实例化入口nnunetv2/utilities/get_network_via_name.py训练/推理命令行总览documentation/how-to/train-models.md、documentation/how-to/run-inference.md赞分享人工智能深度学习计算机视觉医疗健康【免费下载链接】nnUNet项目地址https://gitcode.com/gh_mirrors/nn/nnUNet点击查看免费下载相关推荐3D-TransUNet医学图像分割Transformer与U-Net的完美融合实践3D TransUNet医学图像分割Transformer与U Net的完美融合实践 在当今医学影像分析领域3D TransUNet凭借其创新的架构设计正在MONAI与nnU-Net集成实战指南医学图像分割的强强联合MONAI与nnU Net集成实战指南医学图像分割的强强联合 前言 在医学影像分析领域nnU Net和MONAI都是备受推崇的开源框架。本文将深入探讨这两个示例工程教程PaddleSeg MedicalSeg 实战nnU-Net 医学图像分割全流程以 MSD Lung 数据集为例PaddleSeg MedicalSeg 实战nnU Net 医学图像分割全流程以 MSD Lung 数据集为例 导读 本文以 Medical Segme人工智能计算机视觉预训练上一篇CANN opbase 算子输入 Format 校验错误 EZ0006 Invalid_Input_Format 定位与排查指南下一篇nvtop用户反馈处理从issue到功能改进的流程创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考