本文翻译自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还可以构建一个绑定的部署,其中运行时与模型组合在一个包中。