从一个例子看tvm执行流程
TVM整体流程(参考：TVM介绍)
机器学习模型在用TVM优化编译器框架进行变换时的步骤：
-
从Tensorflow/pytorch或ONNX等框架导入模型
import层是TVM从其他框架中导入模型的地方
注：TVM为每个前端提供的支持水平不尽相同，可尝试将模型转换为ONNX -
转换到Relay
Relay是TVM的高级模型语言，导入到TVM的模型是用Relay表示的。Relay是一种函数式语言(function language)和神经网络的中间表示法(intermediate representation,IR),它支持以下内容：
- 传统的数据流图式表示法
- Functional-style scoping 和 let-binding 使其成为一种功能齐全的可微分语言
- 能够允许用户混合两种编程风格
Relay应用图级(graph-level)优化passes来优化模型
- Lower to Tensor Expression (TE) representation
lower是指较高层表示转换成较低层表示。在high-level经过优化后，Relay 运行 FuseOps，将模型分割成许多小的子图，并将子图 lower 到 TE 表示。
TE 即Tensor Expression,张量表达，是描述张量计算的专属性语言
TE 还提供了一些schedule 原语来指定低级的循环优化，例如平铺（tiling）、矢量化（vectorization）、并行化（parallelization）、unrolling 和 fusion。
为了帮助将 Relay 表示转换为 TE 表示的过程，TVM 包含张量算子清单（Tensor Operator Inventory， TOPI），它有预先定义的常见张量算子的模板（如 conv2d、transpose）。
- Search for the best schedule using the auto-tuning module AutoTVM or AutoScheduler.
schedule 指定在 TE 中定义了算子或子图的低级循环优化。auto-tuning 模块搜索最佳 schedule 并将其与 cost 模型和设备上的测量结果进行比较。
在 TVM 中，有两个 auto-tuning 模块：
- AutoTVM: 基于模板的 auto-tuning 模块。它运行搜索算法为用户定义的模板中的可调节旋钮找到最佳值。对于常见的运算符，其模板已经在 TOPI 中提供。
- AutoScheduler （别名 Ansor） :无模板的auto-tuning 模块。它不需要预先定义的 schedule 模板。相反，它通过分析计算的定义自动生成搜索空间。然后，它在生成的搜索空间中搜索最佳 schedule。
-
Choose the optimal configurations for model compilation.
tuning 后，auto-tuning 模块会生成 JSON 格式的 auto-tuning 记录。这一步为每个子图挑选出最佳的 schedule。 -
Lower to Tensor Intermediate Representation (TIR)，TVM's low-level intermediate representation
TIR 是张量级的中间表示（Tensor Intermediate Representation），TVM 的低层次中间表示。
在根据 tuning 步骤选择最佳配置后，每个 TE 子图被降低到 TIR，并通过低级别的优化 passes 进行优化。
接下来，优化后的 TIR 被 lower 到硬件平台的目标编译器中。这是最后的代码生成阶段，产生可以部署到生产中的优化模型。
TVM 支持几种不同的编译器后端，包括：
- LVM：它可以针对任意的微处理器架构，包括 标准 x86 和 ARM 处理器，AMDGPU 和 NVPTX 代码生成，以及 LLVM 支持的任何其他平台。
- 专门的编译器，如 NVCC，NVIDIA 的编译器。
- 嵌入式和专用目标，通过 TVM 的 Bring Your Own Codegen（BYOC）框架实现。
- Compile down to machine code.
在这个过程结束时，特定的编译器生成的代码可以 lower 为机器码。
TVM 可以将模型编译成可链接的对象模块，然后可以用轻量级的 TVM 运行时来运行，该运行时提供 C 语言的 API 来动态加载模型，以及其他语言的入口，如 Python 和 Rust。TVM 还可以建立捆绑式部署，其中运行时与模型结合在一个包中。
例子 -- 编译Pytorch模型(参考：官网例子)
- 导入库
import tvm
from tvm import relay
import numpy as np
from tvm.contrib.download import download_testdata
# PyTorch imports
import torch
import torchvision
- 加载pytorch预训练模型
model_name = "resnet18"
model = getattr(torchvision.models, model_name)(pretrained=True)
model = model.eval()
# We grab the TorchScripted model via tracing
input_shape = [1, 3, 224, 224]
input_data = torch.randn(input_shape)
scripted_model = torch.jit.trace(model, input_data).eval()
- 加载一张测试图片
经典的猫图片
from PIL import Image
img_url = "https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true"
img_path = download_testdata(img_url, "cat.png", module="data")
img = Image.open(img_path).resize((224, 224))
# Preprocess the image and convert to tensor
from torchvision import transforms
my_preprocess = transforms.Compose(
[
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
)
img = my_preprocess(img)
img = np.expand_dims(img, 0)
- Import the graph to Relay
将Pytorch 图转换为Relay图，Input name是随意的
input_name = "input0"
shape_list = [(input_name, img.shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
- Relay Build
使用给定的输入将graph图编译成llvm目标
target = tvm.target.Target("llvm", host="llvm")
dev = tvm.cpu(0)
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target=target, params=params)
本地执行过程中，执行结果如下：
URLError(OSError(99, 'Cannot assign requested address'))
Download attempt 0/3 failed, retrying.
URLError(OSError(99, 'Cannot assign requested address'))
Download attempt 1/3 failed, retrying.
WARNING:root:Failed to download tophub package for llvm: <urlopen error [Errno 99] Cannot assign requested address>
/home/workspace/tvm/tvm/python/tvm/driver/build_module.py:267: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.
warnings.warn(
WARNING:autotvm:One or more operators have not been tuned. Please tune your model for better performance. Use DEBUG logging level to see more details.
虽有Error报错，但似乎并没有影响(此处也奇怪，为什么会有URLError的报错呢，编译的时候还需要请求什么吗？不清楚，也许看过源码才知道为什么吧)
根据官方的示例结果，应该如下：
/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. "
- TVM上执行可移植图(Execute the portable graph on TVM)
现在，尝试在目标上部署已编译的模型
from tvm.contrib import graph_executor
dtype = "float32"
m = graph_executor.GraphModule(lib["default"](dev))
# Set inputs
m.set_input(input_name, tvm.nd.array(img.astype(dtype)))
# Execute
m.run()
# Get outputs
tvm_output = m.get_output(0)
- 查找合集名称(Look up synset name)
在synset中查找预测的top1的索引
注：其中的imagenet_synsets.txt和imagenet_classes.txt两个文件，示例中给的url链接无效，可自行下载：
https://github.com/Cadene/pretrained-models.pytorch/blob/master/data/ 下的imagenet_synsets.txt和imagenet_classes.txt两个文件
亲测有效，可用
synset_url = "".join(
[
"https://raw.githubusercontent.com/Cadene/",
"pretrained-models.pytorch/master/data/",
"imagenet_synsets.txt",
]
)
synset_name = "imagenet_synsets.txt"
synset_path = download_testdata(synset_url, synset_name, module="data")
with open(synset_path) as f:
synsets = f.readlines()
synsets = [x.strip() for x in synsets]
splits = [line.split(" ") for line in synsets]
key_to_classname = {spl[0]: " ".join(spl[1:]) for spl in splits}
class_url = "".join(
[
"https://raw.githubusercontent.com/Cadene/",
"pretrained-models.pytorch/master/data/",
"imagenet_classes.txt",
]
)
class_name = "imagenet_classes.txt"
class_path = download_testdata(class_url, class_name, module="data")
with open(class_path) as f:
class_id_to_key = f.readlines()
class_id_to_key = [x.strip() for x in class_id_to_key]
# Get top-1 result for TVM
top1_tvm = np.argmax(tvm_output.numpy()[0])
tvm_class_key = class_id_to_key[top1_tvm]
# Convert input to PyTorch variable and get PyTorch result for comparison
with torch.no_grad():
torch_img = torch.from_numpy(img)
output = model(torch_img)
# Get top-1 result for PyTorch
top1_torch = np.argmax(output.numpy())
torch_class_key = class_id_to_key[top1_torch]
print("Relay top-1 id: {}, class name: {}".format(top1_tvm, key_to_classname[tvm_class_key]))
print("Torch top-1 id: {}, class name: {}".format(top1_torch, key_to_classname[torch_class_key]))
最后输出结果：
Relay top-1 id: 281, class name: tabby, tabby cat
Torch top-1 id: 281, class name: tabby, tabby cat
