yolov6训练自己的数据记录+yolov5对比测试

yolov6训练自己的数据记录+yolov5对比测试

美团最近发布了yolov6版本,恩,不管怎么说,还是实验一下效果怎么样吧。以下博文为实验过程记录。
代码:

https://github.com/meituan/YOLOv6

一、电脑配置

1.Ubuntu20
2.Cuda 11.2 + cudnn

二、环境配置

下载代码和环境配置
下载代码解压后,使用pycharm打开yolov6文件夹,因为yolov6文件夹中包含requirements.txt文件,打开后pycharm会自动识别,出现

选好位置,进行安装即可。
注:直接这样安装,pytorch可能会安装成cpu版本,可以先在requirement.txt中,将torch注释一下,后面在pycharm终端中手动安装pytorch的gpu版本。

进入pytorch官网选择适合自己电脑配置的pytorch-gpu版本

三、yolov6测试

1.下载权重文件
在yolov6的github首页进行下载,下载后新建weights文件夹,存放权重。

2.修改代码
省的用指令跑代码麻烦
1.修改infer.py

2.修改inferer.py

‘…/yolov6/utils/Arial.ttf’,这一行多加一个点。
3.测试结果
直接看图

四、yolov5与yolov6做对比

双方都控制在640的尺寸作对比
yolov6参数

yolov5-6.1参数

yolov5s.pt结果如下:

yolov6s.pt结果如下:

对比两者s的模型,整体来看yolov6得分确实比yolov5-6.1要高一点,不足就是在中间那张图里面,yolov6漏检了一根领带。

五、训练自己的数据集

在yolov6的github上有相关介绍,做为本次的主要参考
https://github.com/meituan/YOLOv6/blob/main/docs/Train_custom_data.md
1.数据标注
在yolov6的github的教程中有介绍,使用的是yolo的数据格式。

利用标注工具进行标注即可。
2.组织目录
在data文件夹下分别新建images和labels两个文件夹,并在下一级新建三个文件夹,具体如下图:


3.编写数据配置文件
在data文件夹下,新建dataset.yaml文件,并写入:

train: path/to/data/images/train # train images
val: path/to/data/images/val # val images
test: path/to/data/images/test # test images (optional)

# Classes
nc: 20  # number of classes
names: ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
        'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']  # class names

按照自己的数据集情况进行修改
我的是

train: /home/cj/work/code/python/yolo/YOLOv6/data/images/train # train images
val: /home/cj/work/code/python/yolo/YOLOv6/data/images/val # val images
test: /home/cj/work/code/python/yolo/YOLOv6/data/images/test # test images (optional)

# Classes
nc: 1  # number of classes
names: ['QR']  # class names

4.编写模型配置文件
这一步骤参考:https://github.com/meituan/YOLOv6/blob/main/docs/Train_custom_data.md编写,也可以直接使用已经存在的几个文件

5.开始训练
pycharm终端输入

python3 tools/train.py --batch 32 --conf configs/yolov6s_finetune.py --data data/dataset.yaml --device 0

注意batch值的选择,太大容易不够内存。

遇到的问题:
问题1
RuntimeError: Unable to find a valid cuDNN algorithm to run convolution
原因:是batch太大,内存不够,降低batch数值即可,如果数值很小,还是不行的话,重启电脑释放一下显卡内存。

训练好的权重会保留在runs文件夹下

6.测试
修改一下infer.py文件

在tools/runs下生成结果

OK!搞定。

来源:小俊俊的博客

物联沃分享整理
物联沃-IOTWORD物联网 » yolov6训练自己的数据记录+yolov5对比测试

发表评论