TVM编译模型的基本步骤

image.png

  1. 导入其他框架的模型
    支持的主要框架：Tensorflow, PyTorch, ONNX

  2. 将该模型转换成Relay（TVM的高层级IR）.
    Relay支持特性：

  • 传统数据流式的表示
  • 函数式语言风格的表示

  • 两种风格的混合

    在此步骤中，Relay进行图层次的优化

  1. 将Relay转换成更细粒度的Tensor Expression(TE)
    Relay使用FuseOps将模型划分成小的子图，在此过程中可以使用一些schedule原语进行优化，
    如tiling, vectorization, parallelization, unrolling, and fusion
    TOPI包含一些预定义的常用Operator。

  2. 搜索最佳调度策略（AutoTVM或AutoScheduler）
    TVM提供了如下两种自动调优模块。

  • AutoTVM: 模板化的自动调优模块。

    TOPI中提供了通用算子优化后的调度模板

  • AutoScheduler (a.k.a. Ansor): 无模板的自动调优模块，自动生成搜索空间

  1. 选择模型编译的最优配置。
    自动调优会生成JSON格式的优化记录.

  2. 编译生成Tensor IR (TIR, TVM的低层级IR，相对于Relay）。
    TVM支持的后端包括:

  • LLVM, 通过它可以生成llvm支持的所有硬件如x86, ARM.
  • 特定编译器，如NVCC, NVIDIA的编译器.
  • 通过BYOC(Bring Your Own Codegen)框架实现
  1. 编译生成机器代码。
    TVM可以将模型编译成可链接的对象模块来通过轻量级的运行时来运行。它提供多种语言的支持。
    TVM也支持将模型和运行时统一打包。

Relay的使用

Relay是TVM中定义的High-level IR。
通过TVM中Relay模块，我们可以转换其他框架模型，并生成对应后端代码。

  1. import numpy as np
  2. from tvm import relay
  3. from tvm.relay import testing
  4. import tvm
  5. from tvm import te
  6. from tvm.contrib import graph_executor
  7. import tvm.testing
  9. ######################################################################
  10. # (1) 导入模型：
  11. #     这里我们使用了预编译的resnet_18
  12. batch_size = 1
  13. num_class = 1000
  14. image_shape = (3, 224, 224)
  15. data_shape = (batch_size,) + image_shape
  16. out_shape = (batch_size, num_class)
  18. mod, params = relay.testing.resnet.get_workload(
  19.     num_layers=18, batch_size=batch_size, image_shape=image_shape
  20. )
  21. # set show_meta_data=True if you want to show meta data
  22. print(mod.astext(show_meta_data=False))
  24. ######################################################################
  25. # (2) 编译
  26. # ----------- 
  27. # - 优化等级范围 (0 to 3). 
  28. # - 优化passes: operator fusion, pre-computation, layout transformation...
  29. # Relay进行graph-level优化， TVM进行tensor-level优化
  30. #
  31. opt_level = 3
  32. target = tvm.target.Target(target="llvm", host="llvm") #tvm.target.cuda()
  33. with tvm.transform.PassContext(opt_level=opt_level):
  34.     lib = relay.build(mod, target, params=params)
  36. #####################################################################
  37. #（3） 运行
  38. dev = tvm.device(target.kind.name, 0) #tvm.cuda()
  39. data = np.random.uniform(-1, 1, size=data_shape).astype("float32") 
  40. module = graph_executor.GraphModule(lib["default"](dev)) 
  41. module.set_input("data", data) 
  42. module.run() 
  43. out = module.get_output(0, tvm.nd.empty(out_shape)).numpy()
  45. print(out.flatten()[0:10])
  47. ######################################################################
  48. # （4） 保存并加载模型
  49. from tvm.contrib import utils
  51. temp = utils.tempdir()
  52. path_lib = temp.relpath("deploy_lib.tar")
  53. lib.export_library(path_lib)
  54. print(temp.listdir())
  56. ####################################################
  58. loaded_lib = tvm.runtime.load_module(path_lib)
  59. input_data = tvm.nd.array(data)
  61. module = graph_executor.GraphModule(loaded_lib["default"](dev))
  62. module.run(data=input_data)
  63. out_deploy = module.get_output(0).numpy()
  65. print(out_deploy.flatten()[0:10])
  67. # 比较结果
  68. tvm.testing.assert_allclose(out_deploy, out, atol=1e-5)

TOPI的使用

TOPI 提供了很多实用算子和针对特性后端的调度器。

  2. """ 
  3. TVM Operator Inventory (TOPI)提供了numpy风格的高级抽象算子. 
  4. """
  5. from __future__ import absolute_import, print_function
  7. import tvm
  8. import tvm.testing
  9. from tvm import te
  10. from tvm import topi
  11. import numpy as np
  13. ######################################################################
  14. # 基础实例：按行求和(等价numpy代码:`B = numpy.sum(A, axis=1)`) 
  15. #
  16. n = te.var("n")
  17. m = te.var("m")
  18. A = te.placeholder((n, m), name="A")
  19. k = te.reduce_axis((0, m), "k")
  20. B = te.compute((n,), lambda i: te.sum(A[i, k], axis=k), name="B")
  21. s = te.create_schedule(B.op)
  23. ######################################################################
  24. # 查看生成的IR
  25. print(tvm.lower(s, [A], simple_mode=True))
  27. ######################################################################
  28. # 使用预定义topi方法
  29. C = topi.sum(A, axis=1)
  30. ts = te.create_schedule(C.op)
  31. print(tvm.lower(ts, [A], simple_mode=True))
  33. ######################################################################
  34. # Numpy-style 广播操作
  35. x, y = 100, 10
  36. a = te.placeholder((x, y, y), name="a")
  37. b = te.placeholder((y, y), name="b")
  38. c = a + b  # same as topi.broadcast_add
  39. d = a * b  # same as topi.broadcast_mul
  41. ######################################################################
  42. # Overloaded with the same syntax, TOPI handles broadcasting a primitive (`int`, `float`) to a tensor :code:`d - 3.14`.
  44. ######################################################################
  45. # TOPI定义了针对不同平台的调度器，如CUDA，x86... 
  46. #
  47. e = topi.elemwise_sum([c, d])
  48. f = e / 2.0
  49. g = topi.sum(f)
  50. with tvm.target.Target(target="llvm", host="llvm"):
  51.     # sg = topi.cuda.schedule_reduce(g) 
  52.     sg = topi.x86.schedule_reduce(g)
  53.     print(tvm.lower(sg, [a, b], simple_mode=True))
  55. ######################################################################
  56. # As you can see, scheduled stages of computation have been accumulated and we can examine them by
  57. print(sg.stages)
  59. ######################################################################
  60. # 和numpy运算结果做比较
  61. func = tvm.build(sg, [a, b, g], "llvm")
  62. dev = tvm.device("llvm", 0)
  63. a_np = np.random.uniform(size=(x, y, y)).astype(a.dtype)
  64. b_np = np.random.uniform(size=(y, y)).astype(b.dtype)
  65. g_np = np.sum(np.add(a_np + b_np, a_np * b_np) / 2.0)
  66. a_nd = tvm.nd.array(a_np, dev)
  67. b_nd = tvm.nd.array(b_np, dev)
  68. g_nd = tvm.nd.array(np.zeros(g_np.shape, dtype=g_np.dtype), dev)
  69. func(a_nd, b_nd, g_nd)
  70. tvm.testing.assert_allclose(g_nd.numpy(), g_np, rtol=1e-5)
  72. ######################################################################
  73. # TOPI 提供了常用的神经网络算子，如softmax，conv
  74. ######################################################################
  75. # NOTE: TOPI的实现因后端而异，目标平台和调度器必须一致 
  77. tarray = te.placeholder((512, 512), name="tarray")
  78. softmax_topi = topi.nn.softmax(tarray)
  79. with tvm.target.Target(target="llvm", host="llvm"):
  80.     # sst = topi.cuda.schedule_softmax(softmax_topi)
  81.     sst = topi.x86.schedule_softmax(softmax_topi) 
  82.     print(tvm.lower(sst, [tarray], simple_mode=True))
  85. data = te.placeholder((1, 3, 224, 224))
  86. kernel = te.placeholder((10, 3, 5, 5))
  88. with tvm.target.Target(target="llvm", host="llvm"): 
  89.     conv = topi.nn.conv2d_nchw(data, kernel, 1, 2, 1)
  90.     out = topi.nn.relu(conv)
  91.     # sconv = topi.cuda.schedule_conv2d_nchw([out])
  92.     sconv = topi.x86.schedule_conv2d_nchw([out]) 
  93.     print(tvm.lower(sconv, [data, kernel], simple_mode=True))

模型的编译、运行和优化

  1. """
  2. 使用TVM的python API来进行模型的编译、运行和优化
  3. """
  4. import onnx
  5. from tvm.contrib.download import download_testdata
  6. from PIL import Image
  7. import numpy as np
  8. import tvm.relay as relay
  9. import tvm
  10. from tvm.contrib import graph_executor
  12. ###############################################################################
  13. # 获取一个预训练的模型（使用了一个ONNX格式的模型）和测试图片
  14. model_url = "".join(
  15.     [
  16.         "https://github.com/onnx/models/raw/",
  17.         "master/vision/classification/resnet/model/",
  18.         "resnet50-v2-7.onnx",
  19.     ]
  20. )
  22. model_path = download_testdata(model_url, "resnet50-v2-7.onnx", module="onnx")
  23. onnx_model = onnx.load(model_path)
  25. img_url = "https://s3.amazonaws.com/model-server/inputs/kitten.jpg"
  26. img_path = download_testdata(img_url, "imagenet_cat.png", module="data")
  28. # Resize it to 224x224
  29. resized_image = Image.open(img_path).resize((224, 224))
  30. img_data = np.asarray(resized_image).astype("float32")
  32. # Our input image is in HWC layout while ONNX expects CHW input, so convert the array
  33. img_data = np.transpose(img_data, (2, 0, 1))
  35. # Normalize according to the ImageNet input specification
  36. imagenet_mean = np.array([0.485, 0.456, 0.406]).reshape((3, 1, 1))
  37. imagenet_stddev = np.array([0.229, 0.224, 0.225]).reshape((3, 1, 1))
  38. norm_img_data = (img_data / 255 - imagenet_mean) / imagenet_stddev
  40. # Add the batch dimension, as we are expecting 4-dimensional input: NCHW.
  41. img_data = np.expand_dims(norm_img_data, axis=0)
  43. ###############################################################################
  44. # 使用Relay来编译模型
  46. target = "llvm" 
  47. input_name = "data"
  48. shape_dict = {input_name: img_data.shape}
  50. mod, params = relay.frontend.from_onnx(onnx_model, shape_dict)
  52. with tvm.transform.PassContext(opt_level=3):
  53.     lib = relay.build(mod, target=target, params=params)
  55. dev = tvm.device(str(target), 0)
  56. module = graph_executor.GraphModule(lib["default"](dev))
  58. ######################################################################
  59. # 使用TVM Runtime运行模型
  61. dtype = "float32"
  62. module.set_input(input_name, img_data)
  63. module.run()
  64. output_shape = (1, 1000)
  65. tvm_output = module.get_output(0, tvm.nd.empty(output_shape)).numpy()
  67. #～～～～～～～～～～～
  68. # 评估基本的性能数据
  69. import timeit
  71. timing_number = 10
  72. timing_repeat = 10
  73. unoptimized = (
  74.     np.array(timeit.Timer(lambda: module.run()).repeat(repeat=timing_repeat, number=timing_number))
  75.     * 1000
  76.     / timing_number
  77. )
  78. unoptimized = {
  79.     "mean": np.mean(unoptimized),
  80.     "median": np.median(unoptimized),
  81.     "std": np.std(unoptimized),
  82. }
  84. print(unoptimized)
  86. #～～～～～～～～～～
  87. # 后处理 
  89. from scipy.special import softmax
  91. # Download a list of labels
  92. labels_url = "https://s3.amazonaws.com/onnx-model-zoo/synset.txt"
  93. labels_path = download_testdata(labels_url, "synset.txt", module="data")
  95. with open(labels_path, "r") as f:
  96.     labels = [l.rstrip() for l in f]
  98. # Open the output and read the output tensor
  99. scores = softmax(tvm_output)
  100. scores = np.squeeze(scores)
  101. ranks = np.argsort(scores)[::-1]
  102. for rank in ranks[0:5]:
  103.     print("class='%s' with probability=%f" % (labels[rank], scores[rank]))
  105. #～～～～～～～～～～～
  106. # 输出结果
  107. #
  108. #     # class='n02123045 tabby, tabby cat' with probability=0.610553
  109. #     # class='n02123159 tiger cat' with probability=0.367179
  110. #     # class='n02124075 Egyptian cat' with probability=0.019365
  111. #     # class='n02129604 tiger, Panthera tigris' with probability=0.001273
  112. #     # class='n04040759 radiator' with probability=0.000261
  114. ################################################################################
  115. # 使用autotvm来优化模型
  116. #
  117. import tvm.auto_scheduler as auto_scheduler
  118. from tvm.autotvm.tuner import XGBTuner
  119. from tvm import autotvm
  121. #～～～～～～～～～～～
  122. # 基本配置参数
  123. number = 10 # 尝试的不同优化配置的数目
  124. repeat = 1  # 每个配置重复测量的次数
  125. min_repeat_ms = 0  # 最小运行时间，影响GPU精度调优，CPU设置为0
  126. timeout = 10  # 超时时间
  128. # create a TVM runner
  129. runner = autotvm.LocalRunner(
  130.     number=number,
  131.     repeat=repeat,
  132.     timeout=timeout,
  133.     min_repeat_ms=min_repeat_ms,
  134.     enable_cpu_cache_flush=True,
  135. )
  137. #～～～～～～～～～～～～
  138. # 优化选项
  139. # tuner: 这里使用xgboost作为优化算法
  140. # trials：对于产品级的模型，trials推荐值为CPU 1500, GPU 3000-4000. 这个是模型和处理器相关的。
  141. #         需要多花时间实验。这里作为简单测试，仅设置为10
  142. # Tearly_stopping：尝试最小次数
  143. # measure_option: 构建运行选项 
  144. # tuning_records： 输出的优化记录文件
  145. tuning_option = {
  146.     "tuner": "xgb",
  147.     "trials": 10,
  148.     "early_stopping": 100,
  149.     "measure_option": autotvm.measure_option(
  150.         builder=autotvm.LocalBuilder(build_func="default"), runner=runner
  151.     ),
  152.     "tuning_records": "resnet-50-v2-autotuning.json",
  153. }
  156. # begin by extracting the taks from the onnx model
  157. tasks = autotvm.task.extract_from_program(mod["main"], target=target, params=params)
  159. # Tune the extracted tasks sequentially.
  160. for i, task in enumerate(tasks):
  161.     prefix = "[Task %2d/%2d] " % (i + 1, len(tasks))
  162.     tuner_obj = XGBTuner(task, loss_type="rank")
  163.     tuner_obj.tune(
  164.         n_trial=min(tuning_option["trials"], len(task.config_space)),
  165.         early_stopping=tuning_option["early_stopping"],
  166.         measure_option=tuning_option["measure_option"],
  167.         callbacks=[
  168.             autotvm.callback.progress_bar(tuning_option["trials"], prefix=prefix),
  169.             autotvm.callback.log_to_file(tuning_option["tuning_records"]),
  170.         ],
  171.     )
  173. #～～～～～～～ 
  174. # 输出实例：
  175. #   # [Task  1/24]  Current/Best:   10.71/  21.08 GFLOPS | Progress: (60/1000) | 111.77 s Done.
  176. #   # [Task  1/24]  Current/Best:    9.32/  24.18 GFLOPS | Progress: (192/1000) | 365.02 s Done.
  177. #   # [Task  2/24]  Current/Best:   22.39/ 177.59 GFLOPS | Progress: (960/1000) | 976.17 s Done.
  178. #   # [Task  3/24]  Current/Best:   32.03/ 153.34 GFLOPS | Progress: (800/1000) | 776.84 s Done.
  179. #   ....
  180. #   # [Task 24/24]  Current/Best:   25.03/ 146.14 GFLOPS | Progress: (1000/1000) | 1112.55 s Done.
  182. ################################################################################
  183. # 编译优化后的模型
  184. with autotvm.apply_history_best(tuning_option["tuning_records"]):
  185.     with tvm.transform.PassContext(opt_level=3, config={}):
  186.         lib = relay.build(mod, target=target, params=params)
  188. dev = tvm.device(str(target), 0)
  189. module = graph_executor.GraphModule(lib["default"](dev))
  191. #～～～～～～～～～
  192. # 校验优化模型的结果
  194. dtype = "float32"
  195. module.set_input(input_name, img_data)
  196. module.run()
  197. output_shape = (1, 1000)
  198. tvm_output = module.get_output(0, tvm.nd.empty(output_shape)).numpy()
  200. scores = softmax(tvm_output)
  201. scores = np.squeeze(scores)
  202. ranks = np.argsort(scores)[::-1]
  203. for rank in ranks[0:5]:
  204.     print("class='%s' with probability=%f" % (labels[rank], scores[rank]))
  206. #～～～～～～～～～
  207. # 评估性能
  208. import timeit
  210. timing_number = 10
  211. timing_repeat = 10
  212. optimized = (
  213.     np.array(timeit.Timer(lambda: module.run()).repeat(repeat=timing_repeat, number=timing_number))
  214.     * 1000
  215.     / timing_number
  216. )
  217. optimized = {"mean": np.mean(optimized), "median": np.median(optimized), "std": np.std(optimized)}
  220. print("optimized: %s" % (optimized))
  221. print("unoptimized: %s" % (unoptimized))

使用Ansor优化模型

  1. """ 
  2. 使用TVM的Auto Scheduling feature（Ansor）搜索优化解
  3. """
  5. import os
  6. import numpy as np
  7. import tvm
  8. from tvm import te, auto_scheduler
  10. ################################################################################
  11. # 定义矩阵乘计算
  12. # ----------------------------------
  13. # 1 函数使用 `auto_sceduler` 装饰器
  14. # 2 函数返回input/output tensors列表. Ansor据此构建计算图。
  16. @auto_scheduler.register_workload  # Note the auto_scheduler decorator
  17. def matmul_add(N, L, M, dtype):
  18.     A = te.placeholder((N, L), name="A", dtype=dtype)
  19.     B = te.placeholder((L, M), name="B", dtype=dtype)
  20.     C = te.placeholder((N, M), name="C", dtype=dtype)
  22.     k = te.reduce_axis((0, L), name="k")
  23.     matmul = te.compute(
  24.         (N, M),
  25.         lambda i, j: te.sum(A[i, k] * B[k, j], axis=k),
  26.         name="matmul",
  27.         attrs={"layout_free_placeholders": [B]},  # enable automatic layout transform for tensor B
  28.     )
  29.     out = te.compute((N, M), lambda i, j: matmul[i, j] + C[i, j], name="out")
  31.     return [A, B, C, out]
  34. ################################################################################
  35. # 创建搜索任务
  36. # ----------------------
  37. # 为了充分利用硬件特性，Target越具体越好 
  38. #   - replace "llvm" below with "llvm -mcpu=core-avx2" to enable AVX2
  39. #   - replace "llvm" below with "llvm -mcpu=skylake-avx512" to enable AVX-512
  41. target = tvm.target.Target("llvm")
  42. N = L = M = 1024
  43. task = tvm.auto_scheduler.SearchTask(func=matmul_add, args=(N, L, M, "float32"), target=target)
  45. # Inspect the computational graph
  46. print("Computational DAG:")
  47. print(task.compute_dag)
  49. #～～～～～～～～～～～～～
  50. # 设置Ansor的参数
  51. # ---------------------------------
  52. # `num_measure_trials`： 尝试次数，通常可能1000次才能收敛，为测试方便，设为10 
  53. log_file = "matmul.json"
  54. tune_option = auto_scheduler.TuningOptions(
  55.     num_measure_trials=10,
  56.     measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
  57.     verbose=2,
  58. )
  60. ################################################################################
  61. # 执行自动搜索任务  
  62. task.tune(tune_option)
  63. # Apply the best schedule
  64. sch, args = task.apply_best(log_file)
  66. ################################################################################
  67. # 校验优化后的结果
  68. # --------------------------------- 
  69. # auto-scheduler 执行的优化包括 multi-level tiling, layout transformation, parallelization, vectorization, unrolling, and operator fusion.
  70. print("Lowered TIR:")
  71. print(tvm.lower(sch, args, simple_mode=True))
  73. #～～～～～～～～～～～
  74. # 检验正确性，评估性能
  76. func = tvm.build(sch, args, target)
  77. a_np = np.random.uniform(size=(N, L)).astype(np.float32)
  78. b_np = np.random.uniform(size=(L, M)).astype(np.float32)
  79. c_np = np.random.uniform(size=(N, M)).astype(np.float32)
  80. out_np = a_np.dot(b_np) + c_np
  82. dev = tvm.cpu()
  83. a_tvm = tvm.nd.array(a_np, device=dev)
  84. b_tvm = tvm.nd.array(b_np, device=dev)
  85. c_tvm = tvm.nd.array(c_np, device=dev)
  86. out_tvm = tvm.nd.empty(out_np.shape, device=dev)
  87. func(a_tvm, b_tvm, c_tvm, out_tvm)
  89. # Check results
  90. np.testing.assert_allclose(out_np, out_tvm.numpy(), rtol=1e-3)
  91. # Evaluate execution time.
  92. evaluator = func.time_evaluator(func.entry_name, dev, min_repeat_ms=500)
  93. print(
  94.     "Execution time of this operator: %.3f ms"
  95.     % (np.median(evaluator(a_tvm, b_tvm, c_tvm, out_tvm).results) * 1000)
  96. )
  98. # 输出等价的Python schedule API结果
  100. print("[~~~~~~Equivalent python schedule:~~~~~]")
  101. print(task.print_best(log_file))
  103. ################################################################################
  104. # 恢复搜索
  105. def resume_search(task, log_file):
  106.     print("[~~~~Resume search:~~~~~~~]")
  107.     cost_model = auto_scheduler.XGBModel()
  108.     cost_model.update_from_file(log_file)
  109.     search_policy = auto_scheduler.SketchPolicy(
  110.         task, cost_model, init_search_callbacks=[auto_scheduler.PreloadMeasuredStates(log_file)]
  111.     )
  112.     tune_option = auto_scheduler.TuningOptions(
  113.         num_measure_trials=5, measure_callbacks=[auto_scheduler.RecordToFile(log_file)]
  114.     )
  115.     task.tune(tune_option, search_policy=search_policy)
  118. resume_search(task, log_file)