官网给的:
https://github.com/ultralytics/yolov5/issues/475

指定显卡

指定单显卡:

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

   
   
  • 1

使用pytorch 的DataParallel mode 指定多显卡(不推荐),为啥不推荐,因为反传的时候梯度默认都在0卡上算了,所以显存分配不均衡,在显卡0占用非常多,别的显卡的显存又占用没那么多。训练时间也没增加多少。

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

   
   
  • 1

单机多卡(推荐)

使用pytorch的 DistributedDataParallel Mode 进行单机多卡训练。(推荐)

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

   
   
  • 1

–nproc_per_node specifies how many GPUs you would like to use. In the example above, it is 2.
–batch is the total batch-size. It will be divided evenly to each GPU. In the example above, it is 64/2=32 per GPU.

我机器上三张显卡,我这么用的:

python -m torch.distributed.launch --nproc_per_node 3 train.py --batch-size 6 --data kevin_firesmoke.yaml --img 1280 --epochs 100 --weight /home/kevin_xie/yolov5_version6/weights/yolov5m6.pt --optimizer "Adam" --save-period 20 --device 0,1,2

   
   
  • 1

在这里插入图片描述

多机多卡

参考这里:
https://blog.csdn.net/flyfish1986/article/details/119786227
https://pytorch.org/docs/stable/distributed.html

Logo

旨在为数千万中国开发者提供一个无缝且高效的云端环境,以支持学习、使用和贡献开源项目。

更多推荐