09-训练自己的模型
在上一个章节中。我们用的模型全部是来自yolo官网。
其模型识别的是别人训练的class。通常情况下，这个通用的模型并不能满足我们的需求。
因此我们需要训练自己的模型。本文以yoloV6为例。训练一个检测是否带安全头盔的模型。
1、下载yoloV6源码
待整理后上传【占位】
2、准备训练图片以及标签
在yolov6的源文件中进入data文件夹，新建一个名为kmbox的文件夹，此文件夹用于存放训练图片和标签。如下图所示：
进入kmbox文件夹，再创建两个文件夹，名字分别为images,labels.如下图所示：
在images中创建两个文件夹，train和val。同理在labels中也创建train和val两个文件夹。
image中的train用于存放训练集（图片），val用于存放验证集（图片）。同理labels中存放对应的标签。
图片集如下（你想训练什么）：
标签集如下：
返回data目录，复制一下voc.yaml文件。重命名为kmbox.yaml
打开kmbox.yaml,做如下修改。
返回YOLOV6源码目录。将tools文件夹下的train.py脚本拷贝到源码目录下。取名kmbox.py
打开kmbox.py文件做如下修改。
3、准备训练环境
我们使用conda来搭建一个python3.8环境（windows电脑上操作）：
进入YOLOV6的源码目录。注意文件路径中不要由中文。
安装依赖：pip install -r requirements.txt 在控制台输入左边的指令。等待安装完成。可能需要科学上网或指定国内源。
4、开始训练
输入python kmbox.py回车。
法克报错。大概是cuda不能用，应该是驱动不兼容。解决办法。重新安装torch。
pip install torch==2.3.1+cu118 torchvision==0.18.1+cu118 torchaudio==2.3.1+cu118 -f https://mirror.sjtu.edu.cn/pytorch-wheels/torch_stable.html
你的电脑是哪个版本可能得你自己解决。因为不同显卡对应的cuda不一样。不同torch配套的cuda也不一样。
再次运行python kmbox.py 可以看到已经开始正常训练了。GPU风扇在嗷嗷叫。等他训练完成即可。
如果训练过程中崩溃了。报各种错误。最简单的方法就是关机再开机重来。保证内存充足。炼丹嘛，炸炉是常有的事。炸了重来就行。如果一直炸，那就是炉子不行。再换环境。直到能用为止。
经过漫长等待。训练已经完成。
5、验证模型是否满足识别要求
训练的模型保存在runs文件夹下。train中
咱们取最后一次训练的权重。
为了方便起见（少敲点键盘）。把best_ckpt.pt复制到顶层YOLOv6的顶层目录中。随便拿一张测试图片，取名img.png. 执行下面指令用于验证。
python tools/infer.py --weights best_ckpt.pt --source img.png
执行命令后截图如下：
可以看到最后一行执行结果再runs/inference/exp中。打开对应文件夹看看结果。
法克。bicycle和person是什么鬼？？？别急。肯定是标签错了。打开tools/infer.py。
另外咱们训练图片尺寸是416x416的，这里也可以改成一致。改完后再次运行。上面的指令验证。
OK.已经能正常识别是否带头盔。
如果识别效果不好，那就换图片换标签。重新训练。直到能达到你的满意度为止。重复前面的234过程即可。
我很容易满足。上面已经达到我的预期。进入下一步。
6、导出rknn格式
上面训练出来的模型都是电脑上用的。我们最终需要是在kmbox内部运行。因此需要将模型进行格式转换。
导出模型执行下面这条指令：
python deploy/RKNN/export_onnx_for_rknn.py --weight ./best_ckpt.pt
导出后模型名字为best_ckpt.onnx
最后将best_ckpt.onnx模型转换为rknn。就是咱们盒子能用的格式。即可下载到kmbox内部使用。我们把best_ckpt.onnx拷贝到SDK内部的09-myYOLO-test/tools文件夹中。
执行命令：python convert.py best_ckpt.onnx可将其转换为yolov6.rknn （这是在虚拟机内操作）
至此，yolov6.rknn就是我们需要的能在开发板内部运行的模型。
7、板载验证
编译09-myYOLO-test工程。将得到的可执行文件和前面的yolov6.rknn下载到开发板中
并且准备一张测试图片。test.jpg 。
编译方法请参考前面章节。
运行截图如下：
实际效果图：
以上为yolov6训练自定义数据集到板载验证的全部流程。以此为例。可以训练用户自定义模型。
从打印可以看出。416x416的模型，推理时间大约25ms。比上一章中的640X640模型速度提升了大概20ms。基本上是成倍的增加。如果有兴趣可以测试一下320X320的。看看推理速度又是多少。
整个工程已打包上传。
FAQ常见问题
如何退出程序
1、控制台直接按键盘ctr+c
2、按键盘End快捷键（切换到内部UI快捷键，以你设置的为准）