tvm学习笔记 （三）：载入onnx格式模型

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