【TVM 学习资料】快速入门:编译深度学习模型

Apache TVM: 开源深度学习编译器栈 Apache TVM是一个开源的深度学习编译器栈,旨在弥合生产力导向的深度学习框架与性能和效率导向的硬件后端之间的差距。它为各种硬件平台提供端到端的编译优化,以提高深度学习模型的性能和效率。TVM项目的愿景是汇集机器学习、编译器和系统架构领域的专家和从业者,共同构建一个可访问、可扩展和自动化的开源框架,以优化当前和新兴的机器学习模型在任何硬件平台上的性能。Apache TVM作为一个强大而灵活的深度学习编译器栈,正在为深度学习模型的高效部署和优化提供关键支持。 阅读详情

本篇文章译自英文文档 Quick Start Tutorial for Compiling Deep Learning Models 作者是 Yao Wang,Truman Tian。更多 TVM 中文文档可访问→ TVM 中文站

这个例子展示了如何用 Relay Python 前端构建神经网络,并为装有 TVM 的 NVIDIA GPU 生成 runtime 库。注意,构建 TVM 需要启用 CUDA 和 LLVM。

TVM 支持的硬件后端概述

下图显示了 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
)

# 想显示元数据则设置 show_meta_data=True
print(mod.astext(show_meta_data=False))

访问代码中的 mod.astext 更多信息

输出结果:

