YOLOv5多GPU训练

一、准备

安装yolov5依赖，要求Python>=3.8以及PyTorch>=1.7

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

二、选择预训练模型

这里我们选择YOLOv5s，训练数据集我们用COCO。

三、单GPU训练

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

四、多GPU数据并行训练（不推荐）

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

这种方法相对于单GPU训练方法对于速度的提升有限。

五、多GPU分布式数据并行训练（推荐）

$ python -m torch.distributed.launch --nproc_per_node 2 train.py --batch-size 64 --data coco.yaml --weights yolov5s.pt

–nproc_per_node： 指定多少GPU要被使用。
–batch-size：这里是总的大小，多个GPU这个值将被均分，比如以上示例有2个GPU，那每个GPU batch-size为64/2=32。

指定多GPU：
$ python -m torch.distributed.launch --nproc_per_node 2 train.py --batch-size 64 --data coco.yaml --cfg yolov5s.yaml --weights '' --device 2,3

当单个GPU batch-size<=8时，可以使用以下方式来提高准确率
$ python -m torch.distributed.launch --nproc_per_node 2 train.py --batch-size 64 --data coco.yaml --cfg yolov5s.yaml --weights '' --sync-bn

多设备训练
# On master machine 0
$ python -m torch.distributed.launch --nproc_per_node G --nnodes N --node_rank 0 --master_addr "192.168.1.1" --master_port 1234 train.py --batch-size 64 --data coco.yaml --cfg yolov5s.yaml --weights ''

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