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]
