YOLO Vision 2026:

使用 YOLOv5 进行多 GPU 训练#

本指南介绍如何在单台机器上或跨多台机器使用多个 GPU 训练 YOLOv5。

开始之前#

Python>=3.8.0 环境中克隆仓库并安装 requirements.txt,其中包括 PyTorch>=1.8模型数据集 会自动从最新的 YOLOv5 版本下载。

git clone https://github.com/ultralytics/yolov5 # clone
cd yolov5
pip install -r requirements.txt # install
使用 Docker

所有多 GPU 训练运行都推荐使用 Ultralytics Docker 镜像。请参阅 Docker 快速入门指南Docker Pulls

PyTorch >= 1.9

PyTorch >= 1.9 中,torch.distributed.run 取代了 torch.distributed.launch。详情请参阅 PyTorch 分布式文档

训练#

选择一个预训练模型作为训练起点。这里我们选择 YOLOv5s,这是一个小型且快速的模型。有关所有模型的完整对比,请参阅我们的 README 表格。我们将在 COCO 数据集上使用多 GPU 训练此模型。

YOLOv5 Models

单 GPU#

python train.py --batch 64 --data coco.yaml --weights yolov5s.pt --device 0

多 GPU DataParallel 模式(⚠️ 不推荐)#

将多个 GPU ID 传递给 --device 以启用 DataParallel 模式:

python train.py --batch 64 --data coco.yaml --weights yolov5s.pt --device 0,1

与使用单个 GPU 相比,DataParallel 速度较慢,训练速度几乎没有提升。

多 GPU DistributedDataParallel 模式(✅ 推荐)#

在训练命令前加上 python -m torch.distributed.run --nproc_per_node,然后传入常用参数:

python -m torch.distributed.run --nproc_per_node 2 train.py --batch 64 --data coco.yaml --weights yolov5s.pt --device 0,1
  • --nproc_per_node 是要使用的 GPU 数量。在上面的示例中,它是 2
  • --batch 是总批量大小,会在每个 GPU 之间平均分配。在上面的示例中,每个 GPU 分配到 64 / 2 = 32

上面的命令使用 GPU 0...(N-1)。如需改用环境变量控制设备可见性,请在启动前设置 CUDA_VISIBLE_DEVICES=2,3(或其他列表)。

Use specific GPUs (click to expand)

传入 --device,后跟指定的 GPU ID。下面的示例使用 GPU 2,3

python -m torch.distributed.run --nproc_per_node 2 train.py --batch 64 --data coco.yaml --cfg yolov5s.yaml --weights '' --device 2,3
Use SyncBatchNorm (click to expand)

SyncBatchNorm 可以提高多 GPU 训练的准确率,但会显著降低训练速度。它适用于多 GPU DistributedDataParallel 训练。

最适合在每个 GPU 的批量大小较小时使用(<= 8)。

要启用 SyncBatchNorm,请传入 --sync-bn

python -m torch.distributed.run --nproc_per_node 2 train.py --batch 64 --data coco.yaml --cfg yolov5s.yaml --weights '' --sync-bn
Use Multiple machines (click to expand)

适用于多 GPU DistributedDataParallel 训练。

继续之前,请确保所有机器上的数据集、代码库和其他依赖项保持一致,然后验证这些机器可以通过网络互相访问。

选择一台主机器(其他机器将连接到该机器),记下其地址(master_addr),并选择一个端口(master_port)。下面的示例使用 master_addr = 192.168.1.1master_port = 1234

然后运行:

# On master machine 0
python -m torch.distributed.run --nproc_per_node G --nnodes N --node_rank 0 --master_addr "192.168.1.1" --master_port 1234 train.py --batch 64 --data coco.yaml --cfg yolov5s.yaml --weights ''
# On machine R
python -m torch.distributed.run --nproc_per_node G --nnodes N --node_rank R --master_addr "192.168.1.1" --master_port 1234 train.py --batch 64 --data coco.yaml --cfg yolov5s.yaml --weights ''

