#AI挑战营第二站# Windows下ONNX转RKNN模型过程
[复制链接]
本帖最后由 MioChan 于 2024-4-26 23:53 编辑
第一步:环境配置
在rknn-toolkit git仓库: 的release里下载1.7.5压缩包,解压后能看见多个平台的whl文件,这里因为我用的是win所以选择
rknn_toolkit-1.7.5-cp36-cp36m-win_amd64.whl这个文件。
另外这个库并不支持高版本的python,建议还是新建一个python 3.6 的conda虚拟环境,另外这个库有各种依赖,pytorch、tensorflow之类的都得装才能正常运行,而且版本不能太高。
conda create --name=RKNN python=3.6.13
激活环境
conda activate RKNN
升级一下pip顺便换个清华源
python -m pip install --upgrade
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
手动安装一下这几个库,其他的在安装rknn-toolkit会自动安装
pip install scikit-build
pip install opencv-python==4.3.0.38
pip install torch==1.10.1 torchversion==0.11.2
pip install tensorflow==1.15.5
pip install matplotlib
安装之前下好的whl文件
pip install rknn_toolkit-1.7.5-cp36-cp36m-win_amd64.whl
注意如果命令行的当前路径不是whl所在的文件夹需要在whl文件名前带上其路径,像下面这样
pip install E:\rknn_toolkit-1.7.0-cp36-cp36m-win_amd64.whl
正常来说就可以顺利运行了,但如果是缺少其他库,自己看报错用pip install或conda install缺啥装啥就行了
第二步:编写代码完成转换
from rknn.api import RKNN
# 创建 RKNN 对象
rknn = RKNN()
# 加载 ONNX 模型
ret = rknn.load_onnx(model='mnist_cnn.onnx')
# 模型配置
rknn.config(mean_values=[[128]], std_values=[[128]]) #图片归一化到【-1,1】
# 模型转换
ret = rknn.build(do_quantization=False) # 这里设置是否量化
# 保存 RKNN 模型
ret = rknn.export_rknn('mnist_cnn.rknn')
# 释放 RKNN 对象
rknn.release()
另外,ONNX版本不能超过11,否则转换会报错 ,在之前任务中的导出onnx的代码带上版本设置参数即可
# 导出模型到ONNX格式
torch.onnx.export(model, dummy_input, 'mnist_cnn.onnx',
input_names=['input'], output_names=['output'], # 输入和输出节点名称
do_constant_folding=True, # 是否执行常量折叠优化
export_params=True, # 是否包含权重
opset_version=10, #设置版本
verbose=True) # 是否打印转换细节
|