The model can be converted to Relay and compiled very well.
However, when use the statement module.set_input('input_1', x_test.astype(dtype))
, I will crash.
How do I set the first parameter of module.set_input()
?
Any comments are welcome. Thanks in advance.
My script is as follows:
import keras
from keras.utils import to_categorical
import tvm
import tvm.relay as relay
from tvm.contrib import graph_runtime
import numpy as np
import tensorflow as tf
try:
tf_compat_v1 = tf.compat.v1
except ImportError:
tf_compat_v1 = tf
def get_data(size):
_, (x_test, y_test) = keras.datasets.cifar100.load_data()
x_test = x_test.astype('float32') / 255.0
w, h = 32, 32
x_test = x_test.reshape(x_test.shape[0], w, h, 3)
x_test, y_test = x_test[:size], y_test[:size]
y_test = y_test.transpose([1, 0])[0]
y_test = np.array(y_test)
y_test = to_categorical(y_test)
return x_test, y_test
model_path = "resnet20-cifar100_origin.pb"
with tf_compat_v1.gfile.GFile(model_path, "rb") as f:
graph_def = tf_compat_v1.GraphDef()
graph_def.ParseFromString(f.read())
batch_size = 10
x_test, y_test = get_data(batch_size)
input_shape = (batch_size, 32, 32, 3)
output_shape = (batch_size, y_test[0].size)
shape_dict = {'input_1': input_shape}
irmod, params = relay.frontend.from_tensorflow(graph_def, layout='NHWC', shape=shape_dict)
target = 'llvm'
ctx = tvm.cpu(0)
with tvm.transform.PassContext(opt_level=3):
graph, lib, params = relay.build_module.build(
irmod, target=target, params=params)
module = graph_runtime.create(graph, lib, ctx)
dtype = 'float32'
module.set_input('input_1', x_test.astype(dtype))
module.set_input(**params) # Crash here!!!
module.run()
res_tvm = module.get_output(0, tvm.nd.empty(output_shape)).asnumpy()
Part of the model info: