【TVM教程】利用 TVMC Micro 实验微模子

Apache TVM是一个深度的深度学习编译框架，实用于 CPU、GPU 和各种呆板学习加快芯片。更多 TVM 中文文档可访问 →https://tvm.hyper.ai/
作者：Andrew Reusch, Mehrdad Hessar
本教程先容怎样用 C runtime 主动调优模子。
安装 microTVM Python 依靠项​

TVM 不包罗用于 Python 串行通讯包，因此在利用 microTVM 之前我们必须先安装一个。我们还须要TFLite来加载模子。

1. pip install pyserial==3.5 tflite==2.1
3. # 如果下面的标志为 False，可以跳过下一节（安装 Zephyr）
4. # 安装 Zephyr 约花费20分钟
5. import os
7. use_physical_hw = bool(os.getenv("TVM_MICRO_USE_HW"))

复制代码

安装 Zephyr​

1. # 安装 west 和 ninja
2. python3 -m pip install west
3. apt-get install -y ninja-build
5. # 安装 ZephyrProject
6. ZEPHYR_PROJECT_PATH="/content/zephyrproject"
7. export ZEPHYR_BASE=${ZEPHYR_PROJECT_PATH}/zephyr
8. west init ${ZEPHYR_PROJECT_PATH}
9. cd ${ZEPHYR_BASE}
10. git checkout v3.2-branch
11. cd ..
12. west update
13. west zephyr-export
14. chmod -R o+w ${ZEPHYR_PROJECT_PATH}
16. # 安装 Zephyr SDK
17. cd /content
18. ZEPHYR_SDK_VERSION="0.15.2"
19. wget "https://github.com/zephyrproject-rtos/sdk-ng/releases/download/v${ZEPHYR_SDK_VERSION}/zephyr-sdk-${ZEPHYR_SDK_VERSION}_linux-x86_64.tar.gz"
20. tar xvf "zephyr-sdk-${ZEPHYR_SDK_VERSION}_linux-x86_64.tar.gz"
21. mv "zephyr-sdk-${ZEPHYR_SDK_VERSION}" zephyr-sdk
22. rm "zephyr-sdk-${ZEPHYR_SDK_VERSION}_linux-x86_64.tar.gz"
24. # 安装 python 依赖
25. python3 -m pip install -r "${ZEPHYR_BASE}/scripts/requirements.txt"

复制代码

导入 Python 依靠项​

1. import json
2. import numpy as np
3. import pathlib
5. import tvm
6. from tvm.relay.backend import Runtime
7. import tvm.micro.testing

复制代码

界说模子​

起首在 Relay 中界说一个要在装备上实验的模子，然后从 Relay 模子中创建一个 IRModule，并用随机数添补参数。

1. data_shape = (1, 3, 10, 10)
2. weight_shape = (6, 3, 5, 5)
4. data = tvm.relay.var("data", tvm.relay.TensorType(data_shape, "float32"))
5. weight = tvm.relay.var("weight", tvm.relay.TensorType(weight_shape, "float32"))
7. y = tvm.relay.nn.conv2d(
8.     data,
9.     weight,
10.     padding=(2, 2),
11.     kernel_size=(5, 5),
12.     kernel_layout="OIHW",
13.     out_dtype="float32",
14. )
15. f = tvm.relay.Function([data, weight], y)
17. relay_mod = tvm.IRModule.from_expr(f)
18. relay_mod = tvm.relay.transform.InferType()(relay_mod)
20. weight_sample = np.random.rand(
21.     weight_shape[0], weight_shape[1], weight_shape[2], weight_shape[3]
22. ).astype("float32")
23. params = {"weight": weight_sample}

复制代码

界说 target​

下面界说形貌实验环境的 TVM target，它与其他 microTVM 教程中的 target 界说非常相似。差别之处是用 C Runtime 来天生模子。
在物理硬件上运行时，选择一个 target 和一个形貌硬件的单板。在本教程的 PLATFORM 列表中可以选择多个硬件目的。运行本教程时用 –platform 参数来选择平台。

1. RUNTIME = Runtime("crt", {"system-lib": True})
2. TARGET = tvm.micro.testing.get_target("crt")
4. # 为物理硬件编译
5. # --------------------------------------------------------------------------
6. # 在物理硬件上运行时，选择描述硬件的 TARGET 和 BOARD。
7. # 下面的示例中选择 STM32L4R5ZI Nucleo。
8. if use_physical_hw:
9.     BOARD = os.getenv("TVM_MICRO_BOARD", default="nucleo_l4r5zi")
10.     SERIAL = os.getenv("TVM_MICRO_SERIAL", default=None)
11.     TARGET = tvm.micro.testing.get_target("zephyr", BOARD)

