瑞芯微rv1126上实现yolov5模型部署
第一次接触模型训练和在开发板部署,过程曲折,从开始的一脸懵到最后模型部署成功,查阅了不少资料和学习了不少大佬的经验,在这里记录一下过程和心得。
一、总体思路
首先在本机上使用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。
- 安装labelimg
修改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推理部署,在第七步左右。
到这里部署就基本完成了,只是在第一次在板子上检测图片的时候发现预测结果没有框,后来发现是给的训练数据集太少了,加上模型经过量化,精度下降,所以没有得出理想结果。
解决办法:扩大数据集的规模,用更多的图片进行训练,提高精度。
最后附上一张检测图片:
- 分享
- 举报
-
浏览量:1222次2024-02-05 10:41:25
-
浏览量:1444次2024-01-22 17:46:51
-
浏览量:7748次2022-07-29 18:09:52
-
浏览量:494次2024-02-18 14:24:39
-
浏览量:1794次2024-01-05 17:00:10
-
浏览量:3338次2022-11-22 10:37:11
-
浏览量:3271次2022-11-22 09:58:36
-
浏览量:3957次2022-10-31 15:07:27
-
浏览量:2824次2022-11-22 09:26:53
-
浏览量:1189次2024-02-23 17:41:04
-
浏览量:952次2024-05-24 17:11:01
-
浏览量:1001次2024-05-24 17:28:44
-
浏览量:6894次2024-01-04 14:35:55
-
浏览量:1000次2024-05-24 19:05:42
-
浏览量:1020次2024-05-24 18:51:13
-
浏览量:1184次2024-05-24 17:53:56
-
浏览量:1547次2023-08-08 13:59:23
-
浏览量:5859次2022-05-11 15:14:51
-
浏览量:3863次2021-04-02 09:47:41
-
广告/SPAM
-
恶意灌水
-
违规内容
-
文不对题
-
重复发帖
保持微笑
感谢您的打赏,如若您也想被打赏,可前往 发表专栏 哦~
举报类型
- 内容涉黄/赌/毒
- 内容侵权/抄袭
- 政治相关
- 涉嫌广告
- 侮辱谩骂
- 其他
详细说明