使用 YOLOv5 进行多 GPU 训练#
本指南介绍如何在单台机器上或多台机器之间使用多个 GPU 训练 YOLOv5。
开始之前#
克隆仓库，并在包含 PyTorch>=1.8 的 Python>=3.8.0 环境中安装 requirements.txt。模型和数据集会从 YOLOv5 最新版本自动下载。
git clone https://github.com/ultralytics/yolov5 # clone
cd yolov5
pip install -r requirements.txt # install建议所有多 GPU 训练任务都使用 Ultralytics Docker 镜像。请参阅 Docker 快速入门指南。
在 PyTorch >= 1.9 中，torch.distributed.run 取代了 torch.distributed.launch。详情请参阅 PyTorch 分布式文档。
训练#
选择一个预训练模型作为训练起点。这里我们选择了YOLOv5s，这是一款小型快速模型。查看我们的 README表格，了解所有模型的完整对比。我们将使用多 GPU 在 COCO 数据集上训练该模型。
单 GPU#
python train.py --batch 64 --data coco.yaml --weights yolov5s.pt --device 0多 GPU DataParallel 模式（⚠️ 不推荐）#
向 --device 传入多个 GPU ID，以启用 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（或任何其他列表）。
使用指定 GPU（点击展开）
传入 --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使用 SyncBatchNorm（点击展开）
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使用多台机器（点击展开）
这仅适用于多 GPU DistributedDataParallel 训练。
继续之前，请确保所有机器上的数据集、代码库和其他依赖项一致，然后确认这些机器能够通过网络相互连接。
选择一台主机器（其他机器将连接到该机器），记下它的地址（master_addr），并选择一个端口（master_port）。下面的示例使用 master_addr = 192.168.1.1 和 master_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 = 2、N = 2 和 R = 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 轮次的 DDP 性能分析结果。
代码性能分析
# 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 设备0 (G) | COCO 训练 | COCO 验证 |
|---|---|---|---|---|
| 1x | 16 | 26GB | 20:39 | 0:55 |
| 2x | 32 | 26GB | 11:43 | 0:57 |
| 4x | 64 | 26GB | 5:57 | 0:55 |
| 8x | 128 | 26GB | 3:09 | 0:57 |
如结果所示，使用 DistributedDataParallel 配合多块 GPU，训练速度几乎可以线性提升。使用 8 块 GPU 时，训练速度约为使用单块 GPU 的 6.5 倍，同时每个设备的内存用量保持不变。
支持的环境#
Ultralytics 提供多种开箱即用的环境，每种环境都预装了 CUDA、CUDNN、Python 和 PyTorch 等必备依赖，助你快速启动项目。
- 免费 GPU 笔记本：
- Google Cloud：GCP 快速入门指南
- Amazon：AWS 快速入门指南
- Azure：AzureML 快速入门指南
- Docker：Docker 快速入门指南
项目状态#
此徽章表示所有 YOLOv5 GitHub Actions 持续集成 (CI) 测试均已成功通过。这些 CI 测试会严格检查 YOLOv5 在多个关键方面的功能和性能：训练、验证、推理、导出和基准测试。这些测试每 24 小时以及每次提交新代码后运行一次，确保 YOLOv5 在 macOS、Windows 和 Ubuntu 上保持稳定可靠地运行。
致谢#
感谢 @MagicFrogSJTU 负责了所有繁重的工作，也感谢 @glenn-jocher 一路上的指导。
另请参阅#
- 训练模式 - 了解如何使用 Ultralytics 训练 YOLO 模型
- 超参数调优 - 优化模型性能
- Docker 快速入门指南 - 配置用于训练的 Docker 环境
常见问题#
打开问题前，请先阅读下面的检查清单——这通常能节省时间。
检查清单（点击展开）
- 你是否从头到尾阅读了本指南？
- 你是否重新克隆了代码库？代码每天都会更改。
- 你是否搜索过错误消息？可能已经有人遇到过同样的问题并分享了解决方法。
- 你是否安装了所有依赖项（包括正确版本的 Python 和 PyTorch）？
- 你是否尝试过上面列出的某个受支持环境？
- 你是否尝试过
coco128或coco2017等标准数据集，以找出根本原因？
如果以上检查都没有问题，请按照模板尽可能详细地提交一个 Issue。