复制代码

提取调优任务​

并非上面打印的 Relay 步伐中的全部算子都可以调优，有些算子不是很紧张，以是只界说了一个实现；其他的作为调优任务没故意义。用 extract_from_program，可以天生可调优任务列表。
由于任务提取涉及到运行编译器，以是起首设置编译器的转换 pass；之后在主动调优过程中应用类似的设置。

1. pass_context = tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True})
2. with pass_context:
3.     tasks = tvm.autotvm.task.extract_from_program(relay_mod["main"], {}, TARGET)
4. assert len(tasks) > 0

复制代码

设置 microTVM​

主动调优前，要界说一个模块加载器，然后将它通报给 tvm.autotvm.LocalBuilder。创建一个 tvm.autotvm.LocalBuilder，并用 builder 和 runner 为主动调优器天生多个测试值。
本教程中可以选择用 x86 主机作为示例，或用来自 Zephyr RTOS 的差别 targets。若用 x86，则通报 –platform=host，也可以从 PLATFORM 列表中选择其他选项。

1. module_loader = tvm.micro.AutoTvmModuleLoader(
2.     template_project_dir=pathlib.Path(tvm.micro.get_microtvm_template_projects("crt")),
3.     project_options={"verbose": False},
4. )
5. builder = tvm.autotvm.LocalBuilder(
6.     n_parallel=1,
7.     build_kwargs={"build_option": {"tir.disable_vectorize": True}},
8.     do_fork=True,
9.     build_func=tvm.micro.autotvm_build_func,
10.     runtime=RUNTIME,
11. )
12. runner = tvm.autotvm.LocalRunner(number=1, repeat=1, timeout=100, module_loader=module_loader)
14. measure_option = tvm.autotvm.measure_option(builder=builder, runner=runner)
16. # 为物理硬件编译
17. if use_physical_hw:
18.     module_loader = tvm.micro.AutoTvmModuleLoader(
19.         template_project_dir=pathlib.Path(tvm.micro.get_microtvm_template_projects("zephyr")),
20.         project_options={
21. "board": BOARD,
22. "verbose": False,
23. "project_type": "host_driven",
24. "serial_number": SERIAL,
25. },
26. )
27.     builder = tvm.autotvm.LocalBuilder(
28.         n_parallel=1,
29.         build_kwargs={"build_option": {"tir.disable_vectorize": True}},
30.         do_fork=False,
31.         build_func=tvm.micro.autotvm_build_func,
32.         runtime=RUNTIME,
33. )
34.     runner = tvm.autotvm.LocalRunner(number=1, repeat=1, timeout=100, module_loader=module_loader)
36.     measure_option = tvm.autotvm.measure_option(builder=builder, runner=runner)

复制代码

运行主动调优​

下面在 microTVM 装备上对每个提取的任务单独运行主动调优。

1. autotune_log_file = pathlib.Path("microtvm_autotune.log.txt")
2. if os.path.exists(autotune_log_file):
3.     os.remove(autotune_log_file)
5. num_trials = 10
6. for task in tasks:
7.     tuner = tvm.autotvm.tuner.GATuner(task)
8.     tuner.tune(
9.         n_trial=num_trials,
10.         measure_option=measure_option,
11.         callbacks=[
12.             tvm.autotvm.callback.log_to_file(str(autotune_log_file)),
13.             tvm.autotvm.callback.progress_bar(num_trials, si_prefix="M"),
14. ],
15.         si_prefix="M",
16. )

复制代码

为未调优的步伐计时​

为了方便比力，编译并运行不实验任何主动调优 schedule 的盘算图，TVM 会为每个算子选择一个随机调优的实现，其性能不如调优的算子。

