一、数据集准备
数据集的准备包括数据集适配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
表示使用电脑摄像头进行检测。
检测结果如下:
文章出处登录后可见!
已经登录?立即刷新