YOLOv5多GPU训练

一、准备

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。