#[version = "0.0.5"]
def @main(%data: Tensor[(1, 3, 224, 224), float32] /* ty=Tensor[(1, 3, 224, 224), float32] */, %bn_data_gamma: Tensor[(3), float32] /* ty=Tensor[(3), float32] */, %bn_data_beta: Tensor[(3), float32] /* ty=Tensor[(3), float32] */, %bn_data_moving_mean: Tensor[(3), float32] /* ty=Tensor[(3), float32] */, %bn_data_moving_var: Tensor[(3), float32] /* ty=Tensor[(3), float32] */, %conv0_weight: Tensor[(64, 3, 7, 7), float32] /* ty=Tensor[(64, 3, 7, 7), float32] */, %bn0_gamma: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %bn0_beta: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %bn0_moving_mean: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %bn0_moving_var: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit1_bn1_gamma: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit1_bn1_beta: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit1_bn1_moving_mean: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit1_bn1_moving_var: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit1_conv1_weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */, %stage1_unit1_bn2_gamma: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit1_bn2_beta: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit1_bn2_moving_mean: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit1_bn2_moving_var: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit1_conv2_weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */, %stage1_unit1_sc_weight: Tensor[(64, 64, 1, 1), float32] /* ty=Tensor[(64, 64, 1, 1), float32] */, %stage1_unit2_bn1_gamma: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit2_bn1_beta: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit2_bn1_moving_mean: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit2_bn1_moving_var: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit2_conv1_weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */, %stage1_unit2_bn2_gamma: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit2_bn2_beta: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit2_bn2_moving_mean: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit2_bn2_moving_var: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage1_unit2_conv2_weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */, %stage2_unit1_bn1_gamma: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage2_unit1_bn1_beta: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage2_unit1_bn1_moving_mean: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage2_unit1_bn1_moving_var: Tensor[(64), float32] /* ty=Tensor[(64), float32] */, %stage2_unit1_conv1_weight: Tensor[(128, 64, 3, 3), float32] /* ty=Tensor[(128, 64, 3, 3), float32] */, %stage2_unit1_bn2_gamma: Tensor[(128), float32] /* ty=Tensor[(128), float32] */, %stage2_unit1_bn2_beta: Tensor[(128), float32] /* ty=Tensor[(128), float32] */, %stage2_unit1_bn2_moving_mean: Tensor[(128), float32] /* ty=Tensor[(128), float32] */, %stage2_unit1_bn2_moving_var: Tensor[(128), float32] /* ty=Tensor[(128), float32] */, %stage2_unit1_conv2_weight: Tensor[(128, 128, 3, 3), float32] /* ty=Tensor[(128, 128, 3, 3), float32] */, %stage2_unit1_sc_weight: Tensor[(128, 64, 1, 1), float32] /* ty=Tensor[(128, 64, 1, 1), float32] */, %stage2_unit2_bn1_gamma: Tensor[(128), float32] /* ty=Tensor[(128), float32] */, %stage2_unit2_bn1_beta: Tensor[(128), float32] /* ty=Tensor[(128), float32] */, %stage2_unit2_bn1_moving_mean: Tensor[(128), float32] /* ty=Tensor[(128), float32] */, %stage2_unit2_bn1_moving_var: Tensor[(128), float32] /* ty=Tensor[(128), float32] */, %stage2_unit2_conv1_weight: Tensor[(128, 128, 3, 3), float32] /* ty=Tensor[(128, 128, 3, 3), float32] */, %stage2_unit2_bn2_gamma: Tensor[(128), float32] /* ty=Tensor[(128), float32] */, %stage2_unit2_bn2_beta: Tensor[(128), float32] /* ty=Tensor[(128), float32] */, %stage2_unit2_bn2_moving_mean: Tensor[(128), float32] /* ty=Tensor[(128), float32] */, %stage2_unit2_bn2_moving_var: Tensor[(128), float32] /* ty=Tensor[(128), float32] */, %stage2_unit2_conv2_weight: Tensor[(128, 128, 3, 3), float32] /* ty=Tensor[(128, 128, 3, 3), float32] */, %stage3_unit1_bn1_gamma: Tensor[(128), float32] /* ty=Tensor[(128), float32] */, %stage3_unit1_bn1_beta: Tensor[(128), float32] /* ty=Tensor[(128), float32] */, %stage3_unit1_bn1_moving_mean: Tensor[(128), float32] /* ty=Tensor[(128), float32] */, %stage3_unit1_bn1_moving_var: Tensor[(128), float32] /* ty=Tensor[(128), float32] */, %stage3_unit1_conv1_weight: Tensor[(256, 128, 3, 3), float32] /* ty=Tensor[(256, 128, 3, 3), float32] */, %stage3_unit1_bn2_gamma: Tensor[(256), float32] /* ty=Tensor[(256), float32] */, %stage3_unit1_bn2_beta: Tensor[(256), float32] /* ty=Tensor[(256), float32] */, %stage3_unit1_bn2_moving_mean: Tensor[(256), float32] /* ty=Tensor[(256), float32] */, %stage3_unit1_bn2_moving_var: Tensor[(256), float32] /* ty=Tensor[(256), float32] */, %stage3_unit1_conv2_weight: Tensor[(256, 256, 3, 3), float32] /* ty=Tensor[(256, 256, 3, 3), float32] */, %stage3_unit1_sc_weight: Tensor[(256, 128, 1, 1), float32] /* ty=Tensor[(256, 128, 1, 1), float32] */, %stage3_unit2_bn1_gamma: Tensor[(256), float32] /* ty=Tensor[(256), float32] */, %stage3_unit2_bn1_beta: Tensor[(256), float32] /* ty=Tensor[(256), float32] */, %stage3_unit2_bn1_moving_mean: Tensor[(256), float32] /* ty=Tensor[(256), float32] */, %stage3_unit2_bn1_moving_var: Tensor[(256), float32] /* ty=Tensor[(256), float32] */, %stage3_unit2_conv1_weight: Tensor[(256, 256, 3, 3), float32] /* ty=Tensor[(256, 256, 3, 3), float32] */, %stage3_unit2_bn2_gamma: Tensor[(256), float32] /* ty=Tensor[(256), float32] */, %stage3_unit2_bn2_beta: Tensor[(256), float32] /* ty=Tensor[(256), float32] */, %stage3_unit2_bn2_moving_mean: Tensor[(256), float32] /* ty=Tensor[(256), float32] */, %stage3_unit2_bn2_moving_var: Tensor[(256), float32] /* ty=Tensor[(256), float32] */, %stage3_unit2_conv2_weight: Tensor[(256, 256, 3, 3), float32] /* ty=Tensor[(256, 256, 3, 3), float32] */, %stage4_unit1_bn1_gamma: Tensor[(256), float32] /* ty=Tensor[(256), float32] */, %stage4_unit1_bn1_beta: Tensor[(256), float32] /* ty=Tensor[(256), float32] */, %stage4_unit1_bn1_moving_mean: Tensor[(256), float32] /* ty=Tensor[(256), float32] */, %stage4_unit1_bn1_moving_var: Tensor[(256), float32] /* ty=Tensor[(256), float32] */, %stage4_unit1_conv1_weight: Tensor[(512, 256, 3, 3), float32] /* ty=Tensor[(512, 256, 3, 3), float32] */, %stage4_unit1_bn2_gamma: Tensor[(512), float32] /* ty=Tensor[(512), float32] */, %stage4_unit1_bn2_beta: Tensor[(512), float32] /* ty=Tensor[(512), float32] */, %stage4_unit1_bn2_moving_mean: Tensor[(512), float32] /* ty=Tensor[(512), float32] */, %stage4_unit1_bn2_moving_var: Tensor[(512), float32] /* ty=Tensor[(512), float32] */, %stage4_unit1_conv2_weight: Tensor[(512, 512, 3, 3), float32] /* ty=Tensor[(512, 512, 3, 3), float32] */, %stage4_unit1_sc_weight: Tensor[(512, 256, 1, 1), float32] /* ty=Tensor[(512, 256, 1, 1), float32] */, %stage4_unit2_bn1_gamma: Tensor[(512), float32] /* ty=Tensor[(512), float32] */, %stage4_unit2_bn1_beta: Tensor[(512), float32] /* ty=Tensor[(512), float32] */, %stage4_unit2_bn1_moving_mean: Tensor[(512), float32] /* ty=Tensor[(512), float32] */, %stage4_unit2_bn1_moving_var: Tensor[(512), float32] /* ty=Tensor[(512), float32] */, %stage4_unit2_conv1_weight: Tensor[(512, 512, 3, 3), float32] /* ty=Tensor[(512, 512, 3, 3), float32] */, %stage4_unit2_bn2_gamma: Tensor[(512), float32] /* ty=Tensor[(512), float32] */, %stage4_unit2_bn2_beta: Tensor[(512), float32] /* ty=Tensor[(512), float32] */, %stage4_unit2_bn2_moving_mean: Tensor[(512), float32] /* ty=Tensor[(512), float32] */, %stage4_unit2_bn2_moving_var: Tensor[(512), float32] /* ty=Tensor[(512), float32] */, %stage4_unit2_conv2_weight: Tensor[(512, 512, 3, 3), float32] /* ty=Tensor[(512, 512, 3, 3), float32] */, %bn1_gamma: Tensor[(512), float32] /* ty=Tensor[(512), float32] */, %bn1_beta: Tensor[(512), float32] /* ty=Tensor[(512), float32] */, %bn1_moving_mean: Tensor[(512), float32] /* ty=Tensor[(512), float32] */, %bn1_moving_var: Tensor[(512), float32] /* ty=Tensor[(512), float32] */, %fc1_weight: Tensor[(1000, 512), float32] /* ty=Tensor[(1000, 512), float32] */, %fc1_bias: Tensor[(1000), float32] /* ty=Tensor[(1000), float32] */) -> Tensor[(1, 1000), float32] {
  %0 = nn.batch_norm(%data, %bn_data_gamma, %bn_data_beta, %bn_data_moving_mean, %bn_data_moving_var, epsilon=2e-05f, scale=False) /* ty=(Tensor[(1, 3, 224, 224), float32], Tensor[(3), float32], Tensor[(3), float32]) */;
  %1 = %0.0 /* ty=Tensor[(1, 3, 224, 224), float32] */;
  %2 = nn.conv2d(%1, %conv0_weight, strides=[2, 2], padding=[3, 3, 3, 3], channels=64, kernel_size=[7, 7]) /* ty=Tensor[(1, 64, 112, 112), float32] */;
  %3 = nn.batch_norm(%2, %bn0_gamma, %bn0_beta, %bn0_moving_mean, %bn0_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 64, 112, 112), float32], Tensor[(64), float32], Tensor[(64), float32]) */;
  %4 = %3.0 /* ty=Tensor[(1, 64, 112, 112), float32] */;
  %5 = nn.relu(%4) /* ty=Tensor[(1, 64, 112, 112), float32] */;
  %6 = nn.max_pool2d(%5, pool_size=[3, 3], strides=[2, 2], padding=[1, 1, 1, 1]) /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %7 = nn.batch_norm(%6, %stage1_unit1_bn1_gamma, %stage1_unit1_bn1_beta, %stage1_unit1_bn1_moving_mean, %stage1_unit1_bn1_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 64, 56, 56), float32], Tensor[(64), float32], Tensor[(64), float32]) */;
  %8 = %7.0 /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %9 = nn.relu(%8) /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %10 = nn.conv2d(%9, %stage1_unit1_conv1_weight, padding=[1, 1, 1, 1], channels=64, kernel_size=[3, 3]) /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %11 = nn.batch_norm(%10, %stage1_unit1_bn2_gamma, %stage1_unit1_bn2_beta, %stage1_unit1_bn2_moving_mean, %stage1_unit1_bn2_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 64, 56, 56), float32], Tensor[(64), float32], Tensor[(64), float32]) */;
  %12 = %11.0 /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %13 = nn.relu(%12) /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %14 = nn.conv2d(%13, %stage1_unit1_conv2_weight, padding=[1, 1, 1, 1], channels=64, kernel_size=[3, 3]) /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %15 = nn.conv2d(%9, %stage1_unit1_sc_weight, padding=[0, 0, 0, 0], channels=64, kernel_size=[1, 1]) /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %16 = add(%14, %15) /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %17 = nn.batch_norm(%16, %stage1_unit2_bn1_gamma, %stage1_unit2_bn1_beta, %stage1_unit2_bn1_moving_mean, %stage1_unit2_bn1_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 64, 56, 56), float32], Tensor[(64), float32], Tensor[(64), float32]) */;
  %18 = %17.0 /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %19 = nn.relu(%18) /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %20 = nn.conv2d(%19, %stage1_unit2_conv1_weight, padding=[1, 1, 1, 1], channels=64, kernel_size=[3, 3]) /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %21 = nn.batch_norm(%20, %stage1_unit2_bn2_gamma, %stage1_unit2_bn2_beta, %stage1_unit2_bn2_moving_mean, %stage1_unit2_bn2_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 64, 56, 56), float32], Tensor[(64), float32], Tensor[(64), float32]) */;
  %22 = %21.0 /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %23 = nn.relu(%22) /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %24 = nn.conv2d(%23, %stage1_unit2_conv2_weight, padding=[1, 1, 1, 1], channels=64, kernel_size=[3, 3]) /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %25 = add(%24, %16) /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %26 = nn.batch_norm(%25, %stage2_unit1_bn1_gamma, %stage2_unit1_bn1_beta, %stage2_unit1_bn1_moving_mean, %stage2_unit1_bn1_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 64, 56, 56), float32], Tensor[(64), float32], Tensor[(64), float32]) */;
  %27 = %26.0 /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %28 = nn.relu(%27) /* ty=Tensor[(1, 64, 56, 56), float32] */;
  %29 = nn.conv2d(%28, %stage2_unit1_conv1_weight, strides=[2, 2], padding=[1, 1, 1, 1], channels=128, kernel_size=[3, 3]) /* ty=Tensor[(1, 128, 28, 28), float32] */;
  %30 = nn.batch_norm(%29, %stage2_unit1_bn2_gamma, %stage2_unit1_bn2_beta, %stage2_unit1_bn2_moving_mean, %stage2_unit1_bn2_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 128, 28, 28), float32], Tensor[(128), float32], Tensor[(128), float32]) */;
  %31 = %30.0 /* ty=Tensor[(1, 128, 28, 28), float32] */;
  %32 = nn.relu(%31) /* ty=Tensor[(1, 128, 28, 28), float32] */;
  %33 = nn.conv2d(%32, %stage2_unit1_conv2_weight, padding=[1, 1, 1, 1], channels=128, kernel_size=[3, 3]) /* ty=Tensor[(1, 128, 28, 28), float32] */;
  %34 = nn.conv2d(%28, %stage2_unit1_sc_weight, strides=[2, 2], padding=[0, 0, 0, 0], channels=128, kernel_size=[1, 1]) /* ty=Tensor[(1, 128, 28, 28), float32] */;
  %35 = add(%33, %34) /* ty=Tensor[(1, 128, 28, 28), float32] */;
  %36 = nn.batch_norm(%35, %stage2_unit2_bn1_gamma, %stage2_unit2_bn1_beta, %stage2_unit2_bn1_moving_mean, %stage2_unit2_bn1_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 128, 28, 28), float32], Tensor[(128), float32], Tensor[(128), float32]) */;
  %37 = %36.0 /* ty=Tensor[(1, 128, 28, 28), float32] */;
  %38 = nn.relu(%37) /* ty=Tensor[(1, 128, 28, 28), float32] */;
  %39 = nn.conv2d(%38, %stage2_unit2_conv1_weight, padding=[1, 1, 1, 1], channels=128, kernel_size=[3, 3]) /* ty=Tensor[(1, 128, 28, 28), float32] */;
  %40 = nn.batch_norm(%39, %stage2_unit2_bn2_gamma, %stage2_unit2_bn2_beta, %stage2_unit2_bn2_moving_mean, %stage2_unit2_bn2_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 128, 28, 28), float32], Tensor[(128), float32], Tensor[(128), float32]) */;
  %41 = %40.0 /* ty=Tensor[(1, 128, 28, 28), float32] */;
  %42 = nn.relu(%41) /* ty=Tensor[(1, 128, 28, 28), float32] */;
  %43 = nn.conv2d(%42, %stage2_unit2_conv2_weight, padding=[1, 1, 1, 1], channels=128, kernel_size=[3, 3]) /* ty=Tensor[(1, 128, 28, 28), float32] */;
  %44 = add(%43, %35) /* ty=Tensor[(1, 128, 28, 28), float32] */;
  %45 = nn.batch_norm(%44, %stage3_unit1_bn1_gamma, %stage3_unit1_bn1_beta, %stage3_unit1_bn1_moving_mean, %stage3_unit1_bn1_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 128, 28, 28), float32], Tensor[(128), float32], Tensor[(128), float32]) */;
  %46 = %45.0 /* ty=Tensor[(1, 128, 28, 28), float32] */;
  %47 = nn.relu(%46) /* ty=Tensor[(1, 128, 28, 28), float32] */;
  %48 = nn.conv2d(%47, %stage3_unit1_conv1_weight, strides=[2, 2], padding=[1, 1, 1, 1], channels=256, kernel_size=[3, 3]) /* ty=Tensor[(1, 256, 14, 14), float32] */;
  %49 = nn.batch_norm(%48, %stage3_unit1_bn2_gamma, %stage3_unit1_bn2_beta, %stage3_unit1_bn2_moving_mean, %stage3_unit1_bn2_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 256, 14, 14), float32], Tensor[(256), float32], Tensor[(256), float32]) */;
  %50 = %49.0 /* ty=Tensor[(1, 256, 14, 14), float32] */;
  %51 = nn.relu(%50) /* ty=Tensor[(1, 256, 14, 14), float32] */;
  %52 = nn.conv2d(%51, %stage3_unit1_conv2_weight, padding=[1, 1, 1, 1], channels=256, kernel_size=[3, 3]) /* ty=Tensor[(1, 256, 14, 14), float32] */;
  %53 = nn.conv2d(%47, %stage3_unit1_sc_weight, strides=[2, 2], padding=[0, 0, 0, 0], channels=256, kernel_size=[1, 1]) /* ty=Tensor[(1, 256, 14, 14), float32] */;
  %54 = add(%52, %53) /* ty=Tensor[(1, 256, 14, 14), float32] */;
  %55 = nn.batch_norm(%54, %stage3_unit2_bn1_gamma, %stage3_unit2_bn1_beta, %stage3_unit2_bn1_moving_mean, %stage3_unit2_bn1_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 256, 14, 14), float32], Tensor[(256), float32], Tensor[(256), float32]) */;
  %56 = %55.0 /* ty=Tensor[(1, 256, 14, 14), float32] */;
  %57 = nn.relu(%56) /* ty=Tensor[(1, 256, 14, 14), float32] */;
  %58 = nn.conv2d(%57, %stage3_unit2_conv1_weight, padding=[1, 1, 1, 1], channels=256, kernel_size=[3, 3]) /* ty=Tensor[(1, 256, 14, 14), float32] */;
  %59 = nn.batch_norm(%58, %stage3_unit2_bn2_gamma, %stage3_unit2_bn2_beta, %stage3_unit2_bn2_moving_mean, %stage3_unit2_bn2_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 256, 14, 14), float32], Tensor[(256), float32], Tensor[(256), float32]) */;
  %60 = %59.0 /* ty=Tensor[(1, 256, 14, 14), float32] */;
  %61 = nn.relu(%60) /* ty=Tensor[(1, 256, 14, 14), float32] */;
  %62 = nn.conv2d(%61, %stage3_unit2_conv2_weight, padding=[1, 1, 1, 1], channels=256, kernel_size=[3, 3]) /* ty=Tensor[(1, 256, 14, 14), float32] */;
  %63 = add(%62, %54) /* ty=Tensor[(1, 256, 14, 14), float32] */;
  %64 = nn.batch_norm(%63, %stage4_unit1_bn1_gamma, %stage4_unit1_bn1_beta, %stage4_unit1_bn1_moving_mean, %stage4_unit1_bn1_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 256, 14, 14), float32], Tensor[(256), float32], Tensor[(256), float32]) */;
  %65 = %64.0 /* ty=Tensor[(1, 256, 14, 14), float32] */;
  %66 = nn.relu(%65) /* ty=Tensor[(1, 256, 14, 14), float32] */;
  %67 = nn.conv2d(%66, %stage4_unit1_conv1_weight, strides=[2, 2], padding=[1, 1, 1, 1], channels=512, kernel_size=[3, 3]) /* ty=Tensor[(1, 512, 7, 7), float32] */;
  %68 = nn.batch_norm(%67, %stage4_unit1_bn2_gamma, %stage4_unit1_bn2_beta, %stage4_unit1_bn2_moving_mean, %stage4_unit1_bn2_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 512, 7, 7), float32], Tensor[(512), float32], Tensor[(512), float32]) */;
  %69 = %68.0 /* ty=Tensor[(1, 512, 7, 7), float32] */;
  %70 = nn.relu(%69) /* ty=Tensor[(1, 512, 7, 7), float32] */;
  %71 = nn.conv2d(%70, %stage4_unit1_conv2_weight, padding=[1, 1, 1, 1], channels=512, kernel_size=[3, 3]) /* ty=Tensor[(1, 512, 7, 7), float32] */;
  %72 = nn.conv2d(%66, %stage4_unit1_sc_weight, strides=[2, 2], padding=[0, 0, 0, 0], channels=512, kernel_size=[1, 1]) /* ty=Tensor[(1, 512, 7, 7), float32] */;
  %73 = add(%71, %72) /* ty=Tensor[(1, 512, 7, 7), float32] */;
  %74 = nn.batch_norm(%73, %stage4_unit2_bn1_gamma, %stage4_unit2_bn1_beta, %stage4_unit2_bn1_moving_mean, %stage4_unit2_bn1_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 512, 7, 7), float32], Tensor[(512), float32], Tensor[(512), float32]) */;
  %75 = %74.0 /* ty=Tensor[(1, 512, 7, 7), float32] */;
  %76 = nn.relu(%75) /* ty=Tensor[(1, 512, 7, 7), float32] */;
  %77 = nn.conv2d(%76, %stage4_unit2_conv1_weight, padding=[1, 1, 1, 1], channels=512, kernel_size=[3, 3]) /* ty=Tensor[(1, 512, 7, 7), float32] */;
  %78 = nn.batch_norm(%77, %stage4_unit2_bn2_gamma, %stage4_unit2_bn2_beta, %stage4_unit2_bn2_moving_mean, %stage4_unit2_bn2_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 512, 7, 7), float32], Tensor[(512), float32], Tensor[(512), float32]) */;
  %79 = %78.0 /* ty=Tensor[(1, 512, 7, 7), float32] */;
  %80 = nn.relu(%79) /* ty=Tensor[(1, 512, 7, 7), float32] */;
  %81 = nn.conv2d(%80, %stage4_unit2_conv2_weight, padding=[1, 1, 1, 1], channels=512, kernel_size=[3, 3]) /* ty=Tensor[(1, 512, 7, 7), float32] */;
  %82 = add(%81, %73) /* ty=Tensor[(1, 512, 7, 7), float32] */;
  %83 = nn.batch_norm(%82, %bn1_gamma, %bn1_beta, %bn1_moving_mean, %bn1_moving_var, epsilon=2e-05f) /* ty=(Tensor[(1, 512, 7, 7), float32], Tensor[(512), float32], Tensor[(512), float32]) */;
  %84 = %83.0 /* ty=Tensor[(1, 512, 7, 7), float32] */;
  %85 = nn.relu(%84) /* ty=Tensor[(1, 512, 7, 7), float32] */;
  %86 = nn.global_avg_pool2d(%85) /* ty=Tensor[(1, 512, 1, 1), float32] */;
  %87 = nn.batch_flatten(%86) /* ty=Tensor[(1, 512), float32] */;
  %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 模块库,以及模型的 blobs 参数。在编译过程中,Relay 进行图级优化,而 TVM 进行张量级优化,从而为模型服务提供优化的 runtime 模块。

