跳至内容

使用YOLOv5进行GPU GPU 训练

本指南介绍如何在单台或多台计算机上正确使用多个GPU 来使用YOLOv5 🚀 训练数据集。

开始之前

克隆 repo 并将requirements.txt安装在 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 培训。参见 Docker 快速入门指南 Docker 拉动

专业提示

torch.distributed.run 替换 torch.distributed.launchPyTorch>=1.9.参见 PyTorch 分布式文档 了解详情。

培训

选择一个预训练模型开始训练。这里我们选择YOLOv5s,它是目前最小、最快的模型。有关所有模型的全面比较,请参见我们的 README表格。我们将使用 Multi-GPU 在COCO数据集上训练该模型。

YOLOv5 机型

单人GPU

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

您可以增加 device 以在数据并行模式下使用多个 GPU。

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

这种方法很慢,与只使用 1 个GPU 相比,几乎无法加快训练速度。

您必须通过 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).

使用特定 GPU(点击展开) 您只需传递 `--device`,然后再传递特定的 GPU 即可。例如,在下面的代码中,我们将使用 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](https://pytorch.org/docs/master/generated/torch.nn.SyncBatchNorm.html) could increase [accuracy](https://www.ultralytics.com/glossary/accuracy) for multiple GPU training, however, it will slow down training by a significant factor. It is **only** available for Multiple GPU DistributedDataParallel training. It is best used when the batch-size on **each** GPU is small (<= 8). To use SyncBatchNorm, simply pass `--sync-bn` to the command like below:
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 ...

成果

AWS EC2 P4d 实例上对 YOLOv5l 的 8x A100 SXM4-40GB 进行 1 COCOepoch 的 DDP 分析结果。

剖析代码
# prepare
t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus 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
图形处理器
A100
批量大小 CUDA_mem
device0 (G)
COCO
火车
COCO
val
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

如结果所示,使用多 GPU 的DistributedDataParallel可以提供近乎线性的训练速度。与单个 GPU 相比,使用 8 个GPU 完成训练的速度约快 6.5 倍,同时每个设备的内存使用量保持不变。

常见问题

如果出现错误,请先阅读下面的核对表!(它可以节省您的时间)

核对表(点击展开) - 你认真阅读过这篇文章吗? - 你试过重新克隆代码库吗?代码每天都在变化。 - 您尝试搜索过您的错误吗?可能已经有人在本版本库或其他版本库中遇到过这个问题,并且已经有了解决方案。 - 你是否安装了上面列出的所有需求(包括正确的Python 和PyTorch 版本)? - 你在下面 "环境 "部分列出的其他环境中试过吗? - 您是否尝试过其他数据集,如 coco128 或 coco2017?这样会更容易找到根本原因。 如果您完成了上述所有步骤,请按照模板尽可能详细地提供信息,随时提出问题。

支持的环境

Ultralytics 提供了一系列随时可用的环境,每个环境都预装了基本的依赖项,如 CUDACUDNNPythonPyTorch等基本依赖项,以便启动项目。

项目现状

YOLOv5 CI

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

荣誉

在此,我们要感谢 @MagicFrogSJTU 和 @glenn-jocher 对我们的指导。

另请参见

📅创建于 1 年前 ✏️已更新 7 天前

评论