TVM中如何tune卷积

TVM是针对深度学习而设计的全栈式编译器，以深度学习计算图为输入（可以是tensorflow, mxnet, pytorch, caffe等框架的计算图），经过两层优化：图优化和算子优化，生成底层代码并运行。
本文是关于TVM探究的第一篇记录，但不会详细介绍TVM，只提及与本文内容有关的部分，其余的留作后续做介绍。

TVM的tuner

TVM提供一个auto-tuner来调优算子。传统的机器学习框架往往是借助现成的算子库，比如CPU上的Eigen和MKL，GPU上的cuBLAS和cuDNN等。但是TVM认为这样做的问题在于算子库的算子往往是针对特定场景优化，场景改变（如计算设备变化或输入规模变化）时就不能保证高效，甚至不能使用。所以TVM设想一个动态生成代码的过程，上层只定义算子的数学含义，比如用函数式语言定义算法，称为Compute，底层的具体实现则针对不同的硬件自动完成。目前TVM需要编程者指定想要的实现方法，拿矩阵乘法为例，编程者可以指定这个矩阵乘算法用分块方法，之后TVM会自动完成代码生成。这种实现方法被称为Schedule，其包含的原语种类很多，常用的有20种左右。

虽然有原语很方便，但是有时实现细节需要提供数字信息，比如即使知道要矩阵分块，分块大小仍然有待商榷，如果让编程者决定，他恐怕要尝试多组可能，然后选择最好的。TVM提供了自动工具完成这种参数tune的过程，称为AutoTVM。本文就是探究如何在x86 CPU上tune一个二维卷积，算法为direct卷积（六重循环）。TVM的官方教程中有tune卷积神经网络的例子，但是本文想只tune一个卷积算子，所以入手点在于改造TVM的官方例子。官方例子在这里

官方例子阅读

在官方例子里，大致可以分为两部分，一部分是得到Compute，以及其上定义的Tune Space，整体可以称为template，另一方面是用tuner去tune这个template。这段代码：

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])
        net, 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])
        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':
        # an example for mxnet model
        from mxnet.gluon.model_zoo.vision 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

就是在得到template，但是直接得到了网络，想要探究如何只得到一个算子，需要看看relay.testing.resnet里究竟有什么。通过找寻源码相关内容，可以看到在tvm/relay/testing/resnet.py中

1
2
3

body = layers.conv2d(
            data=data, channels=filter_list[0], kernel_size=(7, 7),
            strides=(2, 2), padding=(3, 3), name="conv0")

一句调用了卷积，那么继续追寻layers.conv2d在做什么，翻开源码，可以发现在tvm/relay/op/nn/nn.py中：

1
2
3

return _make.conv2d(data, weight, strides, padding, dilation,
                        groups, channels, kernel_size, data_layout,
                        kernel_layout, out_layout, out_dtype)

这样就是进入了C语言部分调用。对于TVM来说，这样的调用本质上做的事情是根据function的名字在已注册的function字典中寻找，然后将其返回。所以TVM一定在某个地方定义了conv2d这个函数。
全项目搜索conv2d这个名字，其实只有一句

1	NNVM_REGISTER_OP(conv2d)

最符合要找的目标。跟踪这一句，发现NNVM中定义的conv2d只是个符号，会在从计算图转化为算子时被实例化，而实例化的过程用的topi中定义的算子。于是在topi中寻找conv2d，发现topi中有关于conv2d的Compute定义和Schedule template，而且在tvm/autotvm/task/topi_integration.py中注册了多个卷积算子。到这里就可以着手更改官方例子来只tune一个卷积了。

更改后的代码

def get_operator(data_shape, out_channel, kernel_size, strides, padding, dtype="float32"):
    data = relay.var("data", shape=data_shape, dtype=dtype)
    body = layers.conv2d(data=data, channels=out_channel, kernel_size=kernel_size, strides=strides, padding=padding, name="conv2d")
    return relay.Function(relay.ir_pass.free_vars(body), body)

def get_workload(batch_size, image_shape, out_channel, kernel_size=(3, 3), strides=(1, 1), padding=(1, 1), dtype="float32"):
    data_shape = (batch_size, *image_shape)
    op = get_operator(data_shape, out_channel, kernel_size, strides, padding, dtype=dtype)
    sym, params = create_workload(op)
    return sym, params, data_shape

用来得到卷积算子，模仿样例中的写法

target = "llvm"
batch_size = 1
image_shape = (256, 14, 14)
out_channel = 512
kernel_size = (3, 3)
strides = (1, 1)
padding = (1, 1)
dtype = "float32"
log_file = "topi_conv2d.log"

num_threads = 1
os.environ["TVM_NUM_THREADS"] = str(num_threads)

tuning_option = {
    "log_filename": log_file,
    "tuner": "gridsearch",
    "early_stopping": None,
    "measure_option": autotvm.measure_option(
        builder=autotvm.LocalBuilder(),
        runner=autotvm.LocalRunner(number=10, repeat=1, min_repeat_ms=1000)
    )
}

一些选项，根据具体情况更改

def tune_kernels(tasks, measure_option, tuner="gridsearch", early_stopping=None, log_filename="tuning.log"):
    for i, tsk in enumerate(tasks):
        prefix = "[Task %2d/%2d] " % (i + 1, len(tasks))
        op_name = tsk.workload[0]
        task = autotvm.task.create("topi_nn_conv2d", args=tsk.args, target=target, template_key="direct")
        task.workload = tsk.workload

        # create tuner
        if tuner == 'xgb' or tuner == 'xgb-rank':
            tuner_obj = XGBTuner(task, loss_type='rank')
        elif tuner == 'ga':
            tuner_obj = GATuner(task, pop_size=50)
        elif tuner == 'random':
            tuner_obj = RandomTuner(task)
        elif tuner == 'gridsearch':
            tuner_obj = GridSearchTuner(task)
        else:
            raise ValueError("Invalid tuner: " + tuner)

        # do tuning
        n_trial = len(task.config_space)
        print("trials=", n_trial)
        tuner_obj.tune(n_trial=n_trial,
                       early_stopping=early_stopping,
                       measure_option=measure_option,
                       callbacks=[
                           autotvm.callback.progress_bar(n_trial, prefix=prefix),
                           autotvm.callback.log_to_file(log_filename)])

从官方例子改的，其中最重要的是指明要用哪个卷积算法，这里指定了"topi_nn_conv2d"，其余选择要看/tvm/autotvm/task/task.py中的TASK_TABLE中有哪些，可以自己print出来看看，记得先调用extract_from_program再print，不然TASK_TABLE中没有添加后续注册的任务。

def tune_and_evaluate(tuning_opt):
    op, params, data_shape = get_workload(batch_size, image_shape, out_channel, kernel_size, strides, padding)
    tasks = autotvm.task.extract_from_program(op, target=target, params=params, ops=(relay.op.nn.conv2d,))

    print("Tuning...")
    tune_kernels(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(op, target=target, params=params)

        ctx = tvm.cpu()
        data_tvm = tvm.nd.array((np.random.uniform(size=data_shape)).astype(dtype))
        module = runtime.create(graph, lib, ctx)
        module.set_input("data", data_tvm)
        module.set_input(**params)

        # evaluate
        print("Evaluate inference time cost...")
        ftimer = module.module.time_evaluator("run", ctx, number=10, repeat=1)
        prof_res = np.array(ftimer().results) * 1e3
        print("Time cost is: ", np.mean(prof_res))

最后直接tune就好。

TVM中如何tune卷积

TVM中如何tune卷积

TVM的tuner

官方例子阅读

更改后的代码

FEATURED TAGS

FRIENDS