• 文档 >
  • 核心 ATen 运算符集定义
快捷方式

核心 ATen 运算符集定义

本页面提供了核心 ATen 运算符集(opset)的描述和背景。建议正在为 ExecuTorch 开发新的内核库或代理(delegate)的开发者阅读本页面。同时,也建议将 torch.export 作为先决条件;特别是 torch FX 图、运算符分解(operator decomposition)和函数化(functionalization)等概念。

可在 PyTorch 文档网站的 IRs 页面 上找到已确定为核心 ATen 运算符的运算符列表。

什么是运算符集?

torch.export 对给定的 PyTorch 程序执行完整的图捕获,生成一个描述程序执行计算的图 IR。运算符(即在 Tensor 上执行的操作)是图中的基本计算单元,通常对应于图 IR 中的一个唯一节点。运算符的主要来源是 ATen 库;除了 ATen 运算符外,开发者还可以定义自己的运算符(即自定义运算符)。

“ATen 运算符集”或“ATen opset”是在 PyTorch 程序被捕获到图 IR 后,用于表示该程序的 ATen 运算符集合。

函数式 ATen 运算符集

torch.export 的程序捕获机制会生成一个函数式图,该图只允许函数式运算符(即不修改或别名输入的运算符)。因此,torch.export 生成的图将包含函数式 ATen 运算符集,其中只包含函数式 ATen 运算符。

核心 ATen 运算符集

导出的图可以通过应用运算符分解来进一步转换。这个过程将指定的 ATen 运算符替换为等效的其他 ATen 运算符序列。例如,aten.hardsigmoid 可以被替换为 aten.clamp(aten.clamp(self + 3, min=0), max=6) / 6

如果使用默认的分解设置对 PyTorch 程序进行分解,则生成的图 IR 将包含“核心 ATen”运算符集。这个运算符集将是函数式 ATen 运算符集的一个子集,因为一些运算符将被分解。作为核心 ATen 运算符集一部分的 ATen 运算符(即核心 ATen 运算符)在默认分解设置下不会被分解。通常,核心 ATen 运算符不能通过分解轻松地用其他 ATen 运算符重新表示。

核心 ATen 运算符集背后的主要动机是减少模型导出后 PyTorch 后端和编译器需要处理的运算符数量。ATen 库中定义的运算符数量众多,而且可能会添加新的运算符,或者现有运算符的模式(schema)可能会发生变化。如果没有运算符分解,建立在 torch.export 生成的 IR 之上的后端将不得不处理庞大的运算符表面积以及不断变化的运算符集。核心 ATen 运算符集通过定义一个更小、更易于管理且在设计时考虑了稳定性的运算符集来解决这个问题。

核心 ATen 运算符集的开发

尽管 ExecuTorch 使用核心 ATen 运算符集,但它并非 ExecuTorch 所特有。核心 ATen 运算符集的主要设计目标之一是尽可能通用;绝大多数用例不会希望分解其中包含的运算符。因此,核心 ATen 运算符集所隐含的分解应该对绝大多数用例都有益。

另一个关键考虑是保持运算符集尽可能精简,但又不以牺牲性能或开发者体验为代价,强制进行可能产生重大负面影响的分解。

核心 ATen 运算符集是通过审查 ATen 运算符列表而开发的,该列表是通过调查公开 GitHub 存储库中的模型以及知名的开源模型而创建的。调查过程的目的是获得一个精简的 ATen 运算符列表,作为 ATen 运算符使用频率的代理。这样就可以优先审查最常用的运算符。

每个运算符应被视为核心运算符还是通过核心 ATen 分解表进行分解的决定是基于

  1. 检查运算符的潜在分解;分解应该是使用其他 ATen 运算符对运算符进行相对直接的重新表达。

    • 分解不应看起来像是运算符的直接实现。

    • 分解不应因输入的运行时特性而异。

    • 我们还考虑了分解运算符是否会影响输出的精度、数值有效性或内存布局。

  2. 考虑开发者是否会出于性能或其他原因希望在图中保留该运算符。

    • 例如,也许一个运算符可以被分解,但它可以在大多数平台上映射到单个硬件指令,在这种情况下,将其提升为核心运算符会更可取。

未来工作

在所有 ATen 运算符都经过审查并被指定为“核心”或“默认分解”之前,核心 ATen 运算符集不能被认为是完全完成的。然而,这是一项艰巨的任务,而且有很多不常使用的运算符。这就是为什么采用了通过调查模型来确定哪些运算符最常用的方法,从而优先处理“更高影响力”的运算符。

尽管如此,仍有许多运算符尚未评估。计划是根据需要继续评估其他运算符;PyTorch 社区可以通过在 GitHub 上打开问题或 在 PyTorch 论坛的此帖子下发表评论 来提议其他核心运算符或核心分解。

文档

访问全面的 PyTorch 开发者文档

查看文档

教程

为初学者和高级开发者提供深入的教程

查看教程

资源

查找开发资源并让您的问题得到解答

查看资源