TVM：编译深度学习模型快速上手教程

TVM：编译深度学习模型快速上手教程

本文将展示如何使用 Relay python 前端构建一个神经网络，并使用 TVM 为 Nvidia GPU 生成一个运行时库。 注意我们需要再构建 TVM 时启用了 cuda 和 llvm。

TVM支持的硬件后端总览

在本教程中，我们使用 cuda 和 llvm 作为目标后端。我们先导入 Relay 和 TVM：

import numpy as np

from 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.00106737
 0.00106262 0.00095838 0.00110792 0.00113151]

保存 / 加载编译好的模块

我们可以将 graph，lib 和 parameters 保存到文件中，并在部署的场景下来加载它们。（译者注：这里的代码会将模型保存在临时文件中，想要保存模型，可以自己修改路径）

# save the graph, lib and params into separate files
from tvm.contrib import utils

temp = 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.00106737
 0.00106262 0.00095838 0.00110792 0.00113151]