首先为 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)

输出结果:

/workspace/python/tvm/target/target.py:377: 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.")
/workspace/python/tvm/driver/build_module.py:268: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.
  "target_host parameter is going to be deprecated. "

运行生成库

创建图执行器,然后在 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]

保存和加载编译模块

还可将计算图、库和参数保存到文件中,然后在部署环境中加载。

# 分别将计算图、库和参数保存到不同文件
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']
# 重新加载模块
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(out_deploy.flatten()[0:10])

# 检查来自部署模块的输出和原始输出是否一致
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]

下载 Python 源代码:relay_quick_start.py
下载 Jupyter Notebook:relay_quick_start.ipynb

探索TVM:深度学习模型编译与部署的前沿与实践【文末送书】 TVM是一个开源的机器学习编译器栈,旨在将深度学习模型部署到各种硬件设备上,包括CPU、GPU和各种加速器。 TVM采用模块化设计,将编译器分为前端、中端和后端三个部分。前端负责将深度学习模型从各种框架(如TensorFlow、PyTorch等)转换为TVM的中间表示(Intermediate Representation,IR),中端负责在中间表示上进行一系列的优化,如算子融合、减枝等,以减少计算量并提高运行速度,后端则负责将优化后的中间表示转换为针对特定硬件设备的低级代码。 阅读详情

