初識 TVM

如有圖像或公式顯示錯誤，可以訪問我的個人博客：https://www.wanglichun.tech/2019/11/15/tvm/

筆者也是最近偶然的機會纔開始接觸TVM，使用過後發現，經過auto-tuning後的TVM模型在速度是竟然超過了TensorRT,並且筆者使用的是MXNet框架，TVM對MXNet絕對的友好，對於Pytorch等模型，可以使用ONNX，操作也一樣簡單，使用起來基本類似一鍵操作，本篇文章是筆者對TVM的簡單整理，也算是對TVM的入門。

當然如何想詳細瞭解TVM，還請閱讀TVM的主頁以及論文，文章最後有鏈接。

TVM簡介

隨着深度學習的發展，深度學習的能力可以說是越來越強大，識別率節節攀升，甚至超過人類。於此同時，深度學習框架也變得越來越多，目前比較主流的深度學習框架包括：Pytorch、TensorFlow、Mxnet、Caffe、Keras等。

一般進行深度學習任務包括兩部分，一是訓練出精度比較高的模型，然後將其部署到對應的目標機器上。

針對第一部分，自然我們可以使用各種深度學習框架，通過修改網絡調參等，訓練出精度比較滿意的模型，一般情況，在訓練深度學習模型的時候，都會使用到GPU。

針對部署，這裏的目標機包括服務器、手機、其他硬件設備等等。部署的模型自然是希望越快越好，所以硬件廠商一般會針對自己的硬件設備進行一定的優化，以使模型達到更高的效率，比如Nvidia的TensorRT。但是框架這麼多，硬件平臺這麼多，並不是所有的硬件平臺都像Nvidia提供了硬件加速庫，而即使做了加速，要適應所有的深度學習訓練框架，也是一件比較難的事情。

其實介紹了這麼多總結起來就是兩個問題：

1. 在進行模型部署的時候，我們是否可以對不同框架訓練的模型均生成統一的模型，解決硬件平臺需要適配所有框架的問題？
2. 在進行模型部署的時候，我們是否可以自動化的針對不同的硬件進行優化，進而得到高效的模型？

TVM實際上就是在解決這兩個問題，並且解決的還不錯。

那麼TVM是什麼？

TVM is an open deep learning compiler stack for CPUs, GPUs, and specialized accelerators. It aims to close the gap between the productivity-focused deep learning frameworks, and the performance- or efficiency-oriented hardware backends.

TVM是一個開源的可面向多種硬件設備的深度學習編譯器，它的作用在於打通模型框架、模型表現以及硬件設備的鴻溝，進而得到表現最好的可部署的深度學習模型，實現端到端的深度學習模型部署。

TVM做了哪些工作

針對第一個問題：

TVM將不同前端（深度學習框架）訓練的模型，轉換爲統一的中間語言表示，如果想詳細理解這裏，可以瞭解一下NNVM，NNVM是陳天奇團隊開發的可以針對不同框架進行深度學習編譯的框架，在TVM中，陳天奇團隊進一步優化，實現了NNVM的第二代Relay。Relay是TVM中實現的一種高級IR，可以簡單理解爲另一種計算圖表示。其在TVM所處的位置如下圖所示，並且該部分實現了比如運算融合等操作，可以提升一部分模型效率。

Relay在優化中的位置

針對第二個問題：

TVM設計了對不同的硬件後端，自動優化tensor操作，以達到加速的目的。該部分的實現，TVM使用機器學習的方法進行計算空間的最優化搜索，通過在目標硬件上跑大量trial，來獲得該硬件上相關運算（例如卷積）的最優實現。詳細介紹可以參考TVM主頁以及論文。

TVM tuning可以對不同硬件進行tensor優化

TVM 安裝

不同環境的安裝方法可以參考tvm的官網：https://docs.tvm.ai/install/index.html

對於安裝環境，我還是強烈推薦docker的,會少很多坑。

TVM 使用

TVM的使用可以閱讀一下tvm提供的tutorials：https://docs.tvm.ai/tutorials/

