YOLOv5系列(四十九) 解读多GPU，SyncBatchNorm训练

介绍了如何在单台或多台机器上正确使用 多个 GPU 通过 YOLOv5 训练数据集。

前言

克隆存储库和安装要求.txt在 Python>=3.8.0 环境中，包括 PyTorch>=1.8。模型和数据集从最新的 YOLOv5 版本自动下载。

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

专业提示！建议将 Docker Image 用于所有多 GPU 训练。

专业提示！ 替换为 PyTorch>=1.9。有关详细信息，请参阅文档。torch.distributed.run``torch.distributed.launch

训练

选择要从中开始训练的预训练模型。在这里，我们选择YOLOv5s，这是最小和最快的型号。请参阅我们的 README 表，了解所有型号的完整比较。我们将在 COCO 数据集上使用 Multi-GPU 训练此模型。

单 GPU

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

多GPU数据并行模式（不推荐）

您可以增加 在 DataParallel 模式下使用多个 GPU。device

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

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

多GPU DistributedData并行模式（推荐）

你必须通过 ，然后是通常的论点。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。 是总批大小。它将平均分配给每个 GPU。在上面的示例中，每个 GPU 为 64/2=32。`--batch

上面的代码将使用 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可以提高多 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

使用DistributedDataParallel

这仅适用于 Multiple GPU DistributedDataParallel 训练。 在继续之前，请确保所有计算机上的文件都相同，数据集、代码库等。之后，确保机器可以相互通信。 您将不得不选择一台主机器（其他人将与之交谈的机器）。记下其地址（‘master_addr’）并选择一个端口（‘master_port’）。在下面的示例中，我将使用“master_addr = 192.168.1.1”和“master_port = 1234”。 要使用它，您可以执行以下操作，

1. 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 ''

1. 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 ...

结果

在具有 2 个 A4 SXM8-100GB 的 AWS EC4 P40d 实例上执行 5 个 COCO 纪元的 YOLOv1l 上的 DDP 分析结果。

图形处理器 A100	批处理大小	CUDA_mem设备0 （G）	COCO 火车	COCO val
1 倍	16	26千兆字节（GB）	20:39	0:55
2 倍	32	26千兆字节（GB）	11:43	0:57
4 倍	64	26千兆字节（GB）	5:57	0:55
8 倍	128	26千兆字节（GB）	3:09	0:57