相关推荐

TVM——快速部署模型的终极方案

作者:禅与计算机程序设计艺术 1.简介 随着人工智能(AI)技术的发展,模型的迅速部署已经成为一种必然。在深度学习模型快速落地到商业落地场景中,如何将高效且准确的推理部署到生产环境并保证服务质量一直是一个重要的课题。本文将通过对TVM工具进行阐述,结合相关开源框架的功能特性和实际案例,对当前的模型部署方

AI天才研究院 2717

TVM 学习指南(个人版)

最近粗略的看完了天奇大佬的MLC课程(顺便修了一些语法和拼写错误,也算是做了微弱的贡献hh),对TVM的近期发展有了一些新的认识。之前天奇大佬在《新一代深度学习编译技术变革和展望》一文中(链接:https://zhuanlan.zhihu.com/p/446935289)讲解了TVM Unify也即统一多层抽象的概念。这里的统一多层抽象具体包括AutoTensorization用来解决硬件指令声明和张量程序对接,TVM FFI(PackedFunc)机制使得我们可以灵活地引入任意的算子库和运行库函数并且在各

I good vegetable a! 8520

【从零开始学深度学习编译器】一,深度学习编译器及TVM 介绍

0x1. 为什么需要深度学习编译器? 深度学习编译器这个词语,我们可以先拆成两个部分来看。 首先谈谈深度学习领域。从训练框架角度来看,Google的TensorFlow和FaceBook的Pytorch是全球主流的深度学习框架,另外亚马逊的MxNet,百度的Paddle,旷视的MegEngine,华为的Mindspore以及一流科技的OneFlow也逐渐在被更多人接受和使用。这么多训练框架,我们究竟应该选择哪个?如果追求易用性,可能你会选择Pytorch,如果追求项目部署落地,可能你会选择TensorFlo

qq_33287871的博客 1173

TVM 从入门到精通 | 安装 TVM (Part 1)

手把手教学之从源码安装 TVM

HyperAI超神经 3203

TVM:端到端自动深度学习编译器(244页PPT)

陈天奇给了关于TVM的报告,TVM: An End to End Automated Deep Learning Compiler。

J-Flash手动添加华大HC32F460型号全攻略

在嵌入式开发中,Flash编程算法是连接调试工具与微控制器的重要桥梁,它定义了芯片存储器的擦除、编程和校验等底层操作规范。其原理是通过调试接口将算法代码加载到目标芯片的RAM中执行,从而实现对Flash存储器的访问与控制。这项技术的核心价值在于打破了工具链对芯片型号的固有支持限制,使得开发者能够灵活适配各种新兴或定制化的微控制器。在工程实践中,当面对J-Link等主流调试器官方支持包更新滞后时,手动添加芯片型号成为保障开发效率的关键技能。本文以华大半导体的HC32F460系列MCU为例,详细解析了如何通过配

weixin_29192365的博客 353

初识 TVM

如有图像或公式显示错误,可以访问我的个人博客:https://www.wanglichun.tech/2019/11/15/tvm/ 笔者也是最近偶然的机会才开始接触TVM,使用过后发现,经过auto-tuning后的TVM模型在速度是竟然超过了TensorRT,并且笔者使用的是MXNet框架,TVM对MXNet绝对的友好,对于Pytorch等模型,可以使用ONNX,操作也一样简单,使用起来基本类...

小菜鸟的AI之路 7152

疾病负担研究(GBD)介绍

今天我们来聊一聊第一个公共数据库——GBD数据库,全称叫做global burden of disease(全球疾病负担),该数据库每年都会更新,目前已经更新到GBD2019。 这里直接给大家贴出网址:Global Burden of Disease Study 2019 (GBD 2019) Data Resources | GHDx 上图中红色圈圈标识的就是数据库的数据获取界面的链接,大家直接点击即可. 数据的检索十分方便: Context下拉菜单中可以选择不同的选项,默认是caus

weixin_44240228的博客 1万+

tvm学习笔记(五):tvm工作原理

一、总体流程: TVM的工作流程:首先,将网络表示成统一的表示形式(Intermediate Representation),并进行一些可重用的图优化;然后,利用不同的后端生成对应设备代码,如图1所示。 图1 tvm 工...

sinat_31425585的博客 1万+

pytorch中dataloader的num_workers参数

结论速递 在Windows系统中,num_workers参数建议设为0,在Linux系统则不需担心。 1 问题描述 在之前的任务超大图上的节点表征学习中,使用PyG库用于数据加载的DataLoader,ClusterLoader等类时,会涉及到参数num_workers。在不同系统环境试验该参数的设定值时,会出现不同的结果 colab(Linux),设为12 有warning message,提示num_workers参数设置不合理,不影响继续运行 /usr/local/lib/python3.7/di

qq_40990057的博客 1万+

智能排班考勤系统原型

可实现从业务预测,到工时预估,再到人员分配全流程自动智能排班。可应用于客服行业或工厂的智能排班系统。

gywm15111336483的专栏 996

【从零开始学深度学习编译器】十,TVM的整体把握

大概4个月前开始接触TVM,虽然是0经验开始看,但目前对TVM的IR以及Pass,Codegen,Scheduler等也有了一些基础的认识。所以这篇文章的目的是梳理一下TVM的整体架构,复盘一下自己学到了什么,以及为想学习TVM的小伙伴们提供一个整体思路。这个专题的文章和实验代码都被我汇总放到了https://github.com/BBuf/tvm_learn这个仓库中,当然是希望,感激不尽。仓库目录如下:TVM 学习笔记目录页。

生活需要深度 1242

05、基于ALV表报实现复选框、全选、取消全选等功能

基于Function ALV实现复选框、全选及取消全选等功能

wo520yl1314的博客 2372

【从零开始学深度学习编译器】一,深度学习编译器及TVM 介绍

0x0. 介绍 大家好呀,在过去的半年到一年时间里,我分享了一些算法解读,算法优化,模型转换相关的一些文章。这篇文章是自己开启学习深度学习编译器的第一篇文章,后续也会努力更新这个系列。这篇文章是开篇,所以我不会具体讲解TVM的知识,更多的是介绍一下我自己眼中的深度学习编译器是什么?以及为什么我要选择学习TVM,最后我也会给出一个让读者快速体验TVM效果的一个开发环境搭建的简要教程。 0x1. 为什么需要深度学习编译器? 深度学习编译器这个词语,我们可以先拆成两个部分来看。 首先谈谈深度学习领域。从训练框架角

I good vegetable a! 4091

9.16华为OD机试真题 新系统 - 矩阵螺旋遍历 (Java/Py/C/C++/Js/Go)

9.16华为OD机试真题 新系统 - 矩阵螺旋遍历 (Java/Py/C/C++/Js/Go)

算法大师! 657

图像处理之边缘检测[微分算子、Canny算子和LOG算子]

LOG算子是在经典算子的基础上发展起来的边缘检测算子,根据信噪比求得检测边缘的最优滤波器。边缘总是以强度突变的形式出现,可以定义为图像局部特性的不连续性,如灰度的突变和纹理结构的突变等。图像的边缘包含了物体形状的重要信息,它不仅在分析图像时大幅度地减少了要处理的信息量,而且还保护了目标的边界结构。对于复杂的图像,采用Roberts算子不能较好地得到图像的边缘,而需要采用更加复杂的3X3的算子比如Prewitt算子。对于离散的图像f(x,y),边缘检测算子是用图像的垂直和水平差分来逼近梯度算子,即。

qq_44111805的博客 9839

【TVM系列教程一】深度学习编译器及TVM 介绍

0x0. 介绍大家好呀,在过去的半年到一年时间里,我分享了一些算法解读,算法优化,模型转换相关的一些文章。这篇文章是自己开启学习深度学习编译器的第一篇文章,后续也会努力更新这个系列。这篇文...

算法码上来 2231

多智能体系统在AI物流优化中的实际应用案例

订单波峰波谷难应对(如双11订单量暴增300%);动态干扰难处理(如突发交通拥堵、客户临时改地址);全局优化难实现(单辆车最优≠整体最优)。本文聚焦“多智能体系统(MAS)”这一AI技术,讲解其如何通过“群体智能”破解物流调度难题,并结合实际案例展示落地效果。本文从“生活故事→核心概念→技术原理→实战案例→未来趋势”逐步展开,用“快递员团队协作送外卖”类比MAS运作,结合Python代码模拟任务分配,最后通过某物流企业的真实项目验证效果。

AIGC应用创新大全的博客 1632

2024三掌柜赠书活动第一期:TVM编译器原理与实践

随着人工智能的发展,计算机视觉、自然语言处理和语音识别等领域的需求不断增加。为了更好地满足这些需求,许多深度学习框架被开发出来,其中TVM(TVirtual Machine)是一种优秀的编译器,能够将深度学习模型编译为高效的机器码。而且TVM编译器的核心思想,就是将深度学习模型转化为高效的计算图,并优化图中的计算节点。这样一来,模型运行时的计算时间就会大大减少,同时还可以提高模型的功耗效率。TVM编译器的实现过程可以分为三个主要部分:前端、中间层和后端。

全沾软贱开发攻城狮 1万+

Statistics and Machine Learning Toolbox

Statistics and Machine Learning Toolbox(带书签),MATLAB统计学习工具

精雕软件(JDpaint) 5.5 绿色版

精雕软件(JDpaint) 5.5 绿色版精雕软件(JDpaint) 5.5 绿色版精雕软件(JDpaint) 5.5 绿色版

TVM Android部署终极指南:快速集成深度学习模型到移动应用

TVM(Tensor Virtual Machine)是一个开源的深度学习编译器栈,专门用于优化和部署深度学习模型到各种硬件平台。本文将为您详细介绍如何将TVM集成到Android应用中,实现高效的深度学习模型部署。🎯 ## 为什么选择TVM进行Android部署? TVM提供了跨平台的深度学习模型编译和优化能力,特别适合移动端部署: - **高性能推理**:通过自动调优和优化,实现最佳的

gitblog_00981的博客 972

visio安装包

Microsoft Office Visio 是一款专业的制图软件。利用经过更新的形状、协作工具和数据链接图表,创建专业级图表以简化复杂的信息。

AI编译器入门:TVM与MLIR如何将深度学习模型映射到任意硬件——连接算法与芯片的桥梁

AI编译器是连接深度学习模型与多样化硬件架构的核心桥梁,它通过多层中间表示和自动化优化技术,解决模型部署中的性能与可移植性挑战。本文深入解析TVM和MLIR两大主流框架的设计哲学与实现原理,涵盖从计算图优化、算子融合到硬件特定代码生成的完整编译流程。通过架构分析、代码实例与实战对比,系统阐述如何将高级模型描述高效映射到CPU、GPU、ASIC等异构计算平台,为读者提供从理论到实践的全面指南。

AladdinEdu,你的AI学习实践工作坊。让想法落地,让研究加速。助力高校AI人才成长,点亮创新未来。 2199

win10 跳过U8环境检测的IE WebControls.rar

win10 安装u8的时候 提示IE WebControls 不通过,这个工具可以欺骗u8的检查。 跳过U8环境检测的IE WebControls。提示导入失败后,重新检测通过。

TVM 中文教程:从入门到精通的深度学习编译器指南

TVM是一个开源的深度学习编译器堆栈,它能够将深度学习模型高效地部署到各种硬件后端。作为深度学习领域的重要基础设施,TVM解决了"一次训练,到处部署"的关键挑战。 ## 为什么需要学习TVM 在深度学习应用日益普及的今天,模型部署面临着硬件多样性带来的巨大挑战。不同硬件架构(CPU、GPU、TPU、FPGA等)有着不同的计算特性,而TVM正是解决这一问题的利器。通过学习TVM,开发者可以:

gitblog_00662的博客 886

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 impor

weixin_44966641的博客 2278

Apache TVM:开源深度学习编译器栈的领跑者

Apache TVM:开源深度学习编译器栈的领跑者 【免费下载链接】tvm Open deep learning compiler stack for cpu, gpu and specialized accelerators 项...

gitblog_00753的博客 967

AI编译器与TVM

由于AI芯片的特殊性和高度定制化,为了兼容硬件的多样性,AI模型必须能被高效地映射到各种AI芯片上。AI编译器将深度学习框架描述的AI模型作为输入,将为各种AI芯片生成的优化代码作为输出。AI编译器的目标是通过编译优化的方法将深度学习框架产生的AI模型转化为与特定架构的AI芯片适配的可执行机器码。主流的AI硬件厂商都提供了高度优化的算子库或者推理引擎来实现高效计算,但是依赖算子库的缺点是算子库覆盖的硬件范围有限,而且库函数的更新落后于模型的发展速度,无法充分发挥AI芯片算力。

欢迎订阅 -- 小源笔迹 1603
上一篇: 线上直播 | 聊聊今年的 TVM Community Keynote
下一篇: 姜汉:使用 TVM 做边缘设备的模型编译和推理实践
HyperAI超神经
博客等级 码龄7年 8873粉丝 1753原创
评论
成就一亿技术人!
拼手气红包6.0元
还能输入1000个字符
 
 条评论被折叠 查看
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付:10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值