ARM-CPU卷积网络的自动调谐
为特定的ARM设备自动调谐对于获得最佳性能至关重要。这是一个关于如何调整整个卷积网络的资料。
以模板的形式编写了TVM中ARM CPU的操作实现。模板有许多可调旋钮（平铺系数、矢量化、展开等）。将调整神经网络中的所有卷积和深度卷积算子。在调优之后，生成一个日志文件，其中存储了所有所需操作符的最佳旋钮值。当TVM编译器编译这些运算符时，它将查询此日志文件以获得最佳的旋钮值。
还发布了一些arm设备的预调参数。可以转到arm cpu基准测试来查看结果。
本文不会在Windows或最新版本的macOS上运行。要让它运行，需要将本教程的主体包装在if __name__ == "__main__": 块中。
安装依赖项
要在tvm中使用autotvm包，需要安装一些额外的依赖项。（如果使用python2，请将“3”更改为“2”）：
pip3 install --user psutil xgboost tornado
为了使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_runtime as runtime
Define network
首先需要在转换前端API中定义网络。可以从转换测试. 也可以从MXNet、ONNX和TensorFlow加载模型。
def get_network(name, batch_size):
"""Get the symbol definition and random weight of a network"""
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":
# an example for mxnet model
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
Start RPC Tracker
TVM使用RPC会话与ARM板通信。在调谐过程中，调谐器将生成的代码发送到板上，并测量板上代码的速度。
为了扩大调整范围，TVM使用RPC跟踪器来管理分布式设备。RPC跟踪器是一个集中的控制器节点。可以把所有设备注册到跟踪器上。例如，如果有10部电话，可以将它们全部注册到跟踪器中，并行运行10次测量，从而加快调谐过程。
要启动RPC跟踪器，在主机上运行此命令。在整个整定过程中需要跟踪器，所以需要为这个命令打开一个新的终端：
python -m tvm.exec.rpc_tracker --host=0.0.0.0 --port=9190
The expected output is
INFO:RPCTracker:bind to 0.0.0.0:9190
Register devices to RPC Tracker
现在可以在跟踪器上注册设备。第一步是为ARM设备构建TVM运行时。
对于Linux：按照本节“在设备上构建TVM runtime”在设备上构建TVM runtime。然后通过
python -m tvm.exec.rpc_server --tracker=[HOST_IP]:9190 --key=rk3399
（将[HOST_IP]替换为主机的IP地址）
对于Android：按照这个readme页面在Android设备上安装TVM RPC APK。确保能通过android rpc测试。那么已经注册了设备。在调整过程中，必须进入开发者选项，并启用“在改变屏幕时保持屏幕唤醒”并给手机充电以使其稳定。
注册设备后，可以通过查询rpc_tracker来确认。
python -m tvm.exec.query_rpc_tracker --host=0.0.0.0 --port=9190
例如，如果有2个Huawei mate10 pro、11个Raspberry Pi 3B和2个rk3399，则输出可以是
Queue Status
----------------------------------
key total free pending
----------------------------------
mate10pro 2 2 0
rk3399 2 2 0
rpi3b 11 11 0
可以将多个设备注册到跟踪器中，以加速调谐中的测量。
设置调谐选项
在调整之前，应该应用一些配置。这里以RK3399板为例。在设置中，应该相应地修改target和device_key键。如果您使用安卓手机，请将use_android设置为True。
#### DEVICE CONFIG ####
# Replace "aarch64-linux-gnu" with the correct target of your board.
# This target is used for cross compilation. You can query it by :code:`gcc -v` on your device.
target = tvm.target.Target("llvm -device=arm_cpu -mtriple=aarch64-linux-gnu")
# Also replace this with the device key in your tracker
device_key = "rk3399"
# Set this to True if you use android phone
use_android = False
#### TUNING OPTION ####
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="0.0.0.0",
port=9190,
number=5,
timeout=10,
),
),
}
注意
如何设置调整选项
一般来说，这里提供的默认值工作正常。如果有足够长的时间来进行调试，可以提前停止调试。如果设备运行非常慢，或者conv2d操作员有许多gflop，考虑将超时设置得更大。
如果模型有深度卷积，可以考虑将try_spatial_pack_depthwise设置为True，这通常比默认优化性能更好。例如，在ARM CPU A53 2.0GHz上，发现它可以使Mobilenet V1型号上的深度卷积性能提高1.6倍。
Begin Tuning
现在可以从网络中提取调优任务并开始调优。这里，提供了一个简单的实用函数来优化任务列表。这个函数只是一个按顺序调整初始实现。将在将来引入更复杂的调优调度程序。
# You can skip the implementation of this function for this tutorial.
def tune_tasks(
tasks,
measure_option,
tuner="xgb",
n_trial=1000,
early_stopping=None,
log_filename="tuning.log",
use_transfer_learning=True,
):
# create tmp log file
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))
# create tuner
if tuner == "xgb" or tuner == "xgb-rank":
tuner_obj = XGBTuner(tsk, loss_type="rank")
elif tuner == "xgb_knob":
tuner_obj = XGBTuner(tsk, loss_type="rank", feature_type="knob")
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))
# do tuning
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),
],
)
# pick best records to a cache file
autotvm.record.pick_best(tmp_log_file, log_filename)
os.remove(tmp_log_file)
最后，启动优化作业并评估端到端性能。
def tune_and_evaluate(tuning_opt):
# extract workloads from relay program
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"),)
)
# run tuning tasks
print("Tuning...")
tune_tasks(tasks, **tuning_opt)
# compile kernels with history best records
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)
# export library
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))
# upload module to device
print("Upload...")
remote = autotvm.measure.request_remote(device_key, "0.0.0.0", 9190, timeout=10000)
remote.upload(tmp.relpath(filename))
rlib = remote.load_module(filename)
# upload parameters to device
ctx = remote.context(str(target), 0)
module = runtime.GraphModule(rlib["default"](ctx))
data_tvm = tvm.nd.array((np.random.uniform(size=input_shape)).astype(dtype))
module.set_input("data", data_tvm)
# evaluate
print("Evaluate inference time cost...")
ftimer = module.module.time_evaluator("run", ctx, number=1, repeat=10)
prof_res = np.array(ftimer().results) * 1000 # convert to millisecond
print(
"Mean inference time (std dev): %.2f ms (%.2f ms)"
% (np.mean(prof_res), np.std(prof_res))
)
# We do not run the tuning in our webpage server since it takes too long.
# Uncomment the following line to run it by yourself.
# tune_and_evaluate(tuning_option)
Sample Output
调整需要编译许多程序并从中提取特性。因此建议使用高性能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)
注意
遇到困难？
自动调谐模块容易出错。如果总是看到“0.00/0.00 GFLOPS”，那么一定是出了什么问题。
首先，确保设置了正确的设备配置。然后，可以通过在脚本开头添加这些行来打印调试信息。它将打印每个测量结果，可以在其中找到有用的错误消息。
import logging
logging.getLogger('autotvm').setLevel(logging.DEBUG)
https://tvm.apache.org/docs/tutorials/autotvm/tune_relay_arm.html
最后，请随时寻求帮助https://discus.tvm.apache.org
下载Python源代码：tune_relay_arm.py
下载Jupyter笔记：tune_relay_arm.ipynb