opt_level = 3
target = tvm.target.cuda()
with tvm.transform.PassContext(opt_level=opt_level):
graph, lib, params = relay.build(mod, target, params=params)
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Run the generate library
// ------------------------
// Now we can create graph runtime and run the module on Nvidia GPU.
// create random input
ctx = tvm.gpu()
data = np.random.uniform(-1, 1, size=data_shape).astype("float32")
// create module
module = graph_runtime.create(graph, lib, ctx)
// set input and parameters
module.set_input("data", data)
module.set_input(**params)
// run
module.run()
// get output
out = module.get_output(0, tvm.nd.empty(out_shape)).asnumpy()
// Print first 10 elements of output
print(out.flatten()[0:10])
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Save and Load Compiled Module
// -----------------------------
// We can also save the graph, lib and parameters into files and load them
// back in deploy environment.
////////////////////////////////////////////////////////////////////////////////////////////////////////
// save the graph, lib and params into separate files
from tvm.contrib import util
temp = util.tempdir()
path_lib = temp.relpath("deploy_lib.tar")
lib.export_library(path_lib)
with open(temp.relpath("deploy_graph.json"), "w") as fo:
fo.write(graph)
with open(temp.relpath("deploy_param.params"), "wb") as fo:
fo.write(relay.save_param_dict(params))
print(temp.listdir())
////////////////////////////////////////////////////////////////////////////////////////////////////////
After Change
ctx = tvm.gpu()
data = np.random.uniform(-1, 1, size=data_shape).astype("float32")
// create module
module = graph_runtime.GraphModule(lib["default"](ctx))
// set input and parameters
module.set_input("data", data)
// run
module.run()
