MMDetection 3.x多GPU训练实战从误区到高效配置如果你正在使用MMDetection 3.x进行目标检测任务可能会发现以前在2.x版本中习以为常的多GPU训练方式突然失效了。那种简单设置CUDA_VISIBLE_DEVICES就能让所有GPU运转的日子已经一去不复返。本文将带你深入理解MMDetection 3.x的多GPU训练机制避开常见陷阱并掌握最新的torchrun迁移方法。1. 为什么你的多GPU训练不工作了许多从MMDetection 2.x升级到3.x的用户都会遇到一个令人困惑的现象明明通过CUDA_VISIBLE_DEVICES0,1,2,3指定了多块GPU但实际运行时只有第一块GPU在工作。这不是bug而是框架设计理念的转变。在MMDetection 2.x时代框架内部封装了许多分布式训练的细节提供了--gpus和--gpu-ids这样的便捷参数。但在3.x版本中开发团队决定将这些控制权交还给PyTorch原生的分布式训练机制让用户能够更灵活地配置训练过程。关键区别MMDetection 2.x框架自动处理多GPU分配MMDetection 3.x需要显式使用PyTorch分布式训练工具# MMDetection 2.x的多GPU训练方式已过时 python tools/train.py configs/yolox/yolox_s_8xb8-300e_coco.py --gpus 42. 理解dist_train.sh的工作原理虽然直接使用CUDA_VISIBLE_DEVICES不再有效但MMDetection 3.x仍然保留了tools/dist_train.sh脚本来简化多GPU训练。让我们拆解这个脚本的核心逻辑#!/usr/bin/env bash CONFIG$1 GPUS$2 PYTHONPATH$(dirname $0)/..:$PYTHONPATH \ python -m torch.distributed.launch \ --nproc_per_node$GPUS \ $(dirname $0)/train.py \ $CONFIG \ --launcher pytorch ${:3}这个脚本实际上做了三件事设置Python路径确保能正确导入MMDetection模块使用torch.distributed.launch启动分布式训练将GPU数量和其他参数传递给训练脚本典型用法CUDA_VISIBLE_DEVICES0,1,2,3 ./tools/dist_train.sh configs/yolox/yolox_s_8xb8-300e_coco.py 43. 从torch.distributed.launch迁移到torchrunPyTorch官方已经宣布torch.distributed.launch将被弃用推荐使用更现代的torchrun命令。下面是迁移指南参数/特性torch.distributed.launchtorchrun启动方式python -m torch.distributed.launch直接使用torchrunGPU指定需要CUDA_VISIBLE_DEVICES内置支持自动重启不支持支持训练中断后恢复参数传递需要--nproc_per_node等显式参数更简洁的语法迁移示例# 旧方式将被弃用 CUDA_VISIBLE_DEVICES0,1,2,3 python -m torch.distributed.launch --nproc_per_node4 tools/train.py config.py --launcher pytorch # 新方式推荐 torchrun --nproc_per_node4 --nnodes1 --node_rank0 --master_addr127.0.0.1 --master_port29500 tools/train.py config.py --launcher pytorch关键参数说明--nproc_per_node每个节点使用的GPU数量--nnodes总节点数单机为1--node_rank当前节点序号主节点为0--master_addr和--master_port分布式训练的主节点地址和端口4. 多机多卡训练配置对于需要跨多台机器进行大规模训练的场景MMDetection 3.x同样支持。假设你有2台机器每台8块GPU主节点机器1torchrun --nnodes2 --node_rank0 --master_addr机器1IP --master_port29500 --nproc_per_node8 tools/train.py config.py --launcher pytorch从节点机器2torchrun --nnodes2 --node_rank1 --master_addr机器1IP --master_port29500 --nproc_per_node8 tools/train.py config.py --launcher pytorchSLURM集群配置 如果你在使用SLURM管理的HPC集群可以这样配置srun -p mm_dev --job-namemmdet_train --gresgpu:8 --ntasks16 --ntasks-per-node8 --cpus-per-task5 python tools/train.py config.py --launcherslurm5. 验证与调试技巧确保你的多GPU训练真正生效可以尝试以下方法nvidia-smi监控 在另一个终端运行watch -n 1 nvidia-smi应该看到所有指定GPU的显存和计算利用率都有变化。日志检查 训练开始时日志中应该显示类似信息***************************************** Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. *****************************************性能对比 记录单GPU和多GPU训练每个epoch的时间理想情况下4GPU应该有3-3.5倍的加速。常见问题排查如果只有一块GPU工作检查是否遗漏--launcher pytorch参数遇到端口冲突尝试更改--master_port29500-29599NCCL错误可以尝试设置NCCL_DEBUGINFO环境变量获取更多信息6. 高级配置与优化掌握了基本的多GPU训练后还可以进一步优化梯度累积 当GPU显存不足时可以使用梯度累积模拟更大的batch size# 在config文件中添加 optim_wrapper dict( typeOptimWrapper, optimizerdict(typeSGD, lr0.01, momentum0.9, weight_decay0.0001), clip_gradNone, accumulative_counts4 # 每4个step更新一次权重 )混合精度训练 启用FP16训练可以显著减少显存占用并提高速度# 在config文件中添加 fp16 dict(loss_scale512.)自定义分布式策略 通过修改dist_params可以调整分布式训练参数# 在config文件中添加 env_cfg dict( cudnn_benchmarkFalse, mp_cfgdict(mp_start_methodfork, opencv_num_threads0), dist_cfgdict(backendnccl, timeout1800) )在实际项目中我发现合理配置这些参数可以使8GPU训练的效率提升40%以上特别是对于大模型如Cascade R-CNN或Swin Transformer。