部署YOLOV5自定义模型
0.前言
本章讲述如何训练自定义数据集生成模型,部署到DongshanPI-AICT开发板上。这里假设您已经搭建好YOLOV5-V6.0的环境,搭建环境参考:YOLOV5-V6.0环境搭建。如果您没有阅读过DongshanPI-AICT开发板支持yolov5模型部署文章,请先按照这篇文章进行操作。
下面操作仅演示如何去训练自定义模型、导出模型、转换模型、模型部署。注意:训练模型对于电脑需要有一定的要求,如果电脑性能较弱可能会导致训练效果较差,从而导致模型精度较低。
参考链接:https://docs.ultralytics.com/yolov5/tutorials/train_custom_data/
1.下载数据标注工具
数据标注工具:https://github.com/heartexlabs/labelImg/releases
点击上述红框下载,下载完成后解压压缩包,双击打开labelImg.exe
文件。
打开后等待运行,运行完成后会进入如下标注工作界面。
关于LabelImg更多的使用方法,请访问:https://github.com/heartexlabs/labelImg
由于LabelImg会预先提供一些类供您使用,需要手动删除这些类,使得您可以标注自己的数据集。步骤如下所示:
进入LabelImg程序目录中的data目录中,打开predefined_classes.txt
文件,删除文件中所有预定义的类后保存并退出即可。
2.创建数据集目录
在任意工作目录中创建images
文件夹和labels
文件夹分别存放图像数据集和标注信息。这里我演示仅使用少量图像样本进行标注,在实际项目中需要采集足够的图像进行标注才拿满足模型的准确率和精度。
例如我在100ask-yolov5-image
目录中创建有images
文件夹和labels
文件夹,如下所示,创建images文件,存放图像数据集,创建labels文件夹,该文件夹用于后续存放标注数据。
3.标注图像
打开LabelImg软件后,使用软件打开数据集图像文件夹,如下所示:
打开后,修改输出label的文件夹为我们创建的数据集目录下的labels
文件夹
标注完成后后会在labels
目录下生成classes.txt
(类别)和图像中标注的类别即位置信息。
下面为LabelImg快捷键目录:
Ctrl + u | Load all of the images from a directory |
---|---|
Ctrl + r | Change the default annotation target dir |
Ctrl + s | Save |
Ctrl + d | Copy the current label and rect box |
Ctrl + Shift + d | Delete the current image |
Space | Flag the current image as verified |
w | Create a rect box |
d | Next image |
a | Previous image |
del | Delete the selected rect box |
Ctrl++ | Zoom in |
经过标注大量的图像后,labels
文件夹如下图所示
4.划分训练集和验证集
在模型训练中,需要有训练集和验证集。可以简单理解为网络使用训练集去训练,训练出来的网络使用验证集验证。在总数据集中训练集通常应占80%,验证集应占20%。所以将我们标注的数据集按比例进行分配。
在yolov5-6.0项目目录下创建100ask文件夹(该文件夹名可自定义),在100ask文件夹中创建train文件夹(存放训练集)和创建val文件夹(存放验证集)
在train文件夹中创建images文件夹和labels文件夹。其中images文件夹存放总数据集的80%的图像文件,labels文件夹存放与images中的文件对应的标注文件。
在val文件夹中创建images文件夹和labels文件夹。其中images文件夹存放总数据集的20%的图像文件,labels文件夹存放与images中的文件对应的标注文件。
5.创建数据集配置文件
进入yolov5-6.0\data目录下,创建data.yaml
,文件内容如下所示:
train: 100ask\train\images # train images
val: 100ask\val\images # val images
nc: 3 # number of classes
names: ['T113', 'K510', 'V853'] # class names
6.创建模型配置文件
进入models目录下,拷贝yolov5s.yaml文件,粘贴并models目录下重命名为100ask_my-model.yaml,例如:
修改100ask_my-model.yaml中类的数目为自己训练模型的类数目。
6.修改训练函数
打开yolov5-6.0项目文件夹中的train.py,修改数据配置文件路径,如下图红框所示:
parser.add_argument('--cfg', type=str, default='models/100ask_my-model.yaml', help='model.yaml path')
parser.add_argument('--data', type=str, default=ROOT / 'data/data.yaml', help='dataset.yaml path')
7.训练模型
在conda终端的激活yolov5环境,激活后进入yolov5-6.0项目文件夹。执行python train.py
,如下图所示:
程序默认迭代300次,等待训练完成...
训练完成后结果会保存在runs\train\
目录下最新一次的训练结果,如上图所示,此次训练的最好模型和最后训练的模型保存在以下目录中
runs\train\exp7\weights
8.验证模型
修改val.py函数,修改如下
parser.add_argument('--data', type=str, default=ROOT / 'data/data.yaml', help='dataset.yaml path')
parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'runs/train/exp7/weights/best.pt', help='model.pt path(s)')
修改models文件夹下的yolo.py
class Model(nn.Module):
def __init__(self, cfg='100ask_my-model.yaml', ch=3, nc=None, anchors=None): # model, input channels, number of classes
打开conda终端输入python val.py
执行完成后的结果保存在runs\val\exp
文件下。