image.png
    如上图所示：
    如果有GPU，可以选择cuda
    target = ‘cuda’
    ctx=tvm.gpu(0)

    如果没有GPU，也可以用CPU, llvm ；
    target = ‘llvm’
    ctx=tvm.cpu(0)

    或是自定义的硬件
    target = ‘llvm’
    ctx=tvm.vacc(0)

    详细请参考tutorials/relay_quick_start.py

    1. import numpy as np
    2. from tvm import relay
    3. from tvm.relay import testing
    4. import tvm
    5. from tvm.contrib import graph_runtime
    7. ######################################################################
    8. # Define Neural Network in Relay
    9. batch_size = 1
    10. num_class = 1000
    11. image_shape = (3, 224, 224)
    12. data_shape = (batch_size,) + image_shape
    13. out_shape = (batch_size, num_class)
    15. mod, params = relay.testing.resnet.get_workload(
    16.     num_layers=18, batch_size=batch_size, image_shape=image_shape)
    18. # set show_meta_data=True if you want to show meta data
    19. print(mod.astext(show_meta_data=False))
    21. ######################################################################
    22. # Compilation
    23. opt_level = 3
    24. target = tvm.target.cuda()
    25. with relay.build_config(opt_level=opt_level):
    26.     graph, lib, params = relay.build_module.build(
    27.         mod, target, params=params)
    29. #####################################################################
    30. # Run the generate library
    32. # create random input
    33. ctx = tvm.gpu()
    34. data = np.random.uniform(-1, 1, size=data_shape).astype("float32")
    35. # create module
    36. module = graph_runtime.create(graph, lib, ctx)
    37. # set input and parameters
    38. module.set_input("data", data)
    39. module.set_input(**params)
    40. # run
    41. module.run()
    42. # get output
    43. out = module.get_output(0, tvm.nd.empty(out_shape)).asnumpy()
    45. # Print first 10 elements of output
    46. print(out.flatten()[0:10])
    48. ######################################################################
    49. # Save and Load Compiled Module
    50. # -----------------------------
    51. # We can also save the graph, lib and parameters into files and load them
    52. # back in deploy environment.
    54. ####################################################
    56. # save the graph, lib and params into separate files
    57. from tvm.contrib import util
    59. temp = util.tempdir()
    60. path_lib = temp.relpath("deploy_lib.tar")
    61. lib.export_library(path_lib)
    62. with open(temp.relpath("deploy_graph.json"), "w") as fo:
    63.     fo.write(graph)
    64. with open(temp.relpath("deploy_param.params"), "wb") as fo:
    65.     fo.write(relay.save_param_dict(params))
    66. print(temp.listdir())
    68. ####################################################
    70. # load the module back.
    71. loaded_json = open(temp.relpath("deploy_graph.json")).read()
    72. loaded_lib = tvm.module.load(path_lib)
    73. loaded_params = bytearray(open(temp.relpath("deploy_param.params"), "rb").read())
    74. input_data = tvm.nd.array(np.random.uniform(size=data_shape).astype("float32"))
    76. module = graph_runtime.create(loaded_json, loaded_lib, ctx)
    77. module.load_params(loaded_params)
    78. module.run(data=input_data)
    79. out_deploy = module.get_output(0).asnumpy()
    81. # Print first 10 elements of output
    82. print(out_deploy.flatten()[0:10])
    84. # check whether the output from deployed module is consistent with original one
    85. tvm.testing.assert_allclose(out_deploy, out, atol=1e-3)