主要推薦兩部分：

• compile deep learning models
• auto tuning

其實簡單的使用主要就是這兩塊內容，如果不想細研究其代碼，可以將其當成一個工具使用，通過compile deep learning models,無論你使用什麼樣的框架，都可以生成統一的模型，一般會生成3個東西如下：

這裏一般會做一些層的融合等操作，速度會有一定的提升的，但是不是特別大。這時如果你需要進一步提速可以試試auto tuning,這部分可以參考tutorials以及下面的例子代碼，auto-tune的時間一般比較長，但是效果還是比較顯著的，本地測試，resnet在nvidia 1080ti上可以提高3倍左右。

Demo代碼

TVM的原理很複雜但是使用起來還是比較方便的，下面是使用MXNet進行TVM轉換的demo。

代碼一：生成TVM模型。

import tvm
from tvm import relay
from tvm.relay import testing
from tvm.contrib import graph_runtime
import mxnet as mx
from tvm.contrib.download import download_testdata
import numpy as np
import time

## load mxnet model
prefix = '/Models/resnetv1d-101'
epoch = 13
mx_sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
shape_dict = {'data': (1, 3, 224, 224)}

relay_func, relay_params = relay.frontend.from_mxnet(mx_sym, shape_dict,
        arg_params=arg_params, aux_params=aux_params)


target = 'cuda'
with relay.build_config(opt_level=3):
    graph, lib, params = relay.build(relay_func, target, params=relay_params)
# run forward



from PIL import Image
image = Image.open('test.jpg').resize((224, 224))
def transform_image(im):
    im = np.array(im).astype(np.float32)
    im = np.transpose(im, [2, 0, 1])
    im = im[np.newaxis, :]
    return im
x = transform_image(image)
# let's go
ctx = tvm.gpu(0)
dtype = 'float32'

m = graph_runtime.create(graph, lib, ctx)
## set input data
m.set_input('data', tvm.nd.array(x.astype(dtype)))
## set input params
m.set_input(**params)
t1 = time.time()
m.run()
t2 = time.time()
# get output
outputs = m.get_output(0)
top1 = np.argmax(outputs.asnumpy()[0])
print(outputs, str(t2-t1))

### evaluate inference time

ftimer = m.module.time_evaluator('run', ctx, number=1, repeat=100)
prof_res = np.array(ftimer().results) * 1000
print('time cost : mean:{}'.format(np.mean(prof_res)))





# save model

path_lib = '/Outputs/tvm/deploy_resnet101_v1d_lib.tar'
lib.export_library(path_lib)

with open('/Outputs/tvm/deploy_resnet101_v1d_graph.json', 'w') as f:
    f.write(graph)
with open('/Outputs/tvm/deploy_params', 'wb') as f:
    f.write(relay.save_param_dict(params))
    
    
# load model back

loaded_json = open('/Outputs/tvm/deploy_resnet101_v1d_graph.json').read()
loaded_lib = tvm.module.load(path_lib)
loaded_params = bytearray(open('/Outputs/tvm/deploy_params', 'rb').read())
module = graph_runtime.create(loaded_json, loaded_lib, ctx)
module.load_params(loaded_params)

tvm_data = tvm.nd.array(x.astype(dtype))
module.run(data=tvm_data)
outputs = module.get_output(0)
print(outputs)

代碼二：auto-tuning

這部分耗時較長，一個resnet101模型，在1080ti上面可能要tune1到2天的時間。

import os

import numpy as np
import mxnet as mx
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
import argparse

def get_network(dtype, args):
    """Get the symbol definition and random weight of a network"""
    input_shape = (args.batch_size, 3, 224, 224)

    # 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, dtype=dtype)
    # 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 = (1, 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)
    # else:
    #     raise ValueError("Unsupported network: " + name)

    prefix = '/Models/{}/{}'.format(args.version, args.model_name)
    epoch = args.model_index
    mx_sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)

    mod, params = relay.frontend.from_mxnet(mx_sym, shape={'data': input_shape}, dtype=dtype, arg_params=arg_params,
                                            aux_params=aux_params)
    net = mod["main"]
    net = relay.Function(net.params, relay.nn.softmax(net.body), None, net.type_params, net.attrs)
    mod = relay.Module.from_expr(net)
    return mod, params, input_shape




