电子工程世界-论坛»论坛 › 电子技术交流 › 嵌入式系统 › #AI挑战营第二站#Ubuntu22上将ONNX模型转换成RKNN模型

阅 4036|回 0

244 帖子	0 TA的资源

纯净的硅（初级）

楼主

发表于2024-4-29 22:33 只看该作者

#AI挑战营第二站#Ubuntu22上将ONNX模型转换成RKNN模型 [复制链接]

本帖最后由 qiao--- 于 2024-4-30 17:18 编辑

1.转换环境搭建

我这里的Linux环境是Ubuntu22，如下

我选择的转换方式是用RKNN-Toolkit2 Conda 环境。

首先下载conda的安装包，用下面的命令

 wget https://mirrors.tuna.tsinghua.edu.cn/anaconda/miniconda/Miniconda3-4.6.14-Linux-x86_64.sh

然后安装miniconda

chmod 777 Miniconda3-4.6.14-Linux-x86_64.sh
bash Miniconda3-4.6.14-Linux-x86_64.sh

安装的时候一路按enter键就可以了。

创建 RKNN-Toolkit2 开发 Conda 环境，-n参数表示环境名称，指定 python 版本为3.8（建议版本）

conda create -n RKNN-Toolkit2 python=3.8

我这里已经创建好了，如下

然后获取 RKNN-Toolkit2 安装包，RKNN-Toolkit2 工具在 PC 平台上提供 C 或 Python 接口，简化模型的部署和运行。用户可以通过该工具轻松完成以下功能：模型转换、量化、推理、性能和内存评估、量化精度分析以及模型加密。RKNN 软件栈可以帮助用户快速的将 AI 模型部署到 Rockchip 芯片。

git clone https://github.com/rockchip-linux/rknn-toolkit2.git

接着安装 RKNN-Toolkit2 相关的依赖库，cp38为对应的 Conda 环境 python 版本，实验使用的版本为 3.8 所以使用后缀为cp38的依赖项

pip install tf-estimator-nightly==2.8.0.dev2021122109 
pip install -r rknn-toolkit2/packages/requirements_cp38-1.6.0.txt -i https://pypi.mirrors.ustc.edu.cn/simple/

安装 RKNN-Toolkit2

pip install rknn-toolkit2/packages/rknn_toolkit2-1.6.0+81f21f4d-cp38-cp38-linux_x86_64.whl

安装好后，可以验证是否安装成功，如果没有报错说明安装成功，我这里是安装成功了的如下

下载 rknn_model_zoo,这是瑞芯微官方的用于转换模型的工具，里面有很多转换示例，转换代码我们可以借鉴一下这里面的

git clone https://github.com/airockchip/rknn_model_zoo.git

2.转换模型

前面我们已经将基础环境搭建好了，接下来可以来转换模型了。

我们找到我们训练营第一站训练出来的ONNX模型，将他拷贝到我们自己创建的文件夹myModel内，如下所示

紧接着我们新建一个convert.py文件里面使我们的转换代码，这个可以参考examples文件下的内容

from rknn.api import RKNN

# Create RKNN object
rknn = RKNN(verbose=True)

# pre-process config
print('--> config model')
rknn.config(target_platform='rv1106', mean_values=[[28]], std_values=[[28]])
print('done')

# Load model
print('--> Loading model')
ret = rknn.load_onnx(model='./mnist.onnx')
if ret != 0:
    print('Load model failed!')
    exit(ret)
print('done')


rknn.build(do_quantization=True, dataset='./data.txt')  # 构建RKNN模型，可选参数量化
if ret != 0:
    print('Build model failed!')
    exit(ret)
print('done')

ret = rknn.export_rknn('./mnist.rknn')  # 导出RKNN模型文件
if ret != 0:
    print('Export rknn model failed!')
    exit(ret)
print('done')

# 释放 RKNN 对象
rknn.release()

然后新建一个pic文件夹，里面装我们的dateset，如下