通过YOLO5训练自己的数据集(以交通标志牌数据集TT100k为例)

一、数据集准备

数据集的准备包括数据集适配YOLO格式的重新分配以及相应配置文件的书写,此处可查看博主的TT100K2yolo的重新分配博文,该文章包括数据集划分,配置文件书写,以及最终的数据集层级目录组织,可以直接提供给下一步进行训练。

二、训练

可参见YOLO官方文件里的可参见readme文档

2.1 train.py

修改配置文件参数。

官方代码中这块默认yaml文件是coco128.yaml

在这里插入图片描述

同理,这里直接将上节中的yaml文件rett100k.yaml替换掉coco128.yaml即可。

在这里插入图片描述

这里是在代码中修改,还有一种方式是在训练的时候输入yaml文件地址,可参见readme文档

2.2 训练代码运行

当然首先需要进入YOLOv5目录下,博主是在新建的虚拟环境中配置的yolov5,所以需要先激活pytorch环境:
在这里插入图片描述可以先通过以下命令查看虚拟环境:

conda info -e

博主yolov5的虚拟环境是在jpytorch下的,所以执行:

source activate jpytorch

如此便可以进行训练操作了。

再执行以下语句:

python train.py --data rett100k.yaml --cfg yolov5s.yaml --weights '' --batch-size 64

在这里插入图片描述
这里batch_size可根据自身计算机性能选择:

在这里插入图片描述
注意:博主用的是3080GPU,选择的64,结果都出现了图里上半部分的错误:RuntimeError: Unable to find a valid cuDNN algorithm to run convolution,于是重新选择了40便开始正常训练了;另外训练期间不要在终端用键盘,否则会打断训练,只能重新来过。

在这里插入图片描述

接下坐等其训练结束就好啦

2.3 测试代码运行

最简单的方式就是将图片放置到如图目录下,

在这里插入图片描述

python detect.py --weights runs/train/exp14/weights/best.pt --source test/test --output run/now

在这里插入图片描述
注意权重地址按照如图所示给,选择最后一个exp里边的pt文件。
在这里插入图片描述
这里--source 0表示使用电脑摄像头进行检测。

检测结果如下:

在这里插入图片描述
在这里插入图片描述

文章出处登录后可见!

已经登录?立即刷新

共计人评分,平均

到目前为止还没有投票!成为第一位评论此文章。

(0)
乘风的头像乘风管理团队
上一篇 2023年2月26日 上午9:16
下一篇 2023年2月26日 上午9:17

相关推荐