其中,G 是每台机器上的 GPU 数量,N 是机器数量,R 是机器在 0...(N-1) 中的排名。例如,有两台机器且每台有两个 GPU 时,请在第二台机器上设置 G = 2N = 2R = 1

只有在所有 N 台机器连接后,训练才会开始。输出仅显示在主机器上。

备注#

  • Windows 支持尚未经过测试;建议使用 Linux。

  • --batch 必须是 GPU 数量的倍数。

  • GPU 0 使用的内存略多于其他 GPU,因为它需要维护 EMA 并处理检查点保存。

  • 如果你遇到 RuntimeError: Address already in use,通常意味着多个训练运行正在使用同一个端口。请使用 --master_port 指定其他端口:

    python -m torch.distributed.run --master_port 1234 --nproc_per_node 2 ...

结果#

在配备 8 个 A100 SXM4-40GB 的 AWS EC2 P4d 实例上,对 YOLOv5l 进行 1 个 COCO epoch 的 DDP 性能分析结果。

Profiling code
# prepare
t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --device nvidia.com/gpu=all -v "$(pwd)"/coco:/usr/src/coco $t
pip3 install torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
cd .. && rm -rf app && git clone https://github.com/ultralytics/yolov5 -b master app && cd app
cp data/coco.yaml data/coco_profile.yaml

# profile
python train.py --batch-size 16 --data coco_profile.yaml --weights yolov5l.pt --epochs 1 --device 0
python -m torch.distributed.run --nproc_per_node 2 train.py --batch-size 32 --data coco_profile.yaml --weights yolov5l.pt --epochs 1 --device 0,1
python -m torch.distributed.run --nproc_per_node 4 train.py --batch-size 64 --data coco_profile.yaml --weights yolov5l.pt --epochs 1 --device 0,1,2,3
python -m torch.distributed.run --nproc_per_node 8 train.py --batch-size 128 --data coco_profile.yaml --weights yolov5l.pt --epochs 1 --device 0,1,2,3,4,5,6,7
GPU
A100
批量大小CUDA_mem
device0 (G)
COCO
训练
COCO
验证
1x1626GB20:390:55
2x3226GB11:430:57
4x6426GB5:570:55
8x12826GB3:090:57

如结果所示,使用 DistributedDataParallel 配合多个 GPU,可以使训练速度近似线性扩展。使用 8 个 GPU 时,训练完成速度约为单 GPU 的 6.5 倍,同时保持每个设备的内存使用量不变。

支持的环境#

Ultralytics 提供一系列可直接使用的环境,每个环境都预装了 CUDACUDNNPythonPyTorch 等必要依赖项,帮助你快速启动项目。

项目状态#

YOLOv5 CI

此徽章表示所有 YOLOv5 GitHub Actions 持续集成(CI)测试均已成功通过。这些 CI 测试会严格检查 YOLOv5 在多个关键方面的功能和性能:训练验证推理导出基准测试。它们确保 YOLOv5 在 macOS、Windows 和 Ubuntu 上持续稳定地运行,并且每 24 小时以及每次新提交时都会执行测试。

致谢#

感谢 @MagicFrogSJTU 承担了大量工作,也感谢 @glenn-jocher 一路上的指导。

另请参阅#

常见问题#

在提交 Issue 前请阅读下面的检查清单——这通常可以节省时间。

Checklist (click to expand)
  • 你是否已完整阅读本指南?
  • 你是否重新克隆了代码库?代码每天都会变化。
  • 你是否搜索过错误消息?可能已经有人遇到过相同问题并分享了解决方法。
  • 你是否安装了所有依赖项(包括正确版本的 Python 和 PyTorch)?
  • 你是否尝试过上面列出的受支持环境之一?
  • 你是否尝试过使用 coco128coco2017 等较小的数据集,以隔离根本原因?

如果以上各项都没有问题,请按照模板尽可能详细地提交 Issue。

评论