快捷方式

架构与组件

本页介绍 ExecuTorch 的技术架构及其各个组件。本文档面向将 PyTorch 模型部署到边缘设备的工程师。

背景

为了针对具有多样化硬件、严苛功耗要求和实时处理需求的设备端 AI,单一的整体解决方案并不实用。相反,需要一种模块化、分层和可扩展的架构。ExecuTorch 定义了一个简化的工作流程来准备(导出、转换和编译)和执行 PyTorch 程序,提供了精心设计的开箱即用默认组件和定义明确的自定义入口点。这种架构极大地提高了可移植性,使工程师能够使用高性能、轻量级、跨平台的运行时,轻松集成到不同的设备和平台。

概览

将 PyTorch 模型部署到设备端分为三个阶段:程序准备、运行时准备和程序执行,如下图所示,并带有多个用户入口点。我们将在本文档中分别讨论每个步骤。

图 1. 该图展示了三个阶段:程序准备、运行时准备和程序执行。

程序准备

ExecuTorch 将 PyTorch 的灵活性和可用性扩展到边缘设备。它利用 PyTorch 2 编译器和导出功能(TorchDynamoAOTAutograd量化动态形状控制流等)来准备 PyTorch 程序以便在设备上执行。

程序准备通常简称为 AOT(提前,ahead-of-time),因为程序的导出、转换和编译是在使用 ExecuTorch 运行时(用 C++ 编写)最终运行之前执行的。为了拥有轻量级运行时和较小的执行开销,我们将尽可能多的工作推到 AOT 阶段。

从程序源代码开始,以下是您完成程序准备需要经历的步骤。

程序源代码

  • 像所有 PyTorch 用例一样,ExecuTorch 从模型编写开始,其中创建了标准的 nn.Module Eager Mode PyTorch 程序。

  • 导出特定的辅助工具用于表示高级特性,例如控制流(例如,用于跟踪 if-else 两个分支的辅助函数)和动态形状(例如,数据依赖的动态形状约束)。

导出

为了将程序部署到设备,工程师需要一个图形表示来编译模型以便在各种后端上运行。使用 torch.export(),将生成带有 ATen dialect 的 EXIR(导出中间表示)。所有 AOT 编译都基于此 EXIR,但如以下详细介绍,在 Lowering 路径上可以有多个 dialect。

  • ATen Dialect。PyTorch Edge 基于 PyTorch 的 Tensor 库 ATen,该库为高效执行提供了清晰的契约。ATen Dialect 是由完全符合 ATen 规范的 ATen 节点表示的图形。允许使用自定义运算符,但必须使用 Dispatcher 进行注册。它是扁平化的,没有模块层次结构(在一个更大的模块中的子模块),但源代码和模块层次结构保留在元数据中。这种表示也是 Autograd 安全的。

  • 可以选择性地,可以在将整个 ATen 图转换为 Core ATen 之前应用量化,无论是 QAT(量化感知训练)还是 PTQ(训练后量化)。量化有助于减小模型大小,这对于边缘设备非常重要。

  • Core ATen Dialect。ATen 有数千个运算符。对于一些基本转换和内核库实现来说,它并不理想。ATen Dialect 图中的运算符被分解为基本运算符,以便运算符集 (op set) 更小并且可以应用更基本的转换。Core ATen dialect 也是可序列化的,并且可以转换为 Edge Dialect,如下所述。

边缘编译

上面讨论的导出过程在一个与最终执行代码的边缘设备无关的图形上进行操作。在边缘编译步骤中,我们处理特定于边缘的表示形式。

  • Edge Dialect。所有运算符要么符合带有 dtype 和内存布局信息(表示为 dim_order)的 ATen 运算符,要么是注册的自定义运算符。标量被转换为 Tensor。这些规范允许后续步骤专注于一个更小的 Edge 领域。此外,它还支持基于特定 dtype 和内存布局的选择性构建。

使用 Edge dialect,有两种目标感知的方式将图进一步 Lowering 到 Backend Dialect。此时,特定硬件的委托可以执行许多操作。例如,iOS 上的 Core ML、Qualcomm 上的 QNN 或 Arm 上的 TOSA 可以重写图。此级别的选项包括

  • 后端委托 (Backend Delegate)。这是将图(完整或部分)编译到特定后端的入口点。在此转换期间,编译后的图会替换为语义等效的图。编译后的图稍后将在运行时卸载到后端(也称为委托)以提高性能。

  • 用户自定义 Passes。用户也可以执行目标特定的转换。这方面很好的例子包括内核融合、异步行为、内存布局转换等。

编译为 ExecuTorch 程序

上述 Edge 程序适合编译,但不适合运行时环境。设备端部署工程师可以降低图,使其能够被运行时高效加载和执行。

在大多数 Edge 环境中,动态内存分配/释放具有显著的性能和功耗开销。可以使用 AOT 内存规划和静态执行图来避免这种情况。

  • ExecuTorch 运行时是静态的(在图表示的意义上,但仍然支持控制流和动态形状)。为了避免输出创建和返回,所有函数式运算符表示都转换为 out 变体(输出作为参数传递)。

  • 用户可以选择应用自己的内存规划算法。例如,嵌入式系统可能有多层内存层次结构。用户可以自定义内存规划以适应该内存层次结构。

  • 程序会被输出为 ExecuTorch 运行时能够识别的格式。

最后,输出的程序可以序列化为 flatbuffer 格式。

运行时准备

有了序列化的程序,并提供了内核库(用于运算符调用)或后端库(用于委托调用),模型部署工程师现在可以为运行时准备程序。

ExecuTorch 具有 选择性构建 API,用于构建仅链接程序使用的内核的运行时,这可以显着减小最终应用程序的二进制文件大小。

程序执行

ExecuTorch 运行时使用 C++ 编写,依赖项极少,以实现可移植性和执行效率。由于程序已在 AOT 阶段充分准备,核心运行时组件非常精简,包括:

  • 平台抽象层

  • 日志记录和可选的性能分析

  • 执行数据类型

  • 内核与后端注册表

  • 内存管理

Executor 是加载和执行程序的入口点。执行会从这个非常精简的运行时触发相应的运算符内核或后端执行。

开发者工具

用户应该能够高效地使用上述流程从研究转向生产。生产力至关重要,以便用户编写、优化和部署模型。我们提供 ExecuTorch 开发者工具来提高生产力。开发者工具未在图中显示。相反,它是一个覆盖三个阶段开发者工作流程的工具集。

在程序准备和执行期间,用户可以使用 ExecuTorch 开发者工具来分析、调试或可视化程序。由于端到端流程都在 PyTorch 生态系统内,用户可以将性能数据与图可视化以及对程序源代码和模型层次结构的直接引用进行关联和显示。我们认为这是快速迭代并将 PyTorch 程序 Lowering 到边缘设备和环境的关键组件。

文档

访问全面的 PyTorch 开发者文档

查看文档

教程

获取面向初学者和高级开发者的深度教程

查看教程

资源

查找开发资源并获得问题解答

查看资源