通过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表示使用电脑摄像头进行检测。

检测结果如下：