1、模型转换
import onnx
import numpy as np
import tvm
import tvm.relay as relay
onnx_model = onnx.load('test.onnx')
target = tvm.target.create('llvm')
input_name = '0' # change '1' to '0'
shape_dict = {input_name: (1, 3, 224, 224)}
sym, params = relay.frontend.from_onnx(onnx_model, shape_dict)
with relay.build_config(opt_level=2):
graph, lib, params = relay.build_module.build(sym, target, params=params)
dtype = 'float32'
from tvm.contrib import graph_runtime
print("Output model files")
libpath = "./test.so"
lib.export_library(libpath)
graph_json_path = "./test.json"
with open(graph_json_path, 'w') as fo:
fo.write(graph)
param_path = "./test.params"
with open(param_path, 'wb') as fo:
fo.write(relay.save_param_dict(params))
2、模型部署
import numpy as np
import tvm
import tvm.relay as relay
from tvm.contrib import graph_runtime
import cv2 as cv
test_json = 'test.json'
test_lib = 'test.so'
test_param = 'test.params'
loaded_json = open(test_json).read()
loaded_lib = tvm.module.load(test_lib)
loaded_params = bytearray(open(test_param, "rb").read())
def preprocess(img_src):
img_src= cv.cvtColor(img_src, cv.COLOR_BGR2RGB)
img_src= cv.resize(img_src, (224, 224))
input_data = np.array(img_src).astype(np.float32)
input_data = input_data / 255.0
input_data = np.transpose(input_data, (2, 0, 1))
input_data[0] = (input_data[0] - 0.485)/ 0.229
input_data[1] = (input_data[1] - 0.456)/ 0.224
input_data[2] = (input_data[2] - 0.406)/ 0.225
input_data = input_data[np.newaxis, :].copy()
return input_data
img = cv.imread("29.jpg")
img_input = preprocess(img)
ctx = tvm.cpu(0)
module = graph_runtime.create(loaded_json, loaded_lib, ctx)
module.load_params(loaded_params)
# run the module
module.set_input("0", img_input)
module.run()
out_deploy = module.get_output(0).asnumpy()
print(classes[np.argmax(out_deploy)])
3、遇到问题
转Mobilenet-SSD的onnx模型时遇到问题:
/* an internal invariant was violated while typechecking your program [23:48:07]
/home/mirror/workspace/tvm/src/relay/op/tensor/transform.cc:204: Check failed:
e_dtype == dtype (int64 vs. int32) : relay.concatenate requires all tensors have
the same dtype; */
在讨论区找到一个讨论帖子:
https://discuss.tvm.ai/t/relay-onnx-load-resnet-onnx-to-relay-failed/2411
尝试使用onnx-simplifier工具:
git 地址:
https://github.com/daquexian/onnx-simplifier.git
安装使用:
>> pip3 install onnx-simplifier
>> python3 -m onnxsim input_model output_model
然后再进行模型载入编译就搞定了,感谢大佬们提供的工具~
参考资料:
https://docs.tvm.ai/tutorials/frontend/from_onnx.html#sphx-glr-tutorials-frontend-from-onnx-py