做网站一般几个人,厦门零基础学seo,重庆网站建设公司联系方式,变更icp备案网站信息TVM#xff1a;编译深度学习模型快速上手教程
本文将展示如何使用 Relay python 前端构建一个神经网络#xff0c;并使用 TVM 为 Nvidia GPU 生成一个运行时库。 注意我们需要再构建 TVM 时启用了 cuda 和 llvm。
TVM支持的硬件后端总览 在本教程中#xff0c;我们使用 cu…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 同样支持其他的模型格式如 MXNetCoreMLONNX 和 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_layers18, batch_sizebatch_size, image_shapeimage_shape
)# set show_meta_dataTrue if you want to show meta data
print(mod.astext(show_meta_dataFalse))输出
#[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, units1000) /* tyTensor[(1, 1000), float32] */;%89 nn.bias_add(%88, %fc1_bias, axis-1) /* tyTensor[(1, 1000), float32] */;nn.softmax(%89) /* tyTensor[(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_levelopt_level):lib relay.build(mod, target, paramsparams)输出
/home/areusch/ws/tvm3/python/tvm/target/target.py:259: UserWarning: Try specifying cuda arch by adding archsm_xx to your target.warnings.warn(Try specifying cuda arch by adding archsm_xx to your target.)运行生成的库
现在我们可以创建 graph executor 来将模块运行在 Nvidia GPU 上。
# create random input
dev tvm.cuda()
data np.random.uniform(-1, 1, sizedata_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]保存 / 加载编译好的模块
我们可以将 graphlib 和 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(datainput_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, atol1e-5)输出
[0.00089283 0.00103331 0.0009094 0.00102275 0.00108751 0.001067370.00106262 0.00095838 0.00110792 0.00113151]
