一、下载:
GitHub - meituan/YOLOv6: YOLOv6: a single-stage object detection framework dedicated to industrial applications.https://github.com/meituan/YOLOv6我下载的版本是 YOLOv6-0.4.0。
代码结构如下:
二、数据集结构
可以直接使用yolov5的数据集存放方式,下面是我自己的数据集结构:
---imgDir
---images
---train
---val
---labels
---train
---val
数据集的配置文件,在data文件里面新建一个mydata.yaml文件,然后仿照coco.yaml的内容,写自己数据集的路径配置文件,如下:
三、训练
可以使用终端,参考官网的用命令行训练:
也可以修改在tools文件夹下的train.py文件的相关参数,训练自己的数据集:
主要修改data-path和conf-file参数,data-path修改为自己数据集对应的yaml文件,conf-file修改为自己需要训练的模型类型(如yolov6n, yolov6s等)。然后找到yolov6n.py文件:
将pretrained参数修改为自己存放预训练参数的位置,还有一些其他的配置可以修改。
都修改完成之后就可以运行train.py开始训练。这里我只跑了10个epoch,测试一下:
四、验证和推理
同理,在tools文件夹下也有eval.py和infer.py,只需要修改部分参数即可。
eval.py:
将data参数修改为自己数据集的配置文件,将weights参数修改为刚训练出来的pt文件。
infer.py同样:
weights参数修改为自己训练出来的pt文件,source参数修改为存放要推理图像的文件夹,也可以具体的某一张图像。
如:
parser.add_argument('--source', type=str, default=r'D:\image\images\val\0_12164454943_9.bmp', help='the source path, e.g. image-file/dir.')
五、onnx模型导出
模型的到处代码在:
只需要对应的修改掉里面的参数即可。