MMDetection 3.x 分布式训练实战：从单机多卡到多机多卡

weixin_30634661

394人浏览 · 2026-04-21 08:07:22

weixin_30634661 · 2026-04-21 08:07:22 发布

1. 为什么需要分布式训练

当你用MMDetection训练目标检测模型时，可能会遇到这样的困扰：单张GPU训练速度太慢，数据集稍微大点就要等好几天。这时候就该分布式训练上场了。简单来说，分布式训练就是让多块GPU甚至多台机器的GPU协同工作，把训练任务拆分开来同时处理。

我去年在训练一个改进版YOLOX模型时就深有体会。用单卡训练COCO数据集需要近30小时，换成4卡分布式训练后直接缩短到8小时。这就像搬砖，一个人干活和四个人分工协作，效率完全不在一个量级。

分布式训练主要解决三个问题：

训练加速：数据并行处理，batch size可以翻倍
大模型支持：模型参数可以分散到不同设备
资源利用率：充分利用服务器上的多块GPU

2. 单机多卡配置实战

2.1 环境准备要点

在开始之前，确保你的环境满足这些条件：

至少两块NVIDIA GPU（建议同型号）
正确安装CUDA和cuDNN
PyTorch版本与CUDA匹配
MMDetection 3.x源码完整

我推荐用conda创建独立环境：

conda create -n mmdet python=3.8 -y
conda activate mmdet
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html
pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu113/torch1.12/index.html

2.2 启动命令演变史

MMDetection 3.x相比2.x有个重要变化：移除了--gpus参数。刚开始我也踩了坑，用老方法CUDA_VISIBLE_DEVICES=0,1,2,3发现只有第一块GPU在工作。

正确做法是用dist_train.sh脚本：

CUDA_VISIBLE_DEVICES=0,1,2,3 ./tools/dist_train.sh configs/yolox/yolox_s_8xb8-300e_coco.py 4

这个脚本背后其实是调用了PyTorch的分布式启动器。我们也可以直接写等效命令：

CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch \
    --nproc_per_node=4 \
    tools/train.py \
    configs/yolox/yolox_s_8xb8-300e_coco.py \
    --launcher pytorch

不过要注意，PyTorch官方已经宣布torch.distributed.launch将被弃用，建议改用torchrun：

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun \
    --nproc_per_node=4 \
    tools/train.py \
    configs/yolox/yolox_s_8xb8-300e_coco.py \
    --launcher pytorch

2.3 常见问题排查

第一次跑分布式训练时，我遇到了几个典型问题：

GPU显存不均：某块卡显存爆满而其他卡闲置
- 解决方法：检查数据加载是否均衡，适当调整persistent_workers参数
进程挂起：训练卡在初始化阶段
- 通常是因为端口冲突，可以换--master_port值试试
性能不升反降：4卡比单卡还慢
- 可能是通信开销太大，可以尝试增大batch_size或使用SyncBN

3. 多机多卡进阶配置

3.1 基础网络配置

当单机GPU不够用时，就需要跨机器协作。假设有两台机器，每台8卡：

机器A（主节点）：IP 192.168.1.100
机器B：IP 192.168.1.101

两台机器需要：

共享代码和数据（建议用NFS）
配置SSH免密登录
防火墙开放通信端口

主节点启动命令：

python -m torch.distributed.launch \
    --nnodes=2 \
    --node_rank=0 \
    --master_addr="192.168.1.100" \
    --master_port=29500 \
    --nproc_per_node=8 \
    tools/train.py \
    configs/yolox/yolox_l_8xb8-300e_coco.py \
    --launcher pytorch

从节点命令：

python -m torch.distributed.launch \
    --nnodes=2 \
    --node_rank=1 \
    --master_addr="192.168.1.100" \
    --master_port=29500 \
    --nproc_per_node=8 \
    tools/train.py \
    configs/yolox/yolox_l_8xb8-300e_coco.py \
    --launcher pytorch

3.2 SLURM集群方案

在超算中心等环境中，通常用SLURM调度系统。配置起来反而更简单：

srun -p mm_dev \
    --job-name=mmdet_train \
    --gres=gpu:8 \
    --ntasks=16 \
    --ntasks-per-node=8 \
    --cpus-per-task=5 \
    --kill-on-bad-exit=1 \
    python tools/train.py \
    configs/yolox/yolox_x_8xb8-300e_coco.py \
    --launcher="slurm"

关键参数说明：

--gres=gpu:8：每节点申请8卡
--ntasks=16：总共16个进程
--ntasks-per-node=8：每节点8进程

4. 实战技巧与调优

4.1 学习率调整策略

分布式训练时，学习率需要随GPU数量线性放大。例如单卡用0.01，4卡应该用0.04。但实际我发现更好的公式是：

base_lr = 0.01
lr = base_lr * sqrt(num_gpus)  # 平方根缩放更稳定

在配置文件中可以这样设置：

optimizer = dict(
    type='SGD',
    lr=0.01 * sqrt(4),  # 假设用4卡
    momentum=0.9,
    weight_decay=0.0001)

4.2 通信优化技巧

多机训练时，网络带宽可能成为瓶颈。我总结了几点优化经验：

梯度压缩：使用torch.distributed.algorithms.ddp_comm_hooks
AllReduce异步化：设置broadcast_buffers=False
数据预处理分离：用CPU预处理减轻GPU负担

4.3 监控与调试

推荐使用MMEngine内置的日志系统，配合TensorBoard：

# 在配置文件中添加
visualizer = dict(
    type='DetLocalVisualizer',
    vis_backends=[dict(type='TensorboardVisBackend')],
    name='visualizer')

训练时可以用nvidia-smi配合watch命令实时监控：

watch -n 1 nvidia-smi

遇到问题时，可以先在小规模数据上测试：

CUDA_VISIBLE_DEVICES=0,1 torchrun \
    --nproc_per_node=2 \
    tools/train.py \
    configs/yolox/yolox_tiny_8xb8-300e_coco.py \
    --launcher pytorch \
    --cfg-options data.train.dataset.pipeline.0.img_scale=(320,320)