代码收藏家 yolov5-5.0训练模型+瑞芯微rv1126上实现模型部署
yolov5-5.0训练模型+瑞芯微rv1126上实现模型部署
第一次接触模型训练和在开发板部署，过程曲折，从开始的一脸懵到最后模型部署成功，查阅了不少资料和学习了不少大佬的经验，在这里记录一下过程和心得。
一、总体思路
首先在本机上使用yolov5代码训练自己的数据集，得到训练后的pt模型，然后通过export.py转化出onnx模型，最后通过瑞芯微官方代码，将onnx模型转化为rknn模型，实现在rv1126开发板上实现模型部署。
ps：为什么不直接使用pytorch模型转化为rknn？
答：rknn1.7.1直接加载pt模型会有bug。
二、实现过程
1、yolov5训练
我是在本机windows训练的模型，所以环境也是装在windows上，创建环境我用的是Anaconda，具体怎么下载安装有很多成功的博客，麻烦大家自己搜索。
首先拉取官方代码，我使用的yolov5代码版本是v5.0 release版本，项目地址是：https://github.com/ultralytics/yolov5/tree/v5.0
conda创建环境指令如下：（python版本是3.7）
conda create --name yolov5 python==3.7
配置训练所需的环境
进入虚拟环境：
activate yolov5
推荐使用清华源下载：
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -r requirements.txt
数据标注准备
安装labelimg
pip install labelimg
使用labelimg标注
点击Open Dir -> 选择需要标注的文件夹 → ok
进行标注格式切换，切换到yolo。
点击Create RectBox → 拖拽鼠标框选目标 → 给上标签 → 点击ok。
点击save，保存txt。
修改yolov5项目文件
1、yolov5/data文件夹下新建fire.yaml（名字自取），用来存放训练集和验证集的划分文件（train.txt和val.txt），参考如下：
train: D:\Lenovo\Desktop\v5-5.0\paper_date\images\train\ val: D:\Lenovo\Desktop\v5-5.0\paper_date\images\val\ # Classes nc: 1 # number of classes names: ['plasticbottle'] # class names
2、model下修改yolov5s.yaml的种类数nc=1，（因为我这里只训练了一种类别，具体根据自己需要改）
3、如果显卡显示不足的可以将dataset.py中的num_worker数量设置为0.
4、yolo.py修改detect设置nc的值，并且替换掉dectet模块
参考：https://zhuanlan.zhihu.com/p/450411127 第三点
5、common.py补充缺少的模块定义
参考：https://zhuanlan.zhihu.com/p/450411127 第二点
6、export.py添加参数 parser.add_argument(‘–grid’, action=‘store_true’, help=‘export Detect() layer grid’)
7、（*）这点很重要，因为rknn官方代码不支持onnx输出的Detect层，因此在运行完train.py训练出pt模型后，使用export.py前需要将
yolo.py中的代码x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
修改为
x[i] = x[i].view(bs, self.na, self.no, ny, nx)而当转化出onnx模型后，若想重新训练新的模型，需要将对yolo.py上述改动修改回去，否则会出现报错！！
参考：35、ubuntu20.04搭建瑞芯微的npu仿真环境和测试rv1126的Debain系统下的yolov5+npu检测功能以及RKNN推理部署
train.py参数设置：
上述做完就可以开始愉快的训练和转onnx模型了。
以下是我遇到的报错：
报错1
OMP: Error #15: Initializing libiomp5md.dll, but found libiomp5md.dll already initialized.
在pycharm里调试程序时可以直接通过在程序前添加这两个语句解决
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
报错2
RuntimeError: result type Float can't be cast to the desired output type __int64
参考：https://blog.csdn.net/Thebest_jack/article/details/125649451
报错3
RuntimeError: The size of tensor a (52) must match the size of tensor b (84) at non-singleton dimension 3
参考：https://blog.csdn.net/Thebest_jack/article/details/125649451
（这就是我上面说的训练没有将yolo.py中onnx的改动修改回去发生的报错。）
2、模型转化（onnx–>rknn）
我这里是将rknn官方的代码拆分成两个部分，一是模型转化，二就是在板子上进行后处理以及推理。
这里贴出onnx转化rknn的代码（通过官方代码修改），量化需要的数据集和标签自己修改。
import os
import urllib
import traceback
import time
import sys
from rknn.api import RKNN
ONNX_MODEL = 'best.onnx'
RKNN_MODEL = 'yolov5s.rknn'
DATASET = './dataset.txt'
QUANTIZE_ON = True
BOX_THRESH = 0.5
NMS_THRESH = 0.6
IMG_SIZE = 640
if __name__ == '__main__':
# Create RKNN object
rknn = RKNN()
if not os.path.exists(ONNX_MODEL):
print('model not exist')
exit(-1)
# pre-process config
print('--> Config model')
rknn.config(reorder_channel='0 1 2',
mean_values=[[0, 0, 0]],
std_values=[[255, 255, 255]],
optimization_level=3,
target_platform='rv1126',
output_optimize=1,
quantize_input_node=QUANTIZE_ON)
print('done')
# Load ONNX model
print('--> Loading model')
ret = rknn.load_onnx(model=ONNX_MODEL,outputs=['output', '405', '418'])
if ret != 0:
print('Load yolov5 failed!')
exit(ret)
print('done')
# Build model
print('--> Building model')
ret = rknn.build(do_quantization=False, dataset=DATASET) #do_quantization=QUANTIZE_ON
if ret != 0:
print('Build yolov5 failed!')
exit(ret)
print('done')
# Export RKNN model
print('--> Export RKNN model')
ret = rknn.export_rknn(RKNN_MODEL)
if ret != 0:
print('Export yolov5rknn failed!')
exit(ret)
print('done')
rknn.release()
注：先使用Netron查看已经转化好的onnx模型结构，看到OUTPUTS的值。从而修改上述代码中rknn.load_onnx的参数。如下图：
环境的配置有需要的可以参考这篇博客：RKNN-toolkit-1.7.1 安装踩坑记录(安装成功)
3、部署
将转化好的rknn模型上传到开发板上，官方的代码有点bug，可以参考一个大佬的博客：35、ubuntu20.04搭建瑞芯微的npu仿真环境和测试rv1126的Debain系统下的yolov5+npu检测功能以及RKNN推理部署，在第七步左右。
到这里部署就基本完成了，只是在第一次在板子上检测图片的时候发现预测结果没有框，后来发现是给的训练数据集太少了，加上模型经过量化，精度下降，所以没有得出理想结果。
解决办法：扩大数据集的规模，用更多的图片进行训练，提高精度。
最后附上一张检测图片：
