RK3588使用npu运行onnx模型推理

文章目录

  • 前言
  • 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,理论上我们需要创建一个这样的环境。
RK3588使用npu运行onnx模型推理_第1张图片
所以我们需要用以下语句创建一个名为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

RK3588使用npu运行onnx模型推理_第2张图片

然后下载【rknn-toolkit2】
RK3588使用npu运行onnx模型推理_第3张图片将下载得到的文件放到一个文件夹下,解压
RK3588使用npu运行onnx模型推理_第4张图片
然后在此文件夹下打开控制台,接着激活使用前面创建的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

RK3588使用npu运行onnx模型推理_第5张图片

安装时,会自动下载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下面。RK3588使用npu运行onnx模型推理_第6张图片假如能够顺利跑完,那应该会生成一个yolov5s_relu.rknn。

RK3588使用npu运行onnx模型推理_第7张图片假如我们想转换自己的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()

把模型放到与该脚本同一个文件夹下,修改好脚本中的 platform = ‘rk3588’ 、onnxModel = 'yolox_s.onnx’为你实际的需求。然后执行,就可以得到onnx模型对应的rcnn模型了。
RK3588使用npu运行onnx模型推理_第8张图片

3.通过rknpu2加载rknn模型执行推理

有空再写


参考:
【Miniconda】
【RK3588模型推理总结】
【RKNN-ToolKit2 1.5.0安装教程】

你可能感兴趣的:(深度学习,rknn)