概述

Apache TVM是一个机器学习编译框架，遵循Python优先开发和通用部署原则。它接收预训练的机器学习模型，编译并生成可部署模块，这些模块可以嵌入到任何地方运行。Apache TVM还支持自定义优化流程，以引入新的优化方法、库、代码生成等功能。

核心原则

  • Python优先: 优化过程完全可通过Python自定义。无需重新编译TVM堆栈即可轻松定制优化流程。

  • 可组合性: 优化过程是可组合的。可以轻松地将新的优化通道、库和代码生成器组合到现有流程中。

关键目标

  • 优化机器学习工作负载的性能，组合库和代码生成。

  • 部署 机器学习工作负载到多样化的新环境，包括新运行时和新硬件。

  • 持续改进和定制化 Python中的机器学习部署流水线，通过快速定制库调度、引入自定义运算符和代码生成。

关键流程

以下是使用TVM部署机器学习模型的典型流程。如需可运行的示例，请参考Quick Start

  1. 导入/构建ML模型

    TVM支持从多种框架导入模型，例如PyTorch、TensorFlow等通用ML模型。同时，针对大语言模型场景，我们可以直接使用Relax前端创建模型。

  2. 执行可组合优化 通过 pipelines 进行转换

    该流程封装了一系列转换以实现两个目标：

    • 图优化: 如算子融合和布局重写。

    • 张量程序优化: 将算子映射到底层实现（包括库或代码生成）

    注意

    这两者是目标而非流程的阶段。这两种优化是在同一层级执行的，或者分别在两个阶段单独进行。

  3. 构建与通用部署

    Apache TVM旨在提供通用部署解决方案，通过最小化运行时支持，将机器学习带到任何语言环境。TVM运行时可在非Python环境中运行，因此适用于移动设备、边缘设备甚至裸机设备。此外，TVM运行时自带原生数据结构，并可通过DLPack支持与现有生态系统（PyTorch、TensorFlow、TensorRT等）实现零拷贝交互。