TVM Compiler中文教程：TVM为NVIDIA GPU自动调优卷积网络

TVM为NVIDIA GPU自动调优卷积网络

针对特定设备和工作负载进行自动调优对于获得最佳性能至关重要。这篇教程是关于TVM如何为NVIDIA GPU调优整个卷积网络。

TVM中NVIDIA GPU的算子实现是以模板template形式编写的。模板有许多可调旋钮knobs（平铺tile因子，展开unroll等）。我们将调优神经网络中的所有convolutin和depthwise-convolution算子。调优后，我们生成一个log文件，存储所有需要的算子的最佳旋钮值。当TVM编译器编译这些算子时，它将查询此log文件获得最佳旋钮值。

我们还为一些NVIDIA GPU发布了预先调优的参数。您可以访问NVIDIA GPU Benchmark查看结果。

加载依赖库

安装步骤与TVM Compiler中文教程：使用TVM编写可调模板和使用自动调优器中步骤一样，安装相应的软件包即可。

python代码中导入包：

import os
import numpy as np

import tvm
from tvm import autotvm
from tvm import relay
import tvm.relay.testing
from tvm.autotvm.tuner import XGBTuner, GATuner, RandomTuner, GridSearchTuner
from tvm.contrib.util import tempdir
import tvm.contrib.graph_runtime as runtime

定义网络

首先，我们需要调用relay前端API中定义网络。我们可以从nnvm.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])#获取层数
        net, params = relay.testing.resnet.get_workload(nnm_layers=n_layer, batch_size=batch_size, dtype=dtype)
    elif "vgg" in name:
        n_layer = int(name.split('-')[1])
        net, params = relay.testing.vgg.get_workload(num_layers=n_layer, batch_size=batch_size, dtype=dtype)
    elif name == 'mobilenet':
        net, params = relay.testing.mobilenet.get_workload(batch_size=batch_size, dtype=dtype)
    elif name == 'squeezenet_v1.1':
        net, params = relay.testing.squeezenet.get_workload(batch_size=batch_size, version='1.1', dtype=dtype)
    elif name == 'inception_v3':
        input_shape = (1, 3, 299, 299)
        net, params = relay.testing.inception_v3.get_workload(batch_size=batch_size, dtype=dtype)
    elif name == 'mxnet':
        #导入mxnet模型
        from mxnet.gluon.model_zoo.version import get_model
        block = get_model('resnet18_v1', pretrained=True)
        net,params = relay.frontend.from_mxnet(block, shape={'data':input_shape}, dtype=dtype)
        net = relay.Function(net.params, relay.nn.softmax(net.body),None,net.type_params,net.attrs)
   else:
        raise ValueError("Unsupported network: " + name)
   return net,params,input_shape,output_shape       

设置调优选项

调优之前，我们定义一些配置。

#设备配置
target = tvm.target.cuda()
#调优选项
network = 'resnet-18'
log_file = "%s,log" % network
dtype = 'float32'

tuning_option = {
    'log_filename': log_file,
    'tunner': 'xgb'
    'n_trial': 2000 #试验2000次
    'early_stopping': 600
    #测试选项
    'measure_option': autotvm.measure_option(
        builder=autotvm.LocalBuilder(timeout=10),
        runner=autotvm.RPCRunner(
            '1080ti',
            '0.0.0.0', 9190,#填自己的ip加port
            number=20,repeat=3,timeout=4,min_repeat_ms=150
        )
    ),
}

如何设置调优选项：

通常，此处提供的默认值效果很好。如果你有大量的时间预算，你可以设置n_trial，early_stopping更大，这使得调优运行时间更长。如果你有多个设备，则可以将所有设备用于测量以加速调优过程。

开始调优

现在我们可以从网络中提取调优任务，并开始调优。在这里，我们提供了一个简单的工具函数来调优任务列表。这个函数只是一个初始实现，它按顺序调优它们。我们将在未来推出更复杂的调优调度程序。

def tune_tasks(tasks,
               measure_option,
               tuner='xgb',
               n_trial=1000,
               early_stopping=None,
               log_filename='tuning.log',
               use_transfer_learning=True,
               try_winograd=True):
	if try_winograd:
        for i in range(len(tasks)):
            try:  # try winograd template
                tsk = autotvm.task.create(tasks[i].name, tasks[i].args,
                                          tasks[i].target, tasks[i].target_host, 'winograd')
                input_channel = tsk.workload[1][1]
                if input_channel >= 64:
                    tasks[i] = tsk
            except Exception:
                pass
	#创建tmp log文件
    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 == 'ga':
            tuner_obj = GATuner(tsk, pop_size=100)
        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
        tuner_obj.tune(n_trial=min(n_trial, len(tsk.config_space)),
                       early_stopping=early_stopping,
                       measure_option=measure_option,
                       callbacks=[
                           autotvm.callback.progress_bar(n_trial, prefix=prefix),
                           autotvm.callback.log_to_file(tmp_log_file)])

    # 选择最佳记录到log文件
    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...")
    net, params, input_shape, out_shape = get_network(network, batch_size=1)
    tasks = autotvm.task.extract_from_program(net, target=target, params=params, ops=(relay.op.nn.conv2d,))

    # 运行调优任务
    print("Tuning...")
    tune_tasks(tasks, **tuning_opt)
    
    #使用历史最佳纪录编译内核函数
    with autotvm.apply_history_best(log_file):
        print("Compile...")
        with relay.build_config(opt_level=3):
            graph, lib, params = relay.build_module.build(
                net, target=target, params=params)

        # 导出库
        tmp = tempdir()
        filename = "net.tar"
        lib.export_library(tmp.relpath(filename))

        # 加载参数用于runtime推理
        ctx = tvm.context(str(target), 0)
        module = runtime.create(graph, lib, ctx)
        data_tvm = tvm.nd.array((np.random.uniform(size=input_shape)).astype(dtype))
        module.set_input('data', data_tvm)
        module.set_input(**params)

        # 评估
        print("Evaluate inference time cost...")
        ftimer = module.module.time_evaluator("run", ctx, number=1, repeat=600)
        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)))

