[个人技术博客] YOLOv5 模型训练，并移植至 OpenCV，若干注意事项

最新推荐文章于 2024-06-21 17:36:43 发布

071903410-林章鹏

最新推荐文章于 2024-06-21 17:36:43 发布

阅读量1.2k

点赞数 1

文章标签： python 计算机视觉 pytorch

本文链接：https://blog.csdn.net/qq_35528522/article/details/125424774

版权

本文介绍了如何使用YOLOv5训练自定义数据集，建议至少训练16个epoch，批处理大小让PyTorch自动决定。接着详细讲述了如何将模型导出为ONNX格式供OpenCV使用，并给出了在Python中使用cv2的示例，注意坐标转换和NMS Threshold的调整以避免目标检测错误。

摘要由CSDN通过智能技术生成

制作数据集

手动标记数据的时候，注意方框要与目标物体边缘相切。

模仿COCO128 的实例即可，此处没有太多需要注意的地方。（除了 Windows 路径要使用双反斜杠 \\ ）

# 官方文档使用的命令
# Train YOLOv5s on COCO128 for 5 epochs
python train.py --img 640 --batch 16 --epochs 5 --data coco128.yaml --weights yolov5s.pt

实际上，5 epoch 不足以训练出可靠的模型。由如下图表可知：epoch 7 以后，metrics/precision 才超过 0.8。

笔者建议训练至少 16 epoch。

另外，同时只训练一个模型时，batch_size 建议让 PyTorch 自行决定。最终建议命令如下：

python train.py --img 640 --batch -1 --epochs 16 --data coco128.yaml --wei

关注