【TVM帮助文档学习】TVM及模型优化综述

本文翻译自Introduction — tvm 0.9.dev0 documentation

 下图说明了使用TVM优化编译器框架转换机器学习模型的步骤。

 

 1.从Tensorflow、PyTorch或Onnx等框架导入模型。导入层是TVM可以从其他框架(如Tensorflow、PyTorch或ONNX)中吸收模型的地方。TVM为每个前端提供的支持水平随着我们不断改进开源项目而提高。当导入模型到TVM有问题时，可以尝试将它先转换到ONNX。

2.转换为Relay。Relay是TVM的高级模型语言。导入TVM的模型用Relay表示。Relay是一种用于神经网络的函数式语言和中间表示。它支持:

• 传统的数据流样式表示
• 函数式的作用域，let-binding，这使得它成为一门功能完备的可微分语言
• 允许用户混合两种编程风格的能力

Relay采用图级优化pass对模型进行优化。

 3.低级化（lower）到张量表达式(TE)的表示。低级化是指将高级表示转换为低级表示。在应用高级优化之后，Relay运行FuseOps pass将模型划分为许多小的子图，并将子图低级化为TE表示。张量表达式(TE)是一种描述张量计算的领域特定语言。TE还提供了几个调度原语来优化低级循环，如平铺（tiling）、向量化(vectorization)、并行化(parallelization)、展开(unrolling)和融合(fusion)。TVM包含了一个张量运算符清单(TOPI)， 预定义了常见张量算子模板(例如conv2d，transpose等)，以帮助将Relay表示转换为TE表示

4.使用自动调优模块AutoTVM或AutoScheduler搜索最佳调度。调度为TE表示的算子或子图指定了低级循环优化。自动调优模块搜索最佳调度测率，并将其与成本模型和设备上的测量进行比较。TVM中有两个自动调优模块。

• AutoTVM:基于模板的自动调优模块。它运行搜索算法来查找用户定义模板中可调参数的最佳值。对于通用算子，TOPI中已经提供了它们的模板。
• AutoScheduler(又叫Ansor):一个无模板的自动调优模块。它不需要预定义的调度模板。相反，它通过分析计算定义自动生成搜索空间。然后在生成的搜索空间中搜索最佳调度策略。

5.为模型编译选择最佳配置。调优之后，自动调优模块生成JSON格式的调优记录。这个步骤为每个子图选择最佳的调度。

6.低级化到张量中间表达（TIR）。TIR是TVM的低级IR。根据调优步骤选择最佳优化配置后，每个TE子图转换为低级表达TIR，并使用低级pass行优化。接下来，优化后的TIR进一步低级化到硬件平台的目标编译器。这是最终的代码生成阶段，以生成可部署到生产环境中的优化模型。TVM支持几种不同的编译器后端，包括:

• LLVM，支持任意微处理器架构，包括标准x86和ARM处理器，AMDGPU和NVPTX代码生成，以及LLVM支持的任何其他平台。
• 专门的编译器，如 NVIDIA的编译器NVCC,。
• 通过TVM的自带代码(BYOC)框架实现的嵌入式和专门的目标。

 7.编译成机器码。在这个过程结束时，指定编译器产生的编码将被转换为成机器码。

TVM可以将模型编译成可链接的对象模块，然后可以与轻量级的TVM运行时一起运行。TVM运行时提供了C API来动态加载模型，并为其他语言(如Python和Rust)提供入口点。TVM还可以构建一个绑定的部署，其中运行时与模型组合在一个包中。