#AI挑战营第二站# onnx转rknn

慕容雪花

#AI挑战营第二站# onnx转rknn [复制链接]

本帖最后由慕容雪花于 2024-6-12 13:06 编辑

#AI挑战营第一站# 基于PyTorch，在PC上完成MNIST手写数字识别模型训练 https://bbs.eeworld.com.cn/thread-1278408-1-1.html

为了使用 RKNPU，用户需要首先在计算机上运行 RKNN-Toolkit2 工具，将训练好的模型转换为 RKNN 格式模型。RKNN-Toolkit2 工具在 PC 平台上提供 C 或 Python 接口，简化模型的部署和运行。用户可以通过该工具轻松完成以下功能：模型转换、量化、推理、性能和内存评估、量化精度分析以及模型加密。RKNN 软件栈可以帮助用户快速的将 AI 模型部署到 Rockchip 芯片。整体的框架如下

RV1106 需要使用rknn-toolkit2：

git clone https://github.com/airockchip/rknn-toolkit2

接下来根据Python版本选择合适的whl：

pip install ./rknn-toolkit2/rknn-toolkit2/packages/rknn_toolkit2-2.0.0b0+9bab5682-cp38-cp38-linux_x86_64.whl

接下来验证rknn toolkit是否安装成功：

之后可以通过Python脚本来进行模型转换，参考大佬：https://bbs.eeworld.com.cn/thread-1280156-1-1.html

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_101_model.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()

有个步骤是新建一个data.txt，指向其中一个测试数据：

最后运行脚本：

附件：

model.rknn (27.97 KB, 下载次数: 0)

Jacktang

模型转换、量化、推理、性能和内存评估、量化精度分析以及模型加密，实现的功能不少

#AI挑战营第二站# onnx转rknn [复制链接]

最新回复