# 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,
               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

    # 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 == '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
        n_trial = min(n_trial, len(tsk.config_space))
        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(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, target, log_file, dtype, args):
    # extract workloads from relay program
    print("Extract tasks...")
    mod, params, input_shape = get_network(dtype, args)
    tasks = autotvm.task.extract_from_program(mod["main"], target=target,
                                              params=params, ops=(relay.op.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 relay.build_config(opt_level=3):
            graph, lib, params = relay.build_module.build(
                mod, target=target, params=params)

        # export library
        tmp = tempdir()
        filename = "/Outputs/tvm_autotuning/{}/{}_auto_tune_deploy_batch_{}_lib.tar".format(args.version,args.model_name, args.batch_size)
        lib.export_library(tmp.relpath(filename))

        with open('/Outputs/tvm_autotuning/{}/{}_auto_tune_deploy_batch_{}_graph.json'.format(args.version,args.model_name,args.batch_size) , 'w') as f:
            f.write(graph)
        with open('/Outputs/tvm_autotuning/{}/{}_auto_tune_deploy_batch_{}_params.params'.format(args.version,args.model_name,args.batch_size) , 'wb') as f:
            f.write(relay.save_param_dict(params))

        # load parameters
        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)

        # evaluate
        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)))

# We do not run the tuning in our webpage server since it takes too long.
# Uncomment the following line to run it by yourself.

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='score a model on a dataset')
    parser.add_argument('--version', type=str, default='porno')
    parser.add_argument('--model-name', type=str, default='resnetv1d-101-320x320')
    parser.add_argument('--model-index', type=int, default=16)
    parser.add_argument('--batch-size', type=int, default=1)
    parser.add_argument('--tag', type=str, default='')

    args = parser.parse_args()

    if not os.path.exists(os.path.join('/Outputs/tvm_autotuning/{}'.format(args.version))):
        os.mkdir(os.path.join('/Outputs/tvm_autotuning/{}'.format(args.version)))

    #### DEVICE CONFIG ####
    target = tvm.target.cuda()

    #### TUNING OPTION ####
    log_file = '/Outputs/tvm_autotuning/{}/{}_batch_{}.log'.format(args.version, args.model_name, args.batch_size)
    dtype = 'float32'

    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.LocalRunner(number=20, repeat=3, timeout=4, min_repeat_ms=150),
            # runner=autotvm.RPCRunner(
            #     '1080ti',  # change the device key to your key
            #     '0.0.0.0', 9190,
            #     number=20, repeat=3, timeout=4, min_repeat_ms=150)
        ),
    }


tune_and_evaluate(tuning_option, target, log_file, dtype, args)

TensorRT

這裏簡單介紹一下TensorRT,也是模型加速的利器，並且tvm和tensorRT做的對與模型圖的優化都差不多，可以參考。

TensorRT是Nvidia出品的用於將不同框架訓練的模型部署到GPU的加速引擎，可以自動將不同框架的模型轉換爲TensorRT模型，並進行模型加速。

TensorRT進行模型加速主要有兩點：

• TensorRT支持int8以及FP16計算
• TensorRT對網絡進行重構以及優化:

去掉網絡中的無用層

網絡結構的垂直整合

網絡結構的水平融合

原始網絡

縱向融合

橫向融合

參考資料

TVM官網： https://tvm.ai/

TVM論文：arxiv: https://arxiv.org/abs/1802.04799

tensorRT加速參考文獻：https://blog.csdn.net/xh_hit/article/details/79769599

Nvidia參考文獻：https://devblogs.nvidia.com/production-deep-learning-nvidia-gpu-inference-engine/