最近在学习一些和芯片 AI相关的知识,重点了解了一下TVM,我自己认为TVM在AI应用落地类似的项目中,用途还是非常广泛的,现在把一些重要的笔记贴在下面,有两篇原帖链接也附上,感兴趣的同学可以学习一下。
TVM
1.什么是TVM?
TVM(Tensor Virtual Machine)是一个开源的机器学习编译器框架,由华盛顿大学陈天奇博士主导开发。它旨在为各种硬件后端提供统一的编程模型,使得开发者能够便捷地将深度学习模型编译成不同硬件平台可执行的代码。TVM 支持多种深度学习框架(如 TensorFlow、PyTorch、ONNX 等)的模型导入,并提供了一系列工具来优化和部署模型。
总的来说,TVM 是一个功能强大、灵活且高效的机器学习编译器框架,它通过提供一系列的工具和优化策略,帮助开发者在不同的硬件平台上高效地部署深度学习模型。
2.TVM 的核心特性包括:
多语言支持:TVM 支持多种硬件描述语言,如 CUDA、OpenCL、Metal 等,以及传统的编程语言如 C++。
自动优化:TVM 提供了自动调优工具(如 AutoTVM 和 AutoScheduler),这些工具可以通过搜索最优的计算配置来加速模型的执行。
中间表示(IR):TVM 使用 Relay 作为高级 IR,用于表示和优化计算图。此外,TVM 还引入了 Tensor IR(TIR)作为低级的 IR,用于描述张量级别的计算。
图优化:TVM 提供了图优化工具,可以对计算图进行优化,包括算子融合、内存规划等。
代码生成:TVM 可以将优化后的计算图生成特定硬件平台的机器码。
运行时支持:TVM 提供了轻量级的运行时系统,支持在多种设备上部署和执行编译后的模型。
硬件抽象:TVM 通过 Device API 抽象了硬件操作,使得开发者可以不必关心底层硬件细节。
模块化设计:TVM 的设计允许开发者根据需要选择使用其不同的组件,如编译器、运行时系统、图优化工具等。
3.TVM架构
TVM 的架构包括前端、中端(图优化 Pass 机制)、后端(代码生成)、运行时等部分。它通过定义和演示关键概念,引导用户了解 TVM 的所有主要功能。用户可以通过 TVM 的 Python 接口来编译和优化模型,也可以使用命令行界面进行操作。TVM 还支持交叉编译和远程过程调用(RPC),以及使用 GPU 编译深度学习模型。
在这个网站中 https://tvm.hyper.ai/docs/tutorial/intro/ 介绍了Apache TVM的详细原理和架构以及开源代码,感兴趣的可以过去看看,这里摘要一些重点信息来展示给大家:
TVM 和模型优化概述
下图说明了使用 TVM 优化编译器框架转换时所采取的步骤。
步骤一:从 TensorFlow、PyTorch 或 ONNX 等框架导入模型。在导入阶段中,TVM 可以从其他框架(如 TensorFlow、PyTorch 或 ONNX)中提取模型。 TVM 为前端提供的支持水平会随着我们不断改进这个开源项目而变化。如果在将模型导入 TVM 时遇到问题,可以将其转换为 ONNX。
步骤二:翻译成 TVM 的高级模型语言 Relay。已导入 TVM 的模型在 Relay 中表示。Relay 是神经网络的功能语言和中间表示(IR)。它支持:
1.传统的数据流式表示
2.函数式作用域,let-binding 使其成为一种功能齐全的可微语言
3.允许用户混用两种编程风格的能力
4.Relay 应用图级优化 pass 来优化模型。
步骤三:降级为张量表达式(TE)表示。降级是指将较高级的表示转换为较低级的表示。应用了高级优化之后,Relay 通过运行 FuseOps pass,把模型划分为许多小的子图,并将子图降级为 TE 表示。
张量表达式(TE)是一种用于描述张量计算的领域特定语言。 TE 还提供了几个 schedule 原语来指定底层循环优化,例如循环切分、矢量化、并行化、循环展开和融合。为将 Relay 表示转换为 TE 表示,TVM 包含了一个张量算子清单(TOPI),其中包含常用张量算子的预定义模板(例如,conv2d、transpose)。
步骤四:使用 auto-tuning 模块 AutoTVM 或 AutoScheduler 搜索最佳 schedule。schedule 为 TE 中定义的算子或子图指定底层循环优化。auto-tuning 模块搜索最佳 schedule,并将其与 cost model 和设备上的测量值进行比较。 TVM 中有两个 auto-tuning 模块。
1.AutoTVM:基于模板的 auto-tuning 模块。它运行搜索算法以在用户定义的模板中找到可调 knob 的最佳值。 TOPI 中已经提供了常用算子的模板。
2.AutoScheduler(又名 Ansor):无模板的 auto-tuning 模块。它不需要预定义的 schedule 模板,而是通过分析计算定义自动生成搜索空间,然后在生成的搜索空间中搜索最佳 schedule。
步骤五:为模型编译选择最佳配置。调优后,auto-tuning 模块会生成 JSON 格式的调优记录。此步骤为每个子图选择最佳 schedule。
步骤六:降级为张量中间表示(TIR,TVM 的底层中间表示)。基于调优步骤选择最佳配置后,所有 TE 子图降级为 TIR 并通过底层优化 pass 进行优化。接下来,优化的 TIR 降级为硬件平台的目标编译器。这是生成可部署到生产的优化模型的最终代码生成阶段。 TVM 支持多种不同的编译器后端:
1.LLVM,针对任意微处理器架构,包括标准 x86 和 ARM 处理器、AMDGPU 和 NVPTX 代码生成,以及 LLVM 支持的任何其他平台。
2.特定编译器,例如 NVCC(NVIDIA 的编译器)。
3.嵌入式和特定 target,通过 TVM 的 自定义代码生成(Bring Your Own Codegen, BYOC)框架实现。
步骤七:编译成机器码。compiler-specific 的生成代码最终可降级为机器码。 TVM 可将模型编译为可链接对象模块,然后轻量级 TVM runtime 可以用 C 语言的 API 来动态加载模型,也可以为 Python 和 Rust 等其他语言提供入口点。或将 runtime 和模型放在同一个 package 里时,TVM 可以对其构建捆绑部署。
经过以上的理论科普,我想用我自己的话来总结一下TVM有什么用:
即:TVM可以把你在任何平台上训练好的模型(Windows,Linux,ARM端等等),打包好,并把你优化好的模型放到任何你想要放到的平台上运行,也就是AI技术的落地。
以上只是自己学习过程中的一个笔记,如果大家想做一些关于TVM的例子,了解一下怎么将自己搭建的神经网络模型转换成onnx格式的模型权重,非常推荐大家看这篇大佬的经验贴,https://cloud.tencent.com/developer/article/2346576,写的非常好。