ARM-CPU卷积网络的自动调谐

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