[原创] #AI挑战营第二站#RV1106手写MNIST模型onnx转换rknn

环境部署不再赘述，可以参考之前写的帖子

这次转换MNIST模型，还是一个思路，参考rknn-toolkit2中的yolov5的转换脚本，如rknn-toolkit2/rknn-toolkit2/examples/onnx/yolov5/test.py

1. 删掉所有推理测试的代码

2. 修改dataset.txt中的图片名称

3. IMG_PATH为dataset.txt处理前的图片，使用cv库进行灰度及resize操作

4. ONNX_MODEL为输入的onnx模型

5. RKNN_MODEL为输出的rknn模型

import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN

# Model from https://github.com/airockchip/rknn_model_zoo
ONNX_MODEL = 'mnist.onnx'
RKNN_MODEL = 'mnist.rknn'
IMG_PATH = './num_pre.jpg'
DATASET = './dataset.txt'

IMG_SIZE = 28

QUANTIZE_ON = True

if __name__ == '__main__':

    img = cv2.imread(IMG_PATH)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (IMG_SIZE, IMG_SIZE))
    cv2.imwrite('num.jpg', img)

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

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

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

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')

    # Export RKNN model
    print('--> Export rknn model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export rknn model failed!')
        exit(ret)
    print('done')

    rknn.release()

python3 test.py即可得到rknn模型

 

mnist.rknn (58.38 KB)
(下载次数: 0, 2024-6-5 22:55 上传)