模型训练
训练根据计算资源(显卡)的情况,主要分3种: 单机单卡, 单机多卡,多机多卡。
需要说明的是,Yolo v5的作者采用的不是熟悉的DP模式torch.nn.DataParallel,而是torch.nn.parallel.DistributedDataParallel, 即DDP模式。按照PyTorch官方的说法[3]: 即使在单机多卡的环境,也推荐DDP模式,DDP模式会显著快于DP模型,关于DDP的例子可见官网[4].
- 单机单卡训练:
python train.py --img 640 --batch 16 --epochs 5 --data ./data/coco128.yaml --cfg ./models/yolov5s.yaml --weights ''
- 单机多卡训练:
python train.py --img 640 --batch 16 --epochs 5 --data ./data/coco128.yaml --cfg ./models/yolov5s.yaml --weights '' --device 0,1 # 指定第1,2块卡进行训练
- 多机多卡训练:
python -m torch.distributed.launch --nproc_per_node 2 train.py --weights yolov5s.pt --data ./data/coco128.yaml --cfg ./models/yolov5s.yaml --epochs 3 --img 320 --device 0,1 # DDP
参考资料:https://blog.csdn.net/g11d111/article/details/108872076