ExecuTorch 的高级架构和组件¶
本页介绍 ExecuTorch 的技术架构及其各个组件。本文档面向将 PyTorch 模型部署到边缘设备的工程师。
上下文
为了将设备上 AI 作为目标,同时考虑到各种硬件、关键电源要求和实时处理需求,单一的整体解决方案并不实用。相反,需要一个模块化、分层且可扩展的架构。ExecuTorch 定义了一个简化的工作流程来准备(导出、转换和编译)并执行 PyTorch 程序,并提供精心设计的开箱即用默认组件和明确定义的自定义入口点。这种架构极大地提高了可移植性,使工程师能够使用高性能、轻量级、跨平台的运行时,该运行时可以轻松集成到不同的设备和平台中。
概述¶
将 PyTorch 模型部署到设备上有三个阶段:程序准备、运行时准备和程序执行,如下图所示,包含多个用户入口点。我们将在本文档中分别讨论每个步骤。
图 1. 该图说明了三个阶段 - 程序准备、运行时准备和程序执行。
程序准备¶
ExecuTorch 将 PyTorch 的灵活性和可用性扩展到边缘设备。它利用 PyTorch 2 的编译器和导出功能(TorchDynamo、AOTAutograd、量化、动态形状、控制流 等)来准备在设备上执行 PyTorch 程序。
程序准备通常简称为 AOT(提前),因为导出、转换和编译程序是在使用用 C++ 编写的 ExecuTorch 运行时最终运行之前执行的。为了拥有轻量级运行时并在执行过程中产生较少的开销,我们将尽可能多的工作推迟到 AOT。
从程序源代码开始,以下是您完成程序准备过程的步骤。
程序源代码¶
导出¶
为了将程序部署到设备,工程师需要一个图形表示来编译模型以在各种后端上运行。使用 torch.export()
,将生成一个 EXIR(导出中间表示),其中包含 ATen 方言。所有 AOT 编译都基于此 EXIR,但如下所述,在降低路径中可能包含多种方言。
ATen 方言。PyTorch Edge 基于 PyTorch 的张量库 ATen,该库为高效执行提供了明确的契约。ATen 方言是使用 ATen 节点表示的图形,这些节点完全符合 ATen。允许使用自定义运算符,但必须在调度器中注册。它是扁平化的,没有模块层次结构(较大模块中的子模块),但源代码和模块层次结构保存在元数据中。此表示也适用于自动微分。
可选地,可以将量化(无论是 QAT(量化感知训练)还是 PTQ(训练后量化))应用于整个 ATen 图,然后转换为核心 ATen。量化有助于减小模型大小,这对边缘设备很重要。
核心 ATen 方言。ATen 包含数千个运算符。它不适合一些基本的转换和内核库实现。来自 ATen 方言图的运算符被分解为基本运算符,以便运算符集 (op set) 更小,并且可以应用更多基本的转换。核心 ATen 方言也可以序列化并转换为 Edge 方言,如下所述。
边缘编译¶
上面讨论的导出过程操作的是与最终执行代码的边缘设备无关的图。在边缘编译步骤中,我们处理特定于边缘的表示。
Edge 方言。所有运算符要么与 ATen 运算符兼容,并具有 dtype 加上内存布局信息(表示为
dim_order
),要么是已注册的自定义运算符。标量被转换为张量。这些规范允许后续步骤专注于较小的 Edge 领域。此外,它还支持基于特定 dtype 和内存布局的选择性构建。
使用 Edge 方言,有两种目标感知方法可以进一步将图降低到后端方言。此时,特定硬件的委托可以执行许多操作。例如,iOS 上的 Core ML、Qualcomm 上的 QNN 或 Arm 上的 TOSA 可以重写图。此级别的选项包括:
后端委托。用于将图(完整或部分)编译到特定后端的入口点。在该转换过程中,编译后的图将替换为语义等效的图。为了提高性能,编译后的图将在运行时(也称为
委托
)被卸载到后端。用户定义的流程。用户还可以执行目标特定的转换。此类转换的良好示例包括内核融合、异步行为、内存布局转换等。
编译为 ExecuTorch 程序¶
上面的边缘程序适合编译,但不适合运行时环境。设备部署工程师可以降低图的复杂度,使其能够被运行时高效地加载和执行。
在大多数边缘环境中,动态内存分配/释放会导致显著的性能和功耗开销。可以使用 AOT 内存规划和静态执行图来避免这些问题。
ExecuTorch 运行时是静态的(指图表示,但仍支持控制流和动态形状)。为了避免创建输出和返回值,所有功能性操作符表示将转换为输出变体(输出作为参数传递)。
用户可以选择应用自己的内存规划算法。例如,嵌入式系统可能存在特定的内存层次结构。用户可以针对该内存层次结构定制自己的内存规划。
程序将以 ExecuTorch 运行时能够识别的格式发出。
最后,发出的程序可以序列化为 flatbuffer 格式。
运行时准备¶
有了序列化的程序和提供的内核库(用于操作符调用)或后端库(用于委托调用),模型部署工程师现在可以为运行时准备程序。
ExecuTorch 具有 选择性构建 API,用于构建仅链接到程序使用的内核的运行时,这可以显著减少生成的应用程序的二进制文件大小。
程序执行¶
ExecuTorch 运行时是用 C++ 编写的,依赖项很少,以确保可移植性和执行效率。由于程序经过了良好的 AOT 准备,核心运行时组件很少,包括
平台抽象层
日志记录和可选的分析
执行数据类型
内核和后端注册表
内存管理
执行器是加载程序并执行程序的入口点。执行会从这个非常小的运行时触发相应的操作符内核或后端执行。
开发人员工具¶
对于用户而言,使用上面的流程从研究到生产应该很高效。生产力至关重要,以便用户能够编写、优化和部署他们的模型。我们提供 ExecuTorch 开发人员工具 来提高生产力。开发人员工具不在图中。相反,它是一个工具集,涵盖了所有三个阶段的开发人员工作流程。
在程序准备和执行期间,用户可以使用 ExecuTorch 开发人员工具来分析、调试或可视化程序。由于端到端流程都在 PyTorch 生态系统中,用户可以将性能数据与图可视化以及对程序源代码和模型层次结构的直接引用相关联并显示出来。我们认为这是快速迭代并将 PyTorch 程序降低到边缘设备和环境的关键组成部分。