概述
Apache TVM 是一个机器学习编译框架,遵循Python 优先开发和通用部署的原则。 它接受经过预训练的机器学习模型,编译并生成可部署的模块,这些模块可以嵌入到任何地方运行。Apache TVM 还支持自定义优化过程,以引入新的优化方法、库、代码生成等等。
核心原则
- Python 优先: 优化过程完全可以在 Python 中自定义。无需重新编译 TVM 栈即可轻松定制优化流程。
- 可组合性: 优化过程具有可组合性。可以轻松地将新的优化过程、库以及代码生成器组合进已有的流程中。
核心目标
- 优化机器学习任务的性能,组合使用库与代码生成器。
- 部署机器学习任务至多种全新环境,包括新的运行时和新硬件。
- 持续改进与定制机器学习部署流程,可通过 Python 快速定制库调度,引入定制算子与代码生成方式。
关键流程
以下是使用 TVM 部署机器学习模型的典型流程。可运行示例请参见 快速开始。
- 导入 / 构建一个机器学习模型
TVM 支持从多种框架导入模型用于通用机器学习模型,如 PyTorch、TensorFlow。同时对于大语言模型场景,也可以使用 Relax 前端直接创建模型。
- 通过
pipelines
执行可组合的优化转换
pipeline 封装了一系列转换,以实现两个目标:
- 图优化: 如算子融合、布局重写等。
- 张量程序优化: 将算子映射到底层实现(包括库或代码生成器)。
备注
这两个是优化目标,而非流水线的两个阶段。这两种优化可以在同一层级执行,也可以在两个阶段中分别进行。
- 构建并通用部署
Apache TVM 旨在提供一种通用部署方案,将机器学习带到任何地方,以最少的运行时支持适配各种语言。TVM 的运行时可在非 Python 环境中运行,因此适用于移动端、边缘设备,甚至是裸机设备。此外,TVM 运行时自带原生数据结构,并可通过 DLPack 支持,与现有生态(如 PyTorch、TensorFlow、TensorRT 等)实现零拷贝交换。