yolov5 onnx模型 转为 rknn模型

1、转换为rknn模型环境搭建

onnx模型需要转换为rknn模型才能在rv1126开发板上运行，所以需要先搭建转换环境

模型转换工具 模型转换相关文件下载：

网盘下载链接：百度网盘 请输入提取码  提取码：teuc

将其移动到虚拟机中，找到文件夹中的docker文件 rknn-toolkit-1.7.1-docker.tar.gz 、 model_convert文件夹

加载 模型转换工具docker镜像

docker load --input /home/developer/rknn-toolkit/rknn-toolkit-1.7.1-docker.tar.gz

进入 镜像bash环境

 执行以下指令把工作区域映射进docker镜像，其中/home/developer/rknn-toolkit/model_convert为工作区域，/test为映射到docker镜像，/dev/bus/usb:/dev/bus/usb为映射usb到docker镜像：

docker run -t -i --privileged -v /dev/bus/usb:/dev/bus/usb -v /home/developer/rknn-toolkit/model_convert:/test rknn-toolkit:1.7.1 /bin/bash

两处文件映射，即同步

 2、生成量化图片列表

这一步是把准备好的一些图片，生成图片路径的文本文件，在构建RKNN模型的时候有用处。通过使用真实的样本数据集，RKNN工具可以更好地理解和建模模型的输入数据，从而更好地优化网络结构、权重和量化方案。

在docker环境切换到模型转换工作目录；执行 gen_list.py，会得到一个文本文件pic_path.txt，里面是图片的路径:

cd /test/coco_object_detect
python gen_list.py

gen_list.py内容如下：

import os
import randomdef main(image_dir):save_image_txt = './pic_path.txt'save_val_number = 0img_path_list = []image_list = os.listdir(image_dir)for i in image_list:#if os.path.isdir(image_dir):#print("i:", i)image_path = image_dir + '/' + i#print("image_path:", image_path)img_path_list.append(image_path)#print(img_path_list)print('len of all', len(img_path_list))random.shuffle(img_path_list)with open(save_image_txt, 'w') as F:for i in range(len(img_path_list)):F.write(img_path_list[i]+'\n')if __name__ == '__main__':image_dir = '/test/quant_dataset/coco_data'  # 图片所在路径，大概500张main(image_dir)

3、onnx模型转换为rknn模型

还是在docker环境 模型转换工作目录，运行rknn_convert.py

python rknn_convert.py

这一步如果是在虚拟机上运行的话，8GB的内存条win10系统也要用，分配给虚拟机的没多少，3GB也不够执行这一步。

后来我直接在Ubuntu系统执行这一步，8GB系统用一点还剩6.7GB，CPU和内存直接干满

 

rknn_convert.py 源码：

import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNNONNX_MODEL = 'best.onnx' # onnx 模型的路径
RKNN_MODEL = './yolov5_mask_rv1126.rknn'  # 转换后的 RKNN 模型保存路径
DATASET = './pic_path.txt'   # 数据集文件路径QUANTIZE_ON = True   # 是否进行量化if __name__ == '__main__':# 创建 RKNN 对象rknn = RKNN(verbose=True)# 检查 ONNX 模型文件是否存在if not os.path.exists(ONNX_MODEL):print('model not exist')exit(-1)# 配置模型预处理参数print('--> Config model')rknn.config(reorder_channel='0 1 2', # 表示 RGB 通道mean_values=[[0, 0, 0]], # 每个通道的像素均值，预处理时对应通道减去该值std_values=[[255, 255, 255]], # 每个通道的像素标准差，每个通道除以该值optimization_level=3, # 优化级别target_platform = 'rv1126', #指定目标平台为rv1126output_optimize=1,      # 输出优化为真quantize_input_node=QUANTIZE_ON)  # 对时输入节点进行量化print('done')# 加载 ONNX 模型print('--> Loading model')ret = rknn.load_onnx(model=ONNX_MODEL)if ret != 0:print('Load yolov5 failed!')exit(ret)print('done')# 构建模型print('--> Building model')ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)if ret != 0:print('Build yolov5 failed!')exit(ret)print('done')# 导出 RKNN 模型print('--> Export RKNN model')ret = rknn.export_rknn(RKNN_MODEL)if ret != 0:print('Export yolov5rknn failed!')exit(ret)print('done')