TVM：编译深度学习模型快速上手教程
本文将展示如何使用 Relay python 前端构建一个神经网络，并使用 TVM 为 Nvidia GPU 生成一个运行时库。 注意我们需要再构建 TVM 时启用了 cuda 和 llvm。
TVM支持的硬件后端总览
在本教程中，我们使用 cuda 和 llvm 作为目标后端。我们先导入 Relay 和 TVM：
import numpy as npfrom tvm import relay
from tvm.relay import testing
import tvm
from tvm import te
from tvm.contrib import graph_executor
import tvm.testing
使用Relay定义一个神经网络
首先，我们使用 relay 的 python 前端定义一个神经网络。简单起见，我们这里直接使用 relay 中预定义好的 resnet-18 网络。参数按照 Xavier 初始化进行初始化。Relay 同样支持其他的模型格式如 MXNet，CoreML，ONNX 和 TensorFlow。
本教程中，我们假设我们将会在自己的设备上进行推理，batch size 设为1。输入是尺寸为 224 * 224 的 RGB 彩色图像。我们可以调用 tvm.relay.expr.TupleWrapper.astext() 来查看网络结构，
batch_size = 1
num_class = 1000
image_shape = (3, 224, 224)
data_shape = (batch_size,) + image_shape
out_shape = (batch_size, num_class)mod, params = relay.testing.resnet.get_workload(num_layers=18, batch_size=batch_size, image_shape=image_shape
)# set show_meta_data=True if you want to show meta data
print(mod.astext(show_meta_data=False))
输出：
#[version = "0.0.5"]
def @main(%data: Tensor[(1, 3, 224, 224), float32], %bn_data_gamma: Tensor[(3), float32], %bn_data_beta: Tensor[(3), float32], %bn_data_moving_mean: Tensor[(3), float32], %bn_data_moving_var: Tensor[(3), float32], %conv0_weight: Tensor[(64, 3, 7, 7), float32], %bn0_gamma: Tensor[(64), float32], %bn0_beta: Tensor[(64), float32], %bn0_moving_mean:
...
Tensor[(64), float32], %bn0_moving_var: Tensor[(64), float32], %stage1_unit1_bn1_gamma: %88 = nn.dense(%87, %fc1_weight, units=1000) /* ty=Tensor[(1, 1000), float32] */;%89 = nn.bias_add(%88, %fc1_bias, axis=-1) /* ty=Tensor[(1, 1000), float32] */;nn.softmax(%89) /* ty=Tensor[(1, 1000), float32] */
}
编译
下一步就是使用 Relay/TVM 的流程进行编译。用户可以指定编译的优化等级。目前优化等级可以设置为 0 到 3。优化的 pass 包括算子融合，预先计算，排布变换等。
relay.build() 返回三个组件：json 格式的执行图、目标硬件上专门为此图编译函数的 TVM 模块库以及模型的参数 blob。 在编译过程中，Relay 进行图级优化，TVM 进行张量级优化，从而为模型推理服务提供优化的运行时模块。
我们将首先为 Nvidia GPU 编译。 relay.build() 首先在幕后进行了一些图级优化，例如剪枝、融合等，然后将算子（即优化图的节点）注册到 TVM 实现以生成 tvm.module。 为了生成模块库，TVM 将首先将高层 IR 转换为指定目标后端的低层固有 IR，在本例中为 CUDA。 然后生成机器代码，得到模块库。
opt_level = 3
target = tvm.target.cuda()
with tvm.transform.PassContext(opt_level=opt_level):lib = relay.build(mod, target, params=params)
输出：
/home/areusch/ws/tvm3/python/tvm/target/target.py:259: UserWarning: Try specifying cuda arch by adding 'arch=sm_xx' to your target.warnings.warn("Try specifying cuda arch by adding 'arch=sm_xx' to your target.")
运行生成的库
现在，我们可以创建 graph executor 来将模块运行在 Nvidia GPU 上。
# create random input
dev = tvm.cuda()
data = np.random.uniform(-1, 1, size=data_shape).astype("float32")
# create module
module = graph_executor.GraphModule(lib["default"](dev))
# set input and parameters
module.set_input("data", data)
# run
module.run()
# get output
out = module.get_output(0, tvm.nd.empty(out_shape)).numpy()# Print first 10 elements of output
print(out.flatten()[0:10])
输出：
[0.00089283 0.00103331 0.0009094 0.00102275 0.00108751 0.001067370.00106262 0.00095838 0.00110792 0.00113151]
保存 / 加载编译好的模块
我们可以将 graph，lib 和 parameters 保存到文件中，并在部署的场景下来加载它们。（译者注：这里的代码会将模型保存在临时文件中，想要保存模型，可以自己修改路径）
# save the graph, lib and params into separate files
from tvm.contrib import utilstemp = utils.tempdir()
path_lib = temp.relpath("deploy_lib.tar")
lib.export_library(path_lib)
print(temp.listdir())
输出：
['deploy_lib.tar']
# load the module back.
loaded_lib = tvm.runtime.load_module(path_lib)
input_data = tvm.nd.array(data)module = graph_executor.GraphModule(loaded_lib["default"](dev))
module.run(data=input_data)
out_deploy = module.get_output(0).numpy()# Print first 10 elements of output
print(out_deploy.flatten()[0:10])# check whether the output from deployed module is consistent with original one
tvm.testing.assert_allclose(out_deploy, out, atol=1e-5)
输出：
[0.00089283 0.00103331 0.0009094 0.00102275 0.00108751 0.001067370.00106262 0.00095838 0.00110792 0.00113151]