1. with pass_context:
2.     lowered = tvm.relay.build(relay_mod, target=TARGET, runtime=RUNTIME, params=params)
4. temp_dir = tvm.contrib.utils.tempdir()
5. project = tvm.micro.generate_project(
6. str(tvm.micro.get_microtvm_template_projects("crt")),
7.     lowered,
8.     temp_dir / "project",
9. {"verbose": False},
10. )
12. # 为物理硬件编译
13. if use_physical_hw:
14.     temp_dir = tvm.contrib.utils.tempdir()
15.     project = tvm.micro.generate_project(
16. str(tvm.micro.get_microtvm_template_projects("zephyr")),
17.         lowered,
18.         temp_dir / "project",
19. {
20. "board": BOARD,
21. "verbose": False,
22. "project_type": "host_driven",
23. "serial_number": SERIAL,
24. "config_main_stack_size": 4096,
25. },
26. )
28. project.build()
29. project.flash()
30. with tvm.micro.Session(project.transport()) as session:
31.     debug_module = tvm.micro.create_local_debug_executor(
32.         lowered.get_graph_json(), session.get_system_lib(), session.device
33. )
34.     debug_module.set_input(**lowered.get_params())
35. print("########## Build without Autotuning ##########")
36.     debug_module.run()
37. del debug_module

复制代码

输出结果：

1. /workspace/python/tvm/driver/build_module.py:268: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.
2. "target_host parameter is going to be deprecated. "
3. ########## Build without Autotuning ##########
4. Node Name                                     Ops                                           Time(us)  Time(%)  Shape              Inputs  Outputs  Measurements(us)
5. ---------                                     ---                                           --------  -------  -----              ------  -------  ----------------
6. tvmgen_default_fused_nn_contrib_conv2d_NCHWc  tvmgen_default_fused_nn_contrib_conv2d_NCHWc  310.0 98.73 (1, 2, 10, 10, 3) 2 1 [310.0]
7. tvmgen_default_fused_layout_transform_1       tvmgen_default_fused_layout_transform_1       3.031 0.965 (1, 6, 10, 10) 1 1 [3.031]
8. tvmgen_default_fused_layout_transform         tvmgen_default_fused_layout_transform         0.958 0.305 (1, 1, 10, 10, 3) 1 1 [0.958]
9. Total_time                                    -                                             313.988   -        -                  -       -        -

复制代码

为调优步伐计时​

主动调优完成后，用 Debug Runtime 为整个步伐的实验计时：

1. with tvm.autotvm.apply_history_best(str(autotune_log_file)):
2. with pass_context:
3.         lowered_tuned = tvm.relay.build(relay_mod, target=TARGET, runtime=RUNTIME, params=params)
5. temp_dir = tvm.contrib.utils.tempdir()
6. project = tvm.micro.generate_project(
7. str(tvm.micro.get_microtvm_template_projects("crt")),
8.     lowered_tuned,
9.     temp_dir / "project",
10. {"verbose": False},
11. )
13. # 为物理硬件编译
14. if use_physical_hw:
15.     temp_dir = tvm.contrib.utils.tempdir()
16.     project = tvm.micro.generate_project(
17. str(tvm.micro.get_microtvm_template_projects("zephyr")),
18.         lowered_tuned,
19.         temp_dir / "project",
20. {
21. "board": BOARD,
22. "west_cmd": "west",
23. "verbose": False,
24. "project_type": "host_driven",
25. "serial_number": SERIAL,
26. "config_main_stack_size": 4096,
27. },
28. )
30. project.build()
31. project.flash()
32. with tvm.micro.Session(project.transport()) as session:
33.     debug_module = tvm.micro.create_local_debug_executor(
34.         lowered_tuned.get_graph_json(), session.get_system_lib(), session.device
35. )
36.     debug_module.set_input(**lowered_tuned.get_params())
37. print("########## Build with Autotuning ##########")
38.     debug_module.run()
39. del debug_module

复制代码

输出结果：

1. /workspace/python/tvm/driver/build_module.py:268: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.
2. "target_host parameter is going to be deprecated. "
3. ########## Build with Autotuning ##########
4. Node Name                                     Ops                                           Time(us)  Time(%)  Shape              Inputs  Outputs  Measurements(us)
5. ---------                                     ---                                           --------  -------  -----              ------  -------  ----------------
6. tvmgen_default_fused_nn_contrib_conv2d_NCHWc  tvmgen_default_fused_nn_contrib_conv2d_NCHWc  193.2 98.657 (1, 6, 10, 10, 1) 2 1 [193.2]
7. tvmgen_default_fused_layout_transform_1       tvmgen_default_fused_layout_transform_1       1.778 0.908 (1, 6, 10, 10) 1 1 [1.778]
8. tvmgen_default_fused_layout_transform         tvmgen_default_fused_layout_transform         0.851 0.435 (1, 3, 10, 10, 1) 1 1 [0.851]
9. Total_time                                    -                                             195.83    -        -                  -       -        -

复制代码

下载 Python 源代码：micro_autotune.py
下载 Jupyter notebook：micro_autotune.ipynb