%%shell
# Installs the latest dev build of TVM from PyPI. If you wish to build
# from source, see https://tvm.apache.org/docs/install/from_source.html
pip install apache-tvm --pre
基于 VIDIA GPU 的神经网络自动调度#
原作者: Lianmin Zheng
针对特定设备和工作负载的自动调优对于获得最佳性能至关重要。下面是关于如何用自动调度器调优 NVIDIA GPU 的整个神经网络的教程。
为了自动调优神经网络，需要将网络划分为小的子图，并独立地调优它们。每个子图被视为一个搜索任务。任务调度程序对时间进行切片，并动态地为这些任务分配时间资源。任务调度器预测每个任务对端到端执行时间的影响，并优先考虑能够最大程度减少执行时间的任务。
对于每个子图，使用 tvm/python/topi 中的 compute 声明来获得张量表达式形式的计算 DAG。然后，使用自动调度器来构造 DAG 的搜索空间，并搜索良好的调度（低级优化）。
与基于模板的 autotvm 依赖手动模板定义搜索空间不同，自动调度程序不需要任何调度模板。换句话说，自动调度器只使用 tvm/python/topi 中的 compute，而不使用现有的调度模板。
import numpy as np
import tvm
from tvm import relay, auto_scheduler
import tvm.relay.testing
from tvm.contrib import graph_executor
定义网络#
首先，需要用 relay 前端 AP I定义网络。可以从 tvm.relay.testing 加载一些预定义的网络。还可以从 MXNet、ONNX、PyTorch 和 TensorFlow 加载模型（参见前端教程）。
对于卷积神经网络，尽管自动调度器可以在任何布局下正确工作，但我们发现 NHWC 布局通常能获得最佳性能。我们还通过自动调度器实现了对 NHWC 布局的更多优化。因此，建议将模型转换为 NHWC 布局以使用自动调度器。可以使用 ConvertLayout pass 在 TVM 中进行布局转换。
def get_network(name, batch_size, layout="NHWC", dtype="float32"):
"""Get the symbol definition and random weight of a network"""
# auto-scheduler prefers NHWC layout
if layout == "NHWC":
image_shape = (224, 224, 3)
elif layout == "NCHW":
image_shape = (3, 224, 224)
else:
raise ValueError("Invalid layout: " + layout)
input_shape = (batch_size,) + image_shape
output_shape = (batch_size, 1000)
if name.startswith("resnet-"):
n_layer = int(name.split("-")[1])
mod, params = relay.testing.resnet.get_workload(
num_layers=n_layer,
batch_size=batch_size,
layout=layout,
dtype=dtype,
image_shape=image_shape,
)
elif name.startswith("resnet3d-"):
n_layer = int(name.split("-")[1])
mod, params = relay.testing.resnet.get_workload(
num_layers=n_layer,
batch_size=batch_size,
layout=layout,
dtype=dtype,
image_shape=image_shape,
)
elif name == "mobilenet":
mod, params = relay.testing.mobilenet.get_workload(
batch_size=batch_size, layout=layout, dtype=dtype, image_shape=image_shape
)
elif name == "squeezenet_v1.1":
assert layout == "NCHW", "squeezenet_v1.1 only supports NCHW layout"
mod, params = relay.testing.squeezenet.get_workload(
version="1.1",
batch_size=batch_size,
dtype=dtype,
image_shape=image_shape,
)
elif name == "inception_v3":
input_shape = (batch_size, 3, 299, 299) if layout == "NCHW" else (batch_size, 299, 299, 3)
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
assert layout == "NCHW"
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)
return mod, params, input_shape, output_shape
# Define the neural network and compilation target
network = "resnet-18"
batch_size = 1
layout = "NHWC"
target = tvm.target.Target("cuda")
dtype = "float32"
log_file = "%s-%s-B%d-%s.json" % (network, layout, batch_size, target.kind.name)
提取搜索任务#
接下来，从网络中提取搜索任务及其权重。任务的权重是该任务的子图在整个网络中出现的次数。通过使用权重，可以将网络的端到端延迟近似为 sum(latency[t] * weight[t])，其中 latency[t] 是任务的延迟，weight[t] 是任务的权重。任务调度器会优化这个目标。
# 从网络中提取任务
print("Extract tasks...")
mod, params, input_shape, output_shape = get_network(network, batch_size, layout, dtype=dtype)
tasks, task_weights = auto_scheduler.extract_tasks(mod["main"], params, target)
for idx, task in enumerate(tasks):
print("========== Task %d (workload key: %s) ==========" % (idx, task.workload_key))
print(task.compute_dag)
开始调优#
现在，设置了一些调优选项并启动搜索任务
measure_ctx启动不同的测量进程以提供隔离。它可以在测量期间保护主进程不受 GPU 崩溃的影响，并避免其他运行时冲突。min_repeat_ms定义每次测量中一次“重复”的最小持续时间。这可以预热 GPU，这对于获得准确的测量结果是必要的。通常，建议值 >= 300 ms。num_measure_trials是在调优期间可以使用的度量试验的数量。可以将它设置为小的数字（例如，200）以进行快速演示运行。在实践中，建议将它设置在900 * len(tasks)左右，这通常足以让搜索收敛。例如 resnet-18 中有 24 个任务，所以可以将其设置为 20000。可以根据自己的时间预算调整该参数。此外，使用
RecordToFile将测量记录转储到日志文件中，测量记录可以用于查询历史，恢复搜索，并在以后进行更多的分析。查阅
tvm.auto_scheduler.TuningOptions、tvm.auto_scheduler.LocalRPCMeasureContext获取更多参数。
def run_tuning():
print("开始调优...")
measure_ctx = auto_scheduler.LocalRPCMeasureContext(repeat=1, min_repeat_ms=300, timeout=10)
tuner = auto_scheduler.TaskScheduler(tasks, task_weights)
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=200, # change this to 20000 to achieve the best performance
runner=measure_ctx.runner,
measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
)
tuner.tune(tune_option)
run_tuning()
解释在调优期间的打印信息
在调优过程中，控制台上将打印大量信息。它们用于调试目的。最重要的信息是任务调度器的输出。下表是示例输出。
----------------------------------------------------------------------
------------------------------ [ Task Scheduler ]
----------------------------------------------------------------------
| ID | Latency (ms) | Speed (GFLOPS) | Trials |
-------------------------------------------------
| 0 | 0.005 | 0.88 | 64 |
| 1 | 0.010 | 99.10 | 64 |
| 2 | 0.006 | 0.00 | 64 |
| 3 | 0.145 | 979.78 | 384 |
| 4 | 0.130 | 1097.02 | 384 |
| 5 | 0.143 | 992.69 | 384 |
| 6 | 0.076 | 1526.86 | 192 |
| 7 | 0.115 | 999.44 | 320 |
| 8 | 0.079 | 1449.39 | 320 |
| 9 | 0.122 | 938.73 | 384 |
| 10 | 0.063 | 1832.98 | 192 |
| 11 | 0.072 | 1763.62 | 256 |
| 12 | 0.062 | 2036.40 | 192 |
| 13 | 0.068 | 1874.44 | 192 |
| 14 | 0.049 | 2346.50 | 128 |
| 15 | 0.076 | 1694.31 | 256 |
| 16 | 0.067 | 1933.30 | 448 |
| 17 | 0.076 | 1680.90 | 256 |
| 18 | 0.022 | 98.43 | 64 |
| 19 | 0.076 | 3112.55 | 192 |
| 20 | 0.013 | 2026.44 | 64 |
| 21 | 0.011 | 1136.69 | 64 |
| 22 | 0.013 | 992.47 | 64 |
| 23 | 0.020 | 627.56 | 64 |
-------------------------------------------------
Estimated total latency: 1.587 ms Trials: 4992 Used time : 13296 s Next ID: 3
该表列出了所有任务的延迟和（估计的）速度。它还列出了所有任务的测量试验分配。最后一行打印这些任务的总加权延迟，这可以粗略估计网络的端到端执行时间。最后一行还输出测试的总数、自动调优所花费的总时间和下一个要调优的任务的 id。
也会有一些 “tvm::Error” 和 CUDA 的错误，因为自动调度程序将尝试一些无效的调度。如果可以继续进行调优，则可以安全地忽略它们，因为这些错误与主进程隔离开来。
提前终止调优
可以通过强制终止此进程提前终止调优。只要为日志文件中的每个任务获得至少一个有效的调度，就应该能够进行编译（见下面的部分）。
编译和评估#
在自动调优之后，可以用找到的最佳调度来编译网络。在自动调优期间，所有测量记录都被转储到日志文件中，因此可以读取日志文件并加载最佳调度。
# Compile with the history best
print("Compile...")
with auto_scheduler.ApplyHistoryBest(log_file):
with tvm.transform.PassContext(opt_level=3, config={"relay.backend.use_auto_scheduler": True}):
lib = relay.build(mod, target=target, params=params)
# Create graph executor
dev = tvm.device(str(target), 0)
module = graph_executor.GraphModule(lib["default"](dev))
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...")
print(module.benchmark(dev, repeat=3, min_repeat_ms=500))
其他技巧#
在调优过程中，自动调度器需要编译许多程序并从中提取特征。该部分是 CPU 密集型的，因此建议使用多核的高性能 CPU，以加快搜索速度。
可以使用
python3 -m tvm.auto_scheduler.measure_record --mode distill -i log.json提取大的日志文件，只保存最好的有用的记录。可以从上一个日志文件恢复搜索。在函数
run_tuning中创建任务调度器时，只需要添加新的参数load_log_file，即tuner = auto_scheduler.TaskScheduler(tasks, task_weights, load_log_file=log_file)。如果有多个目标 CPU，您可以将它们都用于测量，从而使测量并行化。请查阅如何使用 RPC 跟踪器和 RPC 服务器。要在自动调度器中使用 RPC 跟踪器，请将
TuningOptions中的运行器替换为auto_scheduler.RPCRunner。