rknn-toolkit2
rknn-toolkit2 copied to clipboard
rockchip-linux
新版本toolkit不支持onnx中lstm操作
尝试用onnx搭建了一个简单的lstm层，转换也有问题，代码如下： import onnx from onnx import helper from onnx import TensorProto import numpy as np from onnxsim import simplify
node_def = onnx.helper.make_node( "LSTM", inputs=["X", "W", "R"], outputs=["Y", "Y_h"], hidden_size=128, )
X = helper.make_tensor_value_info('X', TensorProto.FLOAT, [1, 1, 128]) weight = np.random.ranf(1512128).astype(np.float32) W = helper.make_tensor('W', TensorProto.FLOAT, [1, 512, 128], weight) R = helper.make_tensor('R', TensorProto.FLOAT, [1, 512, 128], weight) Y = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [1, 1, 1, 128]) Y_h = helper.make_tensor_value_info('Y_h', TensorProto.FLOAT, [1, 1, 128])
graph_def = helper.make_graph( [node_def], 'test_conv_mode', [X], # graph inputs [Y, Y_h], # graph outputs initializer=[W, R], )
mode_def = helper.make_model(graph_def, opset_imports=[helper.make_opsetid("", 12)])
onnx.checker.check_model(mode_def)
model_simp, check = simplify(mode_def) onnx.save(model_simp, "lstm.onnx")
模型转换脚本如下：
import os import urllib import traceback import time import sys import numpy as np import cv2 from rknn.api import RKNN
ONNX_MODEL = 'Conv.onnx'
RKNN_MODEL = 'model_p1.rknn'
QUANTIZE_ON = False
if name == 'main':
rknn = RKNN(verbose=True)
print('--> Config model')
rknn.config()
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)
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')
# Init runtime environment
print('--> Init runtime environment')
ret = rknn.init_runtime()
# ret = rknn.init_runtime('rk3566')
if ret != 0:
print('Init runtime environment failed!')
exit(ret)
print('done')
rknn.release()
