日期: 2023 年 5 月 23 日

深入浅出 tvm – (3) 架构 & 设计

Posted on by remaper in 深入浅出 tvm Tagged:
本文将从一个最简单的图像分类模型，了解基本的推理过程，并以此来了解一下 tvm 编译器的基本工作流程，以及对应核心模块的职责
0

1. 模型的定义

由于tvm支持onnx模型格式，所以我们从 resnet50-v2-7.onnx 开始
模型的输出：https://s3.amazonaws.com/onnx-model-zoo/synset.txt，其实就是1000个图像分类
当我们执行模型推理时，输入一个图像，输出是一个 shape=(1,1000) 的数组，表示这个图片对应1000个分类的概率，其中概率最大的那个分类，就是预测的结果

2. 从一个最基本的推理开始

推理的过程就是把模型，编译成目标设备上可运行的代码，根据输入数据，返回预测结果
如下，说下几个关键的地方：
  1. 模型转换：relay.frontend.from_onnx()，把模型从 onnx 格式转换成 relay IRModule 格式
  2. 编译：relay.build()，完成从IRModule到目标设备代码的编译优化，内部包含了 pass 优化、schedule优化
  3. 推理：输入 module.set_input()，运行 module.run()，获取结果 module.get_output()
其中最核心的就是2这几个步骤，基本包括了tvm最核心的逻辑
# https://github.com/onnx/models/raw/main/vision/classification/resnet/model/resnet50-v2-7.onnx
onnx_model = onnx.load('./resnet50-v2-7.onnx')
img_data = get_test_img()
input_name = "data"
shape_dict = {input_name: img_data.shape}

mod, params = relay.frontend.from_onnx(onnx_model, shape_dict)

with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target="llvm", params=params)

module = graph_executor.GraphModule(lib["default"](tvm.device("llvm", 0)))

module.set_input(input_name, tvm.nd.array(img_data.astype('float32')))
module.run()
"https://s3.amazonaws.com/onnx-model-zoo/synset.txt"
out = module.get_output(0, tvm.nd.empty((1, 1000))).numpy()
我们再来看下 relay.build 具体干了啥：
  1. python/tvm/relay/build_module.py: build()
    1. python/tvm/relay/build_module.py:BuildModule::build()
      1. src/relay/backend/build_module.cc: BuildRelay()
BuildRelay 干的事情就是：Compile a Relay IR module to runtime module