使用TVM编译优化Pytorch模型

为什么需要编译优化

1. 当前的深度学习框架为了方便用户使用，定义的算子非常粗粒度，计算图构建的灵活性在一定程度上降低了执行效率

2. 当前的深度学习框架使用cudnn或MKL等通用算子实现，这些算子实现为了兼容多数场景无法达到极致优化

3. 当前的深度学习框架无法直接在异构框架中运行，或在异构框架下无计算优化

下面以一个pytorch-3D-Resnet模型的编译优化为例，为大家简单介绍如何使用TVM进行编译优化。

TVM是什么

TVM是陈天奇基于 Halide的思想提出的深度学习自动代码生成工具，能够将深度学习模型编译成不同后端的的代码，在编译过程中包含了两个层面的优化：计算图级别的优化，算子级别的优化。

计算图级别的优化包括了常见的算子融合、常量折叠等，在算子级别，TVM的Auto Tuning功能可以自动搜索在当前设备环境下的最优算子实现。

Pytorch模型向TVM模型转换

1. 加载pytorch模型

 checkpoint = torch.load('resume_path')
 model.load_state_dict(checkpoint['state_dict'])
 model.eval()

1. 设置输入shape并初始化假输入，由于TVM目前只能转换、优化静态shape的模型，所以此步骤很重要

 input_shape = (2, 3, 16, 112, 112)
 input_data = torch.randn(input_shape)

1. 使用torch.jit.trace，运行模型，记录模型在所有张量上执行的操作，并记录转换为Torch脚本方法。

 scripted_model = torch.jit.trace(torch_model, input_data).eval()

1. 转换成tvm模型

 from tvm import relay
 mod, params = relay.frontend.from_pytorch(scripted_model, [('input', input_shape)])

初次编译优化

TVM编译优化方法及其对应等级，0代表不做任何优化，设置高优化等级默认包含低等级方法

 "SimplifyInference": 0
 "OpFusion": 1
 "FoldConstant": 2
 "FoldScaleAxis": 3
 "AlterOpLayout": 3
 "CanonicalizeOps": 3
 "CanonicalizeCast": 3
 "EliminateCommonSubexpr": 3
 "CombineParallelConv2D": 4
 "CombineParallelDense": 4
 "CombineParallelBatchMatmul": 4
 "FastMath": 4

编译优化

 with tvm.transform.PassContext(opt_level=3):
     lib = relay.build(mod, target='cuda', params=params)

由于缺少默认的算子优化实现方式，在执行上述语句时，会报很多类似如下的警告，使用这种未优化算子，计算性能会有严重下降

 Cannot find config for target=cuda -keys=cuda,gpu -max_num_threads=1024 -thread_warp_size=32, workload=('conv3d_ncdhw.cuda', ('TENSOR', (2, 3, 16, 112, 112), 'float32'), ('TENSOR', (64, 3, 7, 7, 7), 'float32'), (1, 2, 2), (3, 3, 3, 3, 3, 3), (1, 1, 1), 'float32'). A fallback configuration is used, which may bring great performance regression.
 Cannot find config for target=cuda -keys=cuda,gpu -max_num_threads=1024 -thread_warp_size=32, workload=('conv3d_ncdhw.cuda', ('TENSOR', (2, 64, 8, 28, 28), 'float32'), ('TENSOR', (64, 64, 1, 1, 1), 'float32'), (1, 1, 1), (0, 0, 0, 0, 0, 0), (1, 1, 1), 'float32'). A fallback configuration is used, which may bring great performance regression.

然后使用编译后的TVM模型进行推理实验，推理780个样本的总耗时7930.9秒，几乎是10秒一个，看来这种简单直接的调用反而还降低了100多倍执行效率。

使用CUDNN\CUBLAS进行编译

既然没有默认的算子优化方式，那么很自然的可以考虑调用NVIDIA提供的cudnn/cublas来实现。

编译优化的target中添加cudnn，cublas库。

 with tvm.transform.PassContext(opt_level=3):
     lib = relay.build(mod, target='cuda -libs=cudnn,cublas', params=params)

TVM会搜索CUDNN中的8种卷积实现，评估消耗的时间和显存，选出最优实现

 0) CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM - time: xx ms, Memory: xx
 1) CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM - time: xx ms, Memory: xx
 2) CUDNN_CONVOLUTION_FWD_ALGO_GEMM - time: xx ms, Memory: xx
 3) CUDNN_CONVOLUTION_FWD_ALGO_DIRECT - time: xx ms, Memory: xx
 4) CUDNN_CONVOLUTION_FWD_ALGO_FFT - time: xx ms, Memory: xx
 5) CUDNN_CONVOLUTION_FWD_ALGO_FFT_TILING - time: xx ms, Memory: xx
 6) CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD - time: xx ms, Memory: xx
 7) CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD_NONFUSED - time: xx ms, Memory: xx

然后使用编译后的TVM模型进行推理实验，推理780个样本的总耗时65.8秒，这个性能与直接使用pytorch引擎直接推理性能相似。对比无优化的GPU实现，性能提升100多倍，看来CUDNN在深度学习算子的通用优化上已经做的非常好了，但如果业务需求并不满足于这种通用优化策略，想要进一步提升的时候，就轮到TVM的Auto Tuning出场了。

搜索算子实现策略

没有默认的算子优化方式，还可以使用TVM提供的Auto Tuning功能，使用机器学习算法搜索高效的算子实现，搜索出的实现方式性能通常相似或高于经验丰富的手工设计。

1. 设置需要搜索优化策略的算子，这里选择搜索所有3d卷积的实现

 tasks = autotvm.task.extract_from_program(
         mod["main"], target=tvm.target.cuda(0), params=params, ops=(relay.op.get("nn.conv3d"),)
     )

1. 设置一个优化器，并逐个算子优化，这里选择使用XGBoost进行优化策略搜索

 from tvm.autotvm.tuner import XGBTuner
 from tvm import autotvm
 ​
 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),
                  )
 tmp_log_file = 'tuning.log'
 ​
 for i, tsk in enumerate(reversed(tasks)):
     prefix = "[Task %2d/%2d] " % (i + 1, len(tasks))
     tuner_obj = XGBTuner(tsk, loss_type="rank")
     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),
         ],
     )

搜索运行时，控制台输出

 Tuning...
 [Task  1/27]  Current/Best:  236.83/ 503.69 GFLOPS | Progress: (1344/2000) | 2461.30 s Done.
 [Task  2/27]  Current/Best:  124.94/ 629.35 GFLOPS | Progress: (1344/2000) | 2306.06 s Done.
 ...

搜索算子实现通常需要消耗很长时间，在本实验中搜索27个算子，大约消耗了20个小时，搜索完成后，结果保存在tuning.log文件中。

1. 在优化模型时使用搜索出的算子实现文件，在执行过程中不会再打印出前面找不到算子实现的告警信息

 with autotvm.apply_history_best("tuning.log"):
     with tvm.transform.PassContext(opt_level=3):
         lib = relay.build(mod, target='cuda', params=params)

然后使用编译后的TVM模型进行推理实验，推理780个样本的总耗时32.95秒，比CUDNN快了1倍

总结

本次实验在k80中进行，视频时长约10min，切分成780个样本进行推理，总结推理性能如下：

TVM提供的算子实现搜索能力还是非常不错的，虽然比较耗时，但能获得非常棒的效果。对于推理性能要求比较高的项目，TVM是一个不错的选择。