RK3588使用npu运行onnx模型推理

news2024/10/7 7:23:03

文章目录

前言
1.安装rknn-toolkit2
2.将onnx模型转为rknn模型
3.通过rknpu2加载rknn模型执行推理

前言

先说总体流程：
1.在x86主机上安装【rknn-toolkit2】
2.使用【rknn-toolkit2】，把onnx模型转换成rknn格式的模型。
3.在板子部署【rknpu2】，在编程时，调用其提供的接口即可。

1.安装rknn-toolkit2

最好先按照这里【ubuntu安装Miniconda】装个miniconda，注意要选择合适的版本。
conda的好处是可以创建独立的环境，防止环境之间污染。
创建虚拟环境之前，我们先看看rknn-toolkit2对环境的要求是啥。
在rknn-toolkit2的官网看到，在Ubuntu 22.04系统下，使用的是python 3.10，理论上我们需要创建一个这样的环境。
在这里插入图片描述
所以我们需要用以下语句创建一个名为rknn-toolkit2的虚拟环境

# conda create -n rknn-toolkit2 python=3.10
# conda create -n rknn-toolkit2 python=3.8
conda create -n rknn-toolkit2 python=3.6

经过实际测试，不能使用3.10、3.8的python，否则在安装到最后时会报类似下面的错误：

ERROR: Ignored the following versions that require a different python version: 16.2 Reguires-Python =3,7,33,10: 1.6,3 ReguiresPython 2=3,7 ,<3,10: 1.7.0 Reguires-Python >=3.7,<3.10; 1.7.1 Requires-Python >=3.7,<3.10ERROR: Could not find a version that satisfies the reguirement tf-estimator-nightly==2.8.0.dev2021122109 (from tensorflow) (from versions: none)ERROR: No matching distribution found for tf-estimator-nightly==2.8.0.dev2021122109

在这里插入图片描述

然后下载【rknn-toolkit2】
在这里插入图片描述将下载得到的文件放到一个文件夹下，解压

然后在此文件夹下打开控制台，接着激活使用前面创建的conda环境。

conda activate rknn-toolkit2

从packages文件夹中选择合适的文件进行install，我这里是python 3.6的，所以我选择了rknn_toolkit2-1.5.2+b642f30c-cp36-cp36m-linux_x86_64.whl （下面命令中的 -i https://pypi.tuna.tsinghua.edu.cn/simple，是指定源，速度快很多）

cd rknn-toolkit2-1.5.2/packages
pip install rknn_toolkit2-1.5.2+b642f30c-cp36-cp36m-linux_x86_64.whl -i https://pypi.tuna.tsinghua.edu.cn/simple

在这里插入图片描述

安装时，会自动下载tensorflow、torch等组件，会占用较大空间，最好要先预留2G以上的空间。
安装起来慢的话，可以分别试试这几个源：

清华大学 https://pypi.tuna.tsinghua.edu.cn/simple/
阿里云 http://mirrors.aliyun.com/pypi/simple/
中国科技大学 https://pypi.mirrors.ustc.edu.cn/simple/
豆瓣(douban) http://pypi.douban.com/simple/
中国科学技术大学 http://pypi.mirrors.ustc.edu.cn/simple/

安装完成之后，执行一下下面的命令，检验是否已经安装好了。假如没报错的话，就是已经成功安装了。

python
from rknn.api import RKNN

2.将onnx模型转为rknn模型

可以先跑一下自带的例程。在 rknn-toolkit2-1.5.2/examples/onnx/yolov5下面。在这里插入图片描述假如能够顺利跑完，那应该会生成一个yolov5s_relu.rknn。

在这里插入图片描述假如我们想转换自己的onnx模型，可以直接利用下面这个代码（代码来源）：

from rknn.api import RKNN
import os

if __name__ == '__main__':
    platform = 'rk3588'
    onnxModel = 'yolox_s.onnx'

    '''step 1: create RKNN object'''
    rknn = RKNN()

    '''step 2: load the .onnx model'''
    rknn.config(target_platform=platform, optimization_level = 2)
    print('--> Loading model')
    ret = rknn.load_onnx(onnxModel)
    if ret != 0:
        print('load model failed')
        exit(ret)
    print('done')

    '''step 3: building model'''
    print('-->Building model')
    ret = rknn.build(do_quantization=False)
    if ret != 0:
        print('build model failed')
        exit()
    print('done')

    '''step 4: export and save the .rknn model'''
    OUT_DIR = 'rknn_models'
    RKNN_MODEL_PATH = './{}/{}.rknn'.format(OUT_DIR, onnxModel)
    if not os.path.exists(OUT_DIR):
        os.mkdir(OUT_DIR)
    print('--> Export RKNN model: {}'.format(RKNN_MODEL_PATH))
    ret = rknn.export_rknn(RKNN_MODEL_PATH)
    if ret != 0:
        print('Export rknn model failed.')
        exit(ret)
    print('done')

    '''step 5: release the model'''
    rknn.release()