模型部署之——ONNX模型转RKNN

提示：这里可以添加学习目标

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

文章目录

一、加载Docker镜像
二、转换脚本

一、加载Docker镜像

加载rknn官方提供的基于x86架构下模型转换的镜像文件，生成容器，以及执行镜像。

sudo docker load -i rknn-toolkit2:1.3.0-cp36
sudo docker run -v `pwd`/rknn_model_convert:/data -it rknn-toolkit2:1.3.0-cp36 /bin/bash  # 将文件路径rknn_model_convert绑定在docker容器的data文件夹下
docker exec -it de0f9e94348c /bin/bash         #de0f9e94348c 为加载镜像生成容器的id

二、转换脚本

from rknn.api import RKNN
import cv2def export_rknn_inference(img, model_path, Dataset, rknn_path):# Create RKNN object# 只在屏幕打印详细的日志信息 # rknn = RKNN(verbose=True)rknn = RKNN(verbose=True)# pre-process configprint('--> Config model')# mean_values 通道均值# std_values 方差, rknn是除以方差# quant_img_RGB2BGR 该参数是将量化图片格式又RGB转换为BGR，通常caffe训练的模型需要这个操作# quantized_algorithm 量化算法，normal 和 mmse, 不写该参数默认值为 normal, 其中：normal量化速度快, mmse量化速度快,精度稍微比normal保持的好# quantized_method 量化方法 channel, layer可选; layer:每层的 weight 只有一套量化参数; channel:每层的 weight 的每个通道都有一套量化参数。默认使用channel# target_platform 可以用来配置不同的芯片, 目前支持 rk3566、rk3568、rk3588、rv1103、rv1106, 该参数的值大小写不敏感。#rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], quantized_algorithm='normal', quantized_method='channel', target_platform='rk3566')rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588')print('done')# Load ONNX modelprint('--> Loading model')# 从当前目录加载 mobilenet_v2 的 onnx 模型，outputs 为 onnx输出层的名字（outputs可不写）#ret = rknn.load_onnx(model=model_path, outputs=['output1', 'output2', 'output3'])ret = rknn.load_onnx(model=model_path)if ret != 0:print('Load model failed!')exit(ret)print('done')# Build modelprint('--> Building model')# 构建 RKNN 模型# do_quantization 是否做量化(不做量化为float16)# dataset 为量化图片的路径# rknn_batch_size 为 batch_size 默认值为1（可以不写），建议 batch_size 小于 32ret = rknn.build(do_quantization=True, dataset=Dataset, rknn_batch_size=1)if ret != 0:print('Build model failed!')exit(ret)print('done')# Export RKNN modelprint('--> Export rknn model')ret = rknn.export_rknn(rknn_path)if ret != 0:print('Export rknn model failed!')exit(ret)print('done')# Init runtime environmentprint('--> Init runtime environment')ret = rknn.init_runtime(target=None, device_id=None, perf_debug=True)# ret = rknn.init_runtime(target='rk3566')if ret != 0:print('Init runtime environment failed!')exit(ret)print('done')# Inferenceprint('--> Running model')outputs = rknn.inference(inputs=[img])rknn.release()print('done')return outputsif __name__ == '__main__':print('This is main ....')# Set inputsimg_path = '20231116_paper_1042005.jpg'model_input_w = 640model_input_h = 480model_path = './yolox.onnx'Dataset = './test_export_1.txt'rknn_path = './yolox.rknn'origimg = cv2.imread(img_path)origimg = cv2.cvtColor(origimg, cv2.COLOR_BGR2RGB)img = cv2.resize(origimg, (model_input_w , model_input_h ))outputs = export_rknn_inference(img, model_path, Dataset, rknn_path)print("outputs:",outputs)