示例的输出：

调优需要编译许多程序并从中提取功能。因此建议使用高性能CPU。下面列出了一个示例输出。在32T AMD Ryzen Threadripper上获得以下输出大约需要4个小时。调优目标是NVIDIA 1080 Ti。 （你可以在编译过程中看到一些错误。如果调优没有卡住，那就没关系。）

Extract tasks...
Tuning...
[Task  1/12]  Current/Best:  541.83/3570.66 GFLOPS | Progress: (960/2000) | 1001.31 s Done.
[Task  2/12]  Current/Best:    0.56/ 803.33 GFLOPS | Progress: (704/2000) | 608.08 s Done.
[Task  3/12]  Current/Best:  103.69/1141.25 GFLOPS | Progress: (768/2000) | 702.13 s Done.
[Task  4/12]  Current/Best: 2905.03/3925.15 GFLOPS | Progress: (864/2000) | 745.94 sterminate called without an active exception
[Task  4/12]  Current/Best: 2789.36/3925.15 GFLOPS | Progress: (1056/2000) | 929.40 s Done.
[Task  5/12]  Current/Best:   89.06/1076.24 GFLOPS | Progress: (704/2000) | 601.73 s Done.
[Task  6/12]  Current/Best:   40.39/2129.02 GFLOPS | Progress: (1088/2000) | 1125.76 s Done.
[Task  7/12]  Current/Best: 4090.53/5007.02 GFLOPS | Progress: (800/2000) | 903.90 s Done.
[Task  8/12]  Current/Best:    4.78/1272.28 GFLOPS | Progress: (768/2000) | 749.14 s Done.
[Task  9/12]  Current/Best: 1391.45/2325.08 GFLOPS | Progress: (992/2000) | 1084.87 s Done.
[Task 10/12]  Current/Best: 1995.44/2383.59 GFLOPS | Progress: (864/2000) | 862.60 s Done.
[Task 11/12]  Current/Best: 4093.94/4899.80 GFLOPS | Progress: (224/2000) | 240.92 sterminate called without an active exception
[Task 11/12]  Current/Best: 3487.98/4909.91 GFLOPS | Progress: (480/2000) | 534.96 sterminate called without an active exception
[Task 11/12]  Current/Best: 4636.84/4912.17 GFLOPS | Progress: (1184/2000) | 1381.16 sterminate called without an active exception
[Task 11/12]  Current/Best:   50.12/4912.17 GFLOPS | Progress: (1344/2000) | 1602.81 s Done.
[Task 12/12]  Current/Best: 3581.31/4286.30 GFLOPS | Progress: (736/2000) | 943.52 s Done.
Compile...
Evaluate inference time cost...
Mean inference time (std dev): 1.07 ms (0.05 ms)

作为参考基准，resnet-18上MXNet + TensorRT的时间成本为1.30ms。所以我们要快一点。

可能遇到的问题：

自动调优模块容易出错。如果你总是看到“0.00 / 0.00 GFLOPS”，那么肯定有问题。

首先，确保设置正确的设备配置。然后，您可以通过在脚本的开头添加这些行来打印调试信息。它将打印每个测量结果，您可以在其中找到有用的错误消息。

import logging
logging.getLogger('autotvm').setLevel(logging.DEBUG)

最后，请随时向我们的社区寻求帮助https://discuss.tvm.ai

使用多个设备扩展测量范围

如果你有多个设备，则可以将所有设备用于测量。 TVM使用RPC Tracker来管理分布式设备。 RPC Tracker是一个集中式主节点。我们可以将所有设备注册到跟踪器。例如，如果我们有10张GPU卡，我们可以将所有这些卡注册到跟踪器，并行执行10次测量，从而加快调整过程。

要启动RPC跟踪器，请在主机上运行此命令。在整个调优过程中需要跟踪器，因此我们需要为此命令打开一个新终端：

python -m tvm.exec.rpc_tracker --host=0.0.0.0 --port=9190

预计输出：

INFO:RPCTracker:bind to 0.0.0.0:9190

然后打开另一个RPC服务器的新终端。我们需要为每个专用设备启动一台服务器。我们使用字符串key来区分设备类型。你可以选择一个你喜欢的名字。 （注意：对于rocm后端，编译器存在一些内部错误，我们需要在参数列表中添加-no-fork。）

python -m tvm.exec.rpc_server --tracker=0.0.0.0:9190 --key=1080ti

注册设备后，我们可以通过查询rpc_tracker来确认

python -m tvm.exec.query_rpc_tracker --host=0.0.0.0 --port=9190

例如，如果我们有四个1080ti，两个titanx和一个gfx900，将输出：

Queue Status
----------------------------------
key          total  free  pending
----------------------------------
1080ti       4      4     0
titanx       2      2     0
gfx900       1      1     0
----------------------------------

最后，我们需要修改调优选项来使用RPCRunner跟踪器。使用下面的代码替换上面的相应部分。

tuning_option = {
    'log_filename': log_file,

    'tuner': 'xgb',
    'n_trial': 2000,
    'early_stopping': 600,

    'measure_option': autotvm.measure_option(
        builder=autotvm.LocalBuilder(timeout=10),
        runner=autotvm.RPCRunner(
            '1080ti',  #将设备key修改为你的key
            '0.0.0.0', 9190,
            number=20, repeat=3, timeout=4, min_repeat_ms=150),
    ),
}