人工智能深度学习计算机视觉医疗健康【免费下载链接】nnUNet项目地址https://gitcode.com/gh_mirrors/nn/nnUNet点击查看免费下载本文是 nnU-Net v2 官方多 GPU 训练文档documentation/multi_gpu_training.md的深度实战解析。文章围绕如何在多卡环境中高效训练 nnU-Net v2这一核心问题依次讲解最推荐的并行 Fold 方案、基于-num_gpus的单节点 DDP、基于torchrun的多节点 DDP以及集群超时场景下的 SIGUSR1 优雅退出机制并结合 run_training.py、ddp.py 与 nnUNetTrainer.py 等源码给出底层实现证据。读完本文你将掌握三种多卡训练方案的选型依据、完整可运行命令、底层进程组初始化原理以及自定义 Trainer 在 DDP 环境下必须遵守的 rank 使用规范。先想清楚多卡最有效的用法是什么当手头有多张 GPU 时nnU-Net 官方文档给出的第一建议非常明确把不同的训练任务并行放到不同 GPU 上而不是做数据并行。原因是 nnU-Net 使用的网络相对较小数据并行DDP很难达到线性加速并行 Fold 反而是吞吐量最高的方案。真正值得引入 DDP 的场景只有两类自定义 nnU-Net 配置超过了单张 GPU 的显存容量单卡训练时间过长无法接受。所以请记住这条决策主线默认并行 Fold特殊需求才上 DDP。推荐方案每张 GPU 跑一个训练并行 Fold把 5 折交叉验证的各个 fold 直接分派到不同 GPU 上是官方推荐的做法。以 2d 配置、5 张 GPU 为例CUDA_VISIBLE_DEVICES0 nnUNetv2_train DATASET_NAME_OR_ID 2d 0 [--npz] # train on GPU 0 CUDA_VISIBLE_DEVICES1 nnUNetv2_train DATASET_NAME_OR_ID 2d 1 [--npz] # train on GPU 1 CUDA_VISIBLE_DEVICES2 nnUNetv2_train DATASET_NAME_OR_ID 2d 2 [--npz] # train on GPU 2 CUDA_VISIBLE_DEVICES3 nnUNetv2_train DATASET_NAME_OR_ID 2d 3 [--npz] # train on GPU 3 CUDA_VISIBLE_DEVICES4 nnUNetv2_train DATASET_NAME_OR_ID 2d 4 [--npz] # train on GPU 4 wait参数说明DATASET_NAME_OR_ID数据集名称如Dataset002_Heart或数字 ID2d配置名configuration也可以是3d_fullres、3d_lowres等0~45 折交叉验证的 fold 编号--npz可选在最终验证时额外保存 softmax 概率的 .npz 文件用于后续寻找最优 ensemble让命令在后台并行执行wait等待所有后台任务结束。也可以用同样的方式在同一张 GPU 上顺序执行多个训练去掉与wait即可。文档特别指出在 nnU-Net 旧版本中需要等第一个训练把 .npz 预处理数据解压成未压缩的 .npy 文件后才能启动其他训练现在预处理数据改用 blosc2 压缩格式存储不再需要等待所有训练可以同时启动。从源码层面看nnUNetv2_train命令的实际入口是 pyproject.toml 中定义的nnunetv2.run.run_training:run_training_entry。该入口函数run_training.py会解析命令行参数并构建TrainingRunOptions最终经由launch_training进入训练流程。当不传-num_gpus默认num_gpus1时launch_training走的就是普通单进程路径因此并行 Fold 本质上就是多个互不干扰的独立训练进程。Option 1-num_gpus单节点 DDP如果你需要在单台机器内用多张 GPU 同步训练真正的数据并行可以直接给nnUNetv2_train传-num_gpus无需torchrun也不需要修改启动方式。克隆安装的仓库版本和 pip 安装版本均支持# 示例在当前机器上用 2 张 GPU CUDA_VISIBLE_DEVICES0,1 nnUNetv2_train DATASET_NAME_OR_ID 2d 0 [--npz] -num_gpus 2使用-num_gpus的四个注意点GPU 选择如果节点上 GPU 多于实际使用数量用CUDA_VISIBLE_DEVICES限制可见范围Batch size 上限GPU 数量不能超过每个 minibatch 的样本数。若batch_size2-num_gpus 2就是上限整除性确保 batch size 能被 GPU 数整除否则硬件利用不充分扩展性除非你增大了模型、patch 或 batch size否则 DDP 可能比每 GPU 各跑一个 fold更慢。底层实现-num_gpus是如何拉起 DDP 的-num_gpus的完整调用链在 run_training.pylaunch_training在num_gpus 1且没有外部 launcher 时会自动设置MASTER_ADDRlocalhost并通过find_free_network_port()run_training.py寻找空闲端口作为MASTER_PORT然后用torch.multiprocessing.spawn启动num_gpus个 worker每个 worker 由run_intranode_ddp执行先模拟 torchrun 的四个环境变量RANK、WORLD_SIZE、LOCAL_RANK、LOCAL_WORLD_SIZE这样下游代码只需统一读取这四个变量不必关心任务是如何启动的随后调用init_ddp_process_group建立 NCCL 进程组init_ddp_process_grouprun_training.py在 torch 版本支持torch 2.3时向dist.init_process_group传入device_id让 NCCL 立即建立 communicator这样坏掉的互联拓扑会在启动时直接报错而不是在数分钟后以 hang 的形式暴露。DDP 下的 batch size 分配与前景过采样进入 Trainer 后DDP 模式下的 batch size 分配逻辑位于_set_batch_size_and_oversamplennUNetTrainer.py全局 batch size 来自 plans 配置configuration_manager.batch_size并断言global_batch_size world_size否则直接报错Cannot run DDP if the batch size is smaller than the number of GPUs全局 batch size 被尽可能均匀地拆分到各 rank余数部分逐个 rank 加 1保证sum(batch_size_per_GPU) global_batch_size前景过采样比例oversample_percent也会按每个 rank 拿到的样本区间重新计算以复现 DataLoader 的取整语义。也就是说你不需要手动把 plans 里的 batch size 除以 GPU 数——Trainer 会自动按 rank 切分。Option 2torchrun多节点 / 进阶 DDPtorchrun是 PyTorch 标准的 DDP launcher同时支持单节点与多节点。使用torchrun时不要传-num_gpus——进程数和 GPU 分配由 launcher 自己决定源码中的launch_training对此有硬性校验见下文。前提条件使用torchrun必须克隆并安装 nnU-Net 仓库因为命令直接调用仓库内的脚本nnunetv2/run/run_training.pypip 安装方式没有该文件路径。单节点、4 张 GPUtorchrun --nnodes1 --nproc_per_node4 --node_rank0 --master_addrlocalhost --master_port12345 NNUNET_REPO_LOCATION/nnunetv2/run/run_training.py DATASETID CONFIGURATION FOLD [--npz]两个节点、每节点 4 张 GPU在所有节点上执行把MASTER_ADDR替换为 rank-0 节点的主机名或 IPtorchrun --nnodes2 --nproc_per_node4 --rdzv_idYOUR_CUSTOM_INTEGER --rdzv_backendc10d --rdzv_endpointMASTER_ADDR:MASTER_PORT NNUNET_REPO_LOCATION/nnunetv2/run/run_training.py DATASETID CONFIGURATION FOLD [--npz]torchrun的注意点禁用-num_gpusnnU-Net 会检测 launcher 并抛出AssertionError详见下文也可以用torchrun跑单张 GPU例如调试完整的 DDP 链路只要 torchrun 环境变量存在runner 就会进入 DDP 路径在 SLURM 集群上torchrun与srun/sbatch及多节点作业集成良好。底层实现nnU-Net 如何识别外部 launcherlaunch_training的核心分支逻辑run_training.py依赖launched_by_external_launcher()run_training.py只要环境变量中存在TORCHELASTIC_RUN_IDtorchrun 专属标记就判定为外部 launcher或者同时存在完整的RANK WORLD_SIZE LOCAL_RANK三件套兼容 srun、deepspeed 等遵循同一环境变量协议的 launcher之所以要求三个变量齐备是为了防止 shell 里残留的过期RANK把-num_gpus 4悄悄变成单进程任务。在外部 launcher 分支中断言num_gpus 1否则抛出AssertionError提示launcher 决定进程数和 GPU 分配不要同时传-num_gpus从LOCAL_RANK推导当前进程的 CUDA 设备用env://方式加入既有进程组打印 rank/world_size 信息训练结束后统一dist.destroy_process_group()。此外进程拓扑信息由get_ddp_topology()提供ddp.py。它优先读取 torchrun/-num_gpus写入的四个环境变量对第三方 launcherMPI 等则回退到SLURM_LOCALID、OMPI_COMM_WORLD_LOCAL_RANK等调度器变量若仍未设置还会通过一次跨 rank 的 hostname 集体通信all_gather_object推导 local_rank/local_world_size。三种方案如何选需求推荐方案标准 nnU-Net 配置追求最高吞吐每 GPU 一个训练并行 Fold单节点内多卡同步训练-num_gpus多节点 DDP或偏好 PyTorch 标准 launchertorchrunTroubleshooting 速查使用torchrun时务必省略-num_gpus只使用部分 GPU 时正确设置CUDA_VISIBLE_DEVICESbatch size 必须能被 GPU 数整除否则硬件利用不充分或直接报错不要用-device指定 GPU 编号入口参数的帮助信息明确说明-device只接受cuda/cpu/mpsGPU 选择统一交给CUDA_VISIBLE_DEVICESrun_training.py。集群超时SIGUSR1 优雅退出与断点续训nnU-Net 安装了一个SIGUSR1处理器。收到信号后不会打断当前 epoch只是置一个标志位在下一次 epoch 边界nnU-Net 会写入checkpoint_latest.pth并以状态码 0 退出。之后用--c重新提交即可从断点精确续训。Windows 上没有SIGUSR1不安装该处理器。源码级机制处理器注册在 Trainer 初始化阶段且仅当hasattr(signal, SIGUSR1)时执行nnUNetTrainer.pyexit_trainingnnUNetTrainer.py只做一件事设置exit_training_flag。注释解释了原因——信号处理函数可能在任意时刻中断主线程若此时持有日志/打印/文件 I/O 的锁就会死锁on_epoch_endnnUNetTrainer.py在每个 epoch 结束时调用_exit_signal_received()检查标志若收到信号立即保存checkpoint_latest.pth打印日志后sys.exit(0)非 DDP 模式下直接返回本地标志DDP 模式下_exit_signal_receivednnUNetTrainer.py将标志转成 int32 张量做一次all_reduce(MAX)只要任何一个 rank 收到信号所有 rank 一起退出。否则单个 rank 自行退出后其余 rank 会阻塞在 collective 上直到 NCCL watchdog 超时。不同调度器如何发送信号因为到底谁收到信号在不同站点和 launcher 下差异很大文档给出了针对性的配置SLURM#SBATCH --signalB:USR1600在超时前 10 分钟向批处理脚本发送SIGUSR1--signalUSR1600不带B:则发给作业步骤任务。手动发送用scancel --signalUSR1 --full JOBIDLSFbsub -wa USR1 -wt 10请求一个 warning 信号手动发送用bkill -s USR1 JOBIDtorchrun它不会把SIGUSR1转发给 worker——其 agent 只处理TORCHELASTIC_SIGNALS_TO_HANDLE中列出的信号默认SIGTERM,SIGINT,SIGHUP,SIGQUIT而且这些信号会直接杀掉 worker 而非让它们优雅退出。因此必须给workertorchrun 的子进程发信号而不是给 launchertorchrun ... nnunetv2/run/run_training.py DATASETID CONFIGURATION FOLD TORCHRUN_PID$! trap pkill -USR1 -P $TORCHRUN_PID USR1 # -P: only the children, never torchrun itself wait $TORCHRUN_PID特别注意不要用pkill -USR1 -f nnunetv2/run/run_training.py——launcher 自己的命令行里也包含该脚本路径这样会连 torchrun 一起信号到而 torchrun 无法处理SIGUSR1会直接死亡连带 worker 在 checkpoint 完成前一起被杀。由于一个 rank 收到信号即可在任意单个节点上对 worker 发信号就足够了。自定义 Trainer 的 DDP 规范在 DDP 环境中Trainer 的四个拓扑属性由get_ddp_topology()填充nnUNetTrainer.pyself.global_rank进程在整个作业跨所有节点中的索引self.local_rank进程在其所在节点内的索引同时就是 CUDA 设备编号self.device torch.device(typecuda, indexself.local_rank)self.world_size/self.local_world_size对应的进程总数。任何写入nnUNet_results的操作都必须用self.global_rank 0保护而不是self.local_rank 0——多节点作业中每个节点都有一个local_rank 0它们会写同一份文件导致冲突。这一点与save_checkpoint的实现一致nnUNetTrainer.py 中if self.global_rank 0:才执行保存验证集划分val_keys val_keys[self.global_rank::self.world_size]等逻辑同样遵循该约定。另外如果你的自定义 loss 需要在 DDP 下跨 rank 聚合nnU-Net 提供了可微的 SUM all-reduce 原语AllReduceGradddp.py前向把所有 rank 的张量求和反向把梯度求和等价于旧的AllGatherGrad.apply(x).sum(0)但数据移动量减少 world_size 倍。旧的AllGatherGradddp_allgather.py已标记为 deprecated仅因外部自定义 trainer/loss 仍引用而保留。与预训练权重结合-num_gpus机制同样适用于预训练微调入口nnUNetv2_train_pretrained其参数解析同样暴露了-num_gpus内部复用了launch_trainingrun_training_from_pretrained.py因此单节点多卡微调预训练权重与普通训练的使用方式完全一致。赞分享人工智能深度学习计算机视觉医疗健康【免费下载链接】nnUNet项目地址https://gitcode.com/gh_mirrors/nn/nnUNet点击查看免费下载相关推荐SpeechBrain 多 GPU 训练实战从 DDP 单节点到多节点集群部署SpeechBrain 多 GPU 训练实战从 DDP 单节点到多节点集群部署 本篇技术指南基于 SpeechBrain 官方文档 docs/multigpu人工智能深度学习语音音频NLP预训练基于 PyTorch 的 DDP 分布式训练实战用 torchrun 编排多节点多卡应用基于 PyTorch 的 DDP 分布式训练实战用 torchrun 编排多节点多卡应用 本篇技术指南围绕当前仓库 distributed/ddp 示例系统示例工程人工智能深度学习DGL GraphBolt 多 GPU 训练实战用 DistributedItemSampler 与 DDP 训练 GraphSAGE 节点分类模型DGL GraphBolt 多 GPU 训练实战用 DistributedItemSampler 与 DDP 训练 GraphSAGE 节点分类模型 本指南围人工智能机器学习深度学习图计算上一篇如何用M9A游戏助手彻底解放你的游戏时间重返未来1999自动化全攻略下一篇3分钟拯救你的桌面布局PersistentWindows彻底解决Windows多显示器窗口混乱问题创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考