Apache TVM 是一個深度的深度學習編譯框架，適用於 CPU、GPU 和各種機器學習加速芯片。更多 TVM 中文文檔可訪問 →https://tvm.hyper.ai/
作者：Lianmin Zheng, Zhao Wu, Eddie Yan
針對特定 ARM 設備的自動調優對於獲得最佳性能至關重要，本文介紹如何調優整個卷積網絡。
TVM 中 ARM CPU 的算子實現是以 template 形式編寫的，該 template 有許多可調參數（tile 因子，vectorization，unrolling等）。對神經網絡中的所有卷積和深度卷積算子調優後，會生成一個日誌文件，它存儲所有必需算子的最佳參數值。當 TVM 編譯器編譯這些算子時，會查詢這個日誌文件，從而獲取最佳參數值。
我們還發布了一些 ARM 設備的預調參數。可以前往 ARM CPU Benchmark 查看結果。
注意，本教程無法在 Windows 或最新版本的 macOS 上運行。若要運行，需要將本教程的主體包裝在 if name == "__main__": 塊中。
安裝依賴
要在 TVM 中使用 autotvm 包，需要安裝額外的依賴（如果用的是 Python2，請將「3」更改為「2」）：
pip3 install --user psutil xgboost tornado cloudpickle為了讓 TVM 在調優中運行更快，推薦使用 Cython 作為 TVM 的 FFI。在 TVM 的根目錄下，執行如下命令：（若使用 Python2，將「3」改為「2」）：
pip3 install --user cython
sudo make cython3在 Python 代碼中導入包：
import os
import numpy as np
import tvm
from tvm import relay, autotvm
import tvm.relay.testing
from tvm.autotvm.tuner import XGBTuner, GATuner, RandomTuner, GridSearchTuner
from tvm.contrib.utils import tempdir
import tvm.contrib.graph_executor as runtime定義網絡
首先要在 relay 前端 API 中定義網絡，可以從 relay.testing 加載一些預定義的網絡，還可以從 MXNet、ONNX 和 TensorFlow 加載模型。
def get_network(name, batch_size):
"""獲取網絡的符號定義和隨機權重"""
input_shape = (batch_size, 3, 224, 224)
output_shape = (batch_size, 1000)
if "resnet" in name:
n_layer = int(name.split("-")[1])
mod, params = relay.testing.resnet.get_workload(
num_layers=n_layer, batch_size=batch_size, dtype=dtype
)
elif "vgg" in name:
n_layer = int(name.split("-")[1])
mod, params = relay.testing.vgg.get_workload(
num_layers=n_layer, batch_size=batch_size, dtype=dtype
)
elif name == "mobilenet":
mod, params = relay.testing.mobilenet.get_workload(batch_size=batch_size)
elif name == "squeezenet_v1.1":
mod, params = relay.testing.squeezenet.get_workload(
batch_size=batch_size, version="1.1", dtype=dtype
)
elif name == "inception_v3":
input_shape = (batch_size, 3, 299, 299)
mod, params = relay.testing.inception_v3.get_workload(batch_size=batch_size, dtype=dtype)
elif name == "mxnet":
# MXNet 模型的示例
from mxnet.gluon.model_zoo.vision import get_model
block = get_model("resnet18_v1", pretrained=True)
mod, params = relay.frontend.from_mxnet(block, shape={"data": input_shape}, dtype=dtype)
net = mod["main"]
net = relay.Function(
net.params, relay.nn.softmax(net.body), None, net.type_params, net.attrs
)
mod = tvm.IRModule.from_expr(net)
else:
raise ValueError("Unsupported network: " + name)
return mod, params, input_shape, output_shape啓動 RPC Tracker
TVM 使用 RPC session 與 ARM 板進行通信，在調優期間，調優器會將生成的代碼發送到板上並測試板上代碼的速度。
為了加速調優，TVM 使用 RPC Tracker（集中的控制器節點）來管理分佈式設備。例如，若有 10 部手機，可以將它們全部註冊到 Tracker，並行運行 10 次測試，從而加快調優過程。
要啓動 RPC tracker，在主機上運行如下命令。在整個調優過程中都需要 tracker，因此需要為此命令打開一個新終端：
python -m tvm.exec.rpc_tracker --host=0.0.0.0 --port=9190預期輸出：
INFO:RPCTracker:bind to 0.0.0.0:9190將設備註冊到 RPC Tracker
接下來把設備註冊到 Tracker。第一步是為 ARM 設備構建 TVM runtime。
- 對於 Linux：按照 在設備上構建 TVM Runtime 教程操作，然後將設備註冊到 Tracker
python -m tvm.exec.rpc_server --tracker=[HOST_IP]:9190 --key=rk3399（將 [HOST_IP] 換為你的主機 IP 地址）
- 對於 Android：按照此 説明 在 Android 設備上安裝 TVM RPC APK，確保可以通過 Android rpc 測試。在調優期間，打開手機開發者選項並勾選「在更改期間保持屏幕喚醒」，為手機接通電源。
註冊設備後，通過查詢 rpc_tracker 來確認是否註冊成功
python -m tvm.exec.query_rpc_tracker --host=0.0.0.0 --port=9190例如，如果有 2 台華為 mate10 pro、11 台樹莓派 3B 和 2 台 rk3399，則輸出是
Queue Status
----------------------------------
key total free pending
----------------------------------
mate10pro 2 2 0
rk3399 2 2 0
rpi3b 11 11 0
----------------------------------將多個設備註冊到 tracker，從而加快調優測試。
設置調優選項
在調優之前，進行配置。這裏以 RK3399 板為例。根據自己的設備修改 target 和 device_key。若用 Android 手機，請將 use_android 設置為 True。
#### 設備配置 ####
# 將 "aarch64-linux-gnu" 替換為單板的正確 target。
# 此 target 用於交叉編譯。可以通過：code:`gcc -v` 來查詢。
target = tvm.target.Target("llvm -device=arm_cpu -mtriple=aarch64-linux-gnu")
# 根據設備替換 device_key 的值
device_key = "rk3399"
# 若使用 Android 手機，設置 use_android 為 True
use_android = False
#### 調優選項 ####
network = "resnet-18"
log_file = "%s.%s.log" % (device_key, network)
dtype = "float32"
tuning_option = {
"log_filename": log_file,
"tuner": "xgb",
"n_trial": 1500,
"early_stopping": 800,
"measure_option": autotvm.measure_option(
builder=autotvm.LocalBuilder(build_func="ndk" if use_android else "default"),
runner=autotvm.RPCRunner(
device_key,
host="127.0.0.1",
port=9190,
number=5,
timeout=10,
),
),
}開始調優
下面開始從網絡中提取調優任務，並開始調優。接下來我們提供一個簡單的實用函數。它只是一個初始實現，按順序對任務列表進行調優。未來會引入更復雜的調優 scheduler。
# 可跳過此函數的實現。
def tune_tasks(
tasks,
measure_option,
tuner="xgb",
n_trial=1000,
early_stopping=None,
log_filename="tuning.log",
use_transfer_learning=True,
):
# 創建 tmp 日誌文件
tmp_log_file = log_filename + ".tmp"
if os.path.exists(tmp_log_file):
os.remove(tmp_log_file)
for i, tsk in enumerate(reversed(tasks)):
prefix = "[Task %2d/%2d] " % (i + 1, len(tasks))
# 創建調優器
if tuner == "xgb":
tuner_obj = XGBTuner(tsk, loss_type="reg")
elif tuner == "xgb_knob":
tuner_obj = XGBTuner(tsk, loss_type="reg", feature_type="knob")
elif tuner == "xgb_itervar":
tuner_obj = XGBTuner(tsk, loss_type="reg", feature_type="itervar")
elif tuner == "xgb_curve":
tuner_obj = XGBTuner(tsk, loss_type="reg", feature_type="curve")
elif tuner == "xgb_rank":
tuner_obj = XGBTuner(tsk, loss_type="rank")
elif tuner == "xgb_rank_knob":
tuner_obj = XGBTuner(tsk, loss_type="rank", feature_type="knob")
elif tuner == "xgb_rank_itervar":
tuner_obj = XGBTuner(tsk, loss_type="rank", feature_type="itervar")
elif tuner == "xgb_rank_curve":
tuner_obj = XGBTuner(tsk, loss_type="rank", feature_type="curve")
elif tuner == "xgb_rank_binary":
tuner_obj = XGBTuner(tsk, loss_type="rank-binary")
elif tuner == "xgb_rank_binary_knob":
tuner_obj = XGBTuner(tsk, loss_type="rank-binary", feature_type="knob")
elif tuner == "xgb_rank_binary_itervar":
tuner_obj = XGBTuner(tsk, loss_type="rank-binary", feature_type="itervar")
elif tuner == "xgb_rank_binary_curve":
tuner_obj = XGBTuner(tsk, loss_type="rank-binary", feature_type="curve")
elif tuner == "ga":
tuner_obj = GATuner(tsk, pop_size=50)
elif tuner == "random":
tuner_obj = RandomTuner(tsk)
elif tuner == "gridsearch":
tuner_obj = GridSearchTuner(tsk)
else:
raise ValueError("Invalid tuner: " + tuner)
if use_transfer_learning:
if os.path.isfile(tmp_log_file):
tuner_obj.load_history(autotvm.record.load_from_file(tmp_log_file))
# 開始調優
tsk_trial = min(n_trial, len(tsk.config_space))
tuner_obj.tune(
n_trial=tsk_trial,
early_stopping=early_stopping,
measure_option=measure_option,
callbacks=[
autotvm.callback.progress_bar(tsk_trial, prefix=prefix),
autotvm.callback.log_to_file(tmp_log_file),
],
)
# 選擇最佳記錄到緩存文件
autotvm.record.pick_best(tmp_log_file, log_filename)
os.remove(tmp_log_file)最後啓動調優任務，並評估端到端性能。
def tune_and_evaluate(tuning_opt):
# 從 relay 程序中提取工作負載
print("Extract tasks...")
mod, params, input_shape, _ = get_network(network, batch_size=1)
tasks = autotvm.task.extract_from_program(
mod["main"], target=target, params=params, ops=(relay.op.get("nn.conv2d"),)
)
# 運行調優任務
print("Tuning...")
tune_tasks(tasks, **tuning_opt)
# 編譯具有歷史最佳記錄的內核
with autotvm.apply_history_best(log_file):
print("Compile...")
with tvm.transform.PassContext(opt_level=3):
lib = relay.build_module.build(mod, target=target, params=params)
# 導出庫
tmp = tempdir()
if use_android:
from tvm.contrib import ndk
filename = "net.so"
lib.export_library(tmp.relpath(filename), ndk.create_shared)
else:
filename = "net.tar"
lib.export_library(tmp.relpath(filename))
# 上傳模塊到設備
print("Upload...")
remote = autotvm.measure.request_remote(device_key, "127.0.0.1", 9190, timeout=10000)
remote.upload(tmp.relpath(filename))
rlib = remote.load_module(filename)
# 上傳參數到設備
dev = remote.device(str(target), 0)
module = runtime.GraphModule(rlib["default"](dev))
data_tvm = tvm.nd.array((np.random.uniform(size=input_shape)).astype(dtype))
module.set_input("data", data_tvm)
# 評估
print("Evaluate inference time cost...")
print(module.benchmark(dev, number=1, repeat=10))
# 不在網頁服務器中運行調優，因為它耗時很久。
# 取消註釋運行下一行
# tune_and_evaluate(tuning_option)樣本輸出
調優需要編譯許多程序，並從中提取特徵，所以推薦使用高性能的 CPU。下面列出了一個輸出示例。在 32T AMD Ryzen Threadripper 設備上，大約耗時 2 個小時。
Extract tasks...
Tuning...
[Task 1/12] Current/Best: 22.37/ 52.19 GFLOPS | Progress: (544/1000) | 406.59 s Done.
[Task 2/12] Current/Best: 6.51/ 18.77 GFLOPS | Progress: (608/1000) | 325.05 s Done.
[Task 3/12] Current/Best: 4.67/ 24.87 GFLOPS | Progress: (480/1000) | 372.31 s Done.
[Task 4/12] Current/Best: 11.35/ 46.83 GFLOPS | Progress: (736/1000) | 602.39 s Done.
[Task 5/12] Current/Best: 1.01/ 19.80 GFLOPS | Progress: (448/1000) | 262.16 s Done.
[Task 6/12] Current/Best: 2.47/ 23.76 GFLOPS | Progress: (672/1000) | 563.85 s Done.
[Task 7/12] Current/Best: 14.57/ 33.97 GFLOPS | Progress: (544/1000) | 465.15 s Done.
[Task 8/12] Current/Best: 1.13/ 17.65 GFLOPS | Progress: (576/1000) | 365.08 s Done.
[Task 9/12] Current/Best: 14.45/ 22.66 GFLOPS | Progress: (928/1000) | 724.25 s Done.
[Task 10/12] Current/Best: 3.22/ 15.36 GFLOPS | Progress: (864/1000) | 564.27 s Done.
[Task 11/12] Current/Best: 11.03/ 32.23 GFLOPS | Progress: (736/1000) | 635.15 s Done.
[Task 12/12] Current/Best: 8.00/ 21.65 GFLOPS | Progress: (1000/1000) | 1111.81 s Done.
Compile...
Upload...
Evaluate inference time cost...
Mean inference time (std dev): 162.59 ms (0.06 ms)下載 Python 源代碼：tune_relay_arm.py
下載 Jupyter Notebook：tune_relay_arm.ipynb
