ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

PyTorch TorchScript ONNX Exporter 符号函数体系:Opset 继承、Symbolic 注册与导出原理解析

PyTorch TorchScript ONNX Exporter 符号函数体系:Opset 继承、Symbolic 注册与导出原理解析 PyTorch TorchScript ONNX Exporter 符号函数体系Opset 继承、Symbolic 注册与导出原理解析【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch本文围绕 PyTorch 仓库中 torch/onnx/_internal/torchscript_exporter/README.md 展开系统梳理传统 TorchScript 版 ONNX 导出器的符号函数symbolic functions设计为何以 Opset 9 为基准、如何在各symbolic_opset{version}.py中扩展算子、如何使用registration.onnx_symbolic注册新符号函数、参数命名约束、GraphContext上下文机制以及标量逐点运算pointwise by scalar的转换细节。读完本文你将掌握阅读与扩展这套符号函数代码的完整方法论并能基于仓库源码定位到每个机制的实现证据。背景这套导出器的定位与弃用状态在 PyTorch 2.9 之前TorchScript 是 ONNX 导出的核心中间表示模型先被torch.jit.trace/torch.jit.script捕获为 TorchScript 图再由符号函数把图中每个 ATen 算子逐一翻译为 ONNX 算子。这套机制的代码就托管在 torch/onnx/_internal/torchscript_exporter/ 目录下。需要特别说明的是该目录中的 README.md 开头即给出了明确的弃用声明这套 exporter 是legacy的自 PyTorch 2.9 起已进入弃用状态deprecated并将随 TorchScript 一同移除。因此本文描述的是历史机制与现存源码读者不应在新代码中依赖它仓库中该目录的源码仍然完整可读是理解 ONNX 符号化翻译这一经典设计的最佳教材新代码应转向基于 torch.export / dynamo 的现代化导出路径如 torch/onnx/_internal/exporter 相关实现本文不展开。尽管已被弃用其核心思想——为每个 opset 版本维护一份算子符号函数并以某一基准版本为基础做增量覆盖——至今仍是 ONNX 导出乃至各类图翻译工具的共同范式。Symbolic Functions 与 Opset 版本体系为什么 Opset 9 是基准版本README 明确说明Opset 9 是基础版本base version选择它有两个原因Opset 9 是 PyTorch 导出支持的第一个 opset 版本历史包袱最小Opset 9 比更早的版本更健壮Opset 7/8 存在局限性某些基础算子无法在 ONNX 中表达。与其让整个导出器建立在这些缺陷之上不如以 Opset 9 为基线把这些早期版本的特殊情况单独处理。同时README 也划定了支持边界Opset 7 之前的向后支持不在路线图内。这意味着你不会在仓库中看到针对 Opset 6 及以下的专门符号函数。继承机制未覆盖的算子自动回落到 Opset 9对于除 9 之外的其它 opset 版本默认情况下它们直接继承 symbolic_opset9.py 中定义的符号函数。也就是说你只需要为每个新 opset 编写发生变化的算子没有在symbolic_opset{N}.py中重新定义的算子导出时自动使用 Opset 9 的版本。这个继承链在源码结构上一目了然。仓库中按版本号组织了完整的文件系列symbolic_opset7.py、symbolic_opset8.pysymbolic_opset9.py基准symbolic_opset10.py 至 symbolic_opset20.py每个 opset 文件顶部的_onnx_symbolic别名集中声明了该文件服务的 opset 版本例如 symbolic_opset10.py_onnx_symbolic functools.partial(registration.onnx_symbolic, opset10)文件内所有算子都用这个别名装饰保证注册到的 opset 版本一致。如何在更高版本上扩展算子以topk为例README 给出了两条具体示例路径symbolic_opset10.py中的topk演示如何在 Opset 10 上更新算子symbolic_opset8.py中的upsample_nearest2d演示如何在 Opset 8 上为早期版本提供反向覆盖。看 symbolic_opset10.py 中 topk 的实际实现_onnx_symbolic(aten::topk) symbolic_helper.parse_args(v, v, i, i, i, none) def topk(g: jit_utils.GraphContext, self, k, dim, largest, sorted, outNone): return symbolic_helper._topk_helper( g, self, k, dim, largestlargest, sortedsorted, outout )这段代码展示了扩展一个算子所需的全部要素_onnx_symbolic(aten::topk)把该函数注册为aten::topk在 Opset 10 下的符号函数覆盖 Opset 9 的默认实现symbolic_helper.parse_args(v, v, i, i, i, none)声明参数类型v为 Value/tensori为 intnone为可选用于在分发前做参数规范化与解析函数体通常委托给symbolic_helper中的共享助手_topk_helper避免在多个 opset 文件中重复逻辑。编写 Symbolic 文件注册、命名与 Inplace 处理用registration.onnx_symbolic注册新函数README 指出注册新的符号函数应使用内部装饰器registration.onnx_symbolic并以def reshape(g, self, shape):为参考示例。该装饰器定义在 registration.pydef onnx_symbolic( name: str, opset: OpsetVersion | Sequence[OpsetVersion], decorate: Sequence[Callable] | None None, custom: bool False, ) - Callable: Registers a symbolic function. Usage:: onnx_symbolic( aten::symbolic_b, opset10, decorate[quantized_aten_handler(scale1 / 128, zero_point0)], ) symbolic_helper.parse_args(v, v, b) def symbolic_b(g: _C.Graph, x: _C.Value, y: _C.Value, arg1: bool) - _C.Value: ... 关键行为name必须是domain::op形式的限定名如aten::add否则抛出ValueErroropset可以是单个版本也可以是版本序列——为多版本一次性注册时可直接传入元组decorate是额外装饰器列表如量化的quantized_aten_handler装饰器内部把函数按(name, opset_version)写入全局registrySymbolicRegistry并返回原始函数装饰器只对注册实例生效。配套的custom_onnx_symbolicregistration.py则用于注册覆盖既有符号函数的自定义版本对应customTrue分支供高级扩展场景使用。再回到示例本身——symbolic_opset9.py 中reshape的实现_onnx_symbolic(aten::reshape) symbolic_helper.quantized_args(True) def reshape(g: jit_utils.GraphContext, self, shape): return symbolic_helper._reshape_helper(g, self, shape)注意它只负责翻译真正构造 ONNXReshape节点的逻辑收敛在_reshape_helper中符号函数保持薄封装便于被reshape_assymbolic_opset9.py等其它算子复用_onnx_symbolic(aten::reshape_as) symbolic_helper.quantized_args(True) def reshape_as(g: jit_utils.GraphContext, self, other): shape g.op(Shape, other) return reshape(g, self, shape)参数名必须与native_functions.yaml精确匹配这是编写符号函数时最容易踩坑的约束README 强调符号函数的参数名必须与 aten/src/ATen/native/native_functions.yaml 中的名称完全一致因为算子分发dispatch是按关键字参数进行的。换句话说符号函数签名中的self、other、dim、alpha等名字不是随意取的而是与 ATen 原生算子的 schema 一一对应。如果写错参数名分发时关键字匹配就会失败。从源码结构看symbolic_opset9.py中的add函数symbolic_opset9.py就是这种命名约束的直接体现——其签名def add(g, self, other, alphaNone)与 ATen 的addschema 完全对齐。Inplace 算子透明分发到非 inplace 版本ONNX 图本身不区分 inplace 语义但 PyTorch 图中存在add_、relu_这类就地算子。README 解释了处理方式inplace 算子由_jit_pass_onnx_remove_inplace_ops_for_onnx这个 pass 检测出来之后在run_symbolic_function中被透明地分发到其非 inplace 版本即符号函数只需编写一次非 inplace 版本inplace 变体自动复用。这一设计保证了符号函数集合的简洁性你不需要为每个算子的_变体重复编写翻译逻辑。Tensor 类型依赖的原则宁少勿多README 特别给出一节关于 Tensor 类型的忠告一般情况下应避免依赖 trace graph 中 Tensor Value 的类型信息但有时不可避免例如 ONNX 规范本身要求区分类型TensorType对象提供了这些属性的访问器属性在静态已知时返回值否则返回nullopt应遵循依赖最少、最不具体的信息的优先级完全不依赖 tensor 属性最好优于依赖维度数量次之优于依赖具体 shape最差这样做的根本目的是让符号函数对不同的输入图更健壮一个图里同一算子可能遇到各种张量形状符号函数若过早假设具体 shape导出就会在边界图上失败。因此能通过 ONNX 动态 shape 机制表达的就不要硬编码这也是_reshape_helper、_topk_helper等共享助手大量使用g.op(Shape, ...)这类运行时取 shape操作的原因参见 symbolic_opset9.py 中_shape_as_tensor与_reshape_from_tensor的实现。符号函数的第一个参数GraphContextREADME 说明符号函数的第一个参数永远是GraphContext对象。GraphContext包含torch.Graph对象上定义的所有方法外加符号函数所需的上下文信息。GraphContext定义在 jit_utils.py是一个 dataclass关键字段包括字段含义graph正在构造的_C.Graphblock正在构造的_C.Blockopset当前 opset 版本号original_node正在被转换的原始节点params_dict从图初始值名称到 IValue 的映射envTorch 域图 Value 到 ONNX 域图 Value 的映射values_in_envenv 中所有值的集合用于常数时间查找new_nodes所有新增节点列表用于保证元数据传播到新节点值得注意的是GraphContext通过__getattr__把_C.Graph的方法全部中继出来jit_utils.py因此旧代码中期望_C.Graph的符号函数无需改动即可兼容。文件顶部还声明了一个重要边界该类不面向外部使用接口可能演进请勿在torch.onnx之外依赖它。符号函数之所以需要上下文是因为仅凭算子自身的输入输出不足以完成翻译。README 给出的典型例子是prim::Loop——它需要访问原始节点的子块sub-block来递归翻译循环体。一般符号函数只需要原始节点的输入和属性GraphContext则把这类超出常规的需求一并满足。符号函数构造 ONNX 节点靠g.op(...)方法jit_utils.py。它的属性命名约定是{attr_name}_{type_specifier}例如alpha_f表示 float 类型的alpha属性kernel_shape_i、ceil_mode_i表示 int 属性可对照 symbolic_opset10.py 中 MaxPool 节点的构造pads_i、strides_i等则是 int 列表属性。outputs参数用于声明多输出算子如outputs2的 MaxPool返回值相应变为 Value 元组。Export inplace丢弃注解是 sound 的但会丢失信息README 用单独一节讨论 inplace 导出语义导出 inplace 注解比不导出更好因为它是对 ONNX 导出目标runtime有用的信息可以帮助目标端更高效地执行但ONNX 目前并未正式形式化 inplace 语义因此丢弃 inplace 注解是 sound语义安全的——结果图仍然是正确等价的——但代价是丢失了这一信息。这与上文inplace 算子透明分发到非 inplace 版本的策略是同一个决策的两面既然 ONNX 没有 inplace 概念那么把所有 inplace 算子规范化canonicalize为非 inplace 版本是既安全又统一的处理方式。Pointwise by scalar张量加常量的翻译细节README 最后用较大篇幅剖析了一个看似简单、实则涉及多个机制的案例张量与常量相加如x 2在 ONNX 翻译中会发生什么。机制一标量已不再是标量当标量到达符号函数时它不再是 Python 的 long/float而是一个numel 1的 PyTorch 张量作者注理想情况下应是零维张量但这一改动当时尚未落地。关键问题是这个标量张量的 dtype 与用户 Python 代码中写的字面量完全一致却可能与参与运算的张量 dtype 不匹配。PyTorch 会对标量做隐式类型转换type promotionONNX不会做这种隐式转换因此转换必须由导出器自己完成。承担这一职责的是symbolic_helper._if_scalar_type_as()——在符号函数层把标量转成目标张量的类型_jit_pass_onnx_scalar_type_analysis——在图优化 pass 层面做标量类型分析。两者配合保证x 2在 ONNX 图里得到与 PyTorch eager 语义一致的类型结果。机制二重载分发上的惊人巧合README 指出了一个有趣的事实tensortensor 与 tensorscalar 两种重载的分发形式不同——两个张量相加时分发为add(*[self, other], **{alpha: alpha})张量与标量相加时分发为add(*[self], **{other: other, alpha: alpha})也就是说标量重载把other以关键字参数形式传入而other恰好就是该算子在 ATen schema 中的属性名。利用参数名与标量属性名对齐这个巧合只需要编写一个函数就能同时覆盖两个重载——这也是前文参数名必须与 native_functions.yaml 精确匹配这一约束在实际编码中的价值所在。看 symbolic_opset9.py 中add的签名def add(g, self, other, alphaNone)即可印证other既是 tensor 重载的位置参数也是 scalar 重载的关键字参数一个函数两种用途。机制三rounding mode 之类的语义分支symbolic_opset10.py中的divsymbolic_opset10.py是另一个展示符号函数如何按语义分支的例子当存在额外参数如rounding_mode时_div_rounding_mode会区分floor等模式否则回退到 Opset 9 的true_divide。而_floor_dividesymbolic_opset10.py则展示了浮点与整数除法的不同处理浮点走true_divideFloor整数则用Div/Mod/Xor/Where组合修正向负无穷取整的语义差异——这正是某些基础算子无法在早期 ONNX 中直接表达需要特殊处理的典型体现。小结从这套设计中学到什么尽管这套 TorchScript 版 ONNX exporter 已随 PyTorch 2.9 进入弃用通道但它沉淀的设计原则仍然值得借鉴以基准 opset 增量覆盖组织版本兼容Opset 9 为基各版本文件只写差异未覆盖算子自动回落文件证据symbolic_opset7.py 至 symbolic_opset20.py符号函数保持薄封装复杂逻辑收敛在symbolic_helper共享助手中便于跨算子、跨版本复用如_reshape_helper、_topk_helper参数名与 ATen schema 严格对齐是关键字分发正确性的前提对照 native_functions.yaml上下文对象GraphContext在保留torch.Graph全部方法的同时携带导出所需的额外状态解决了prim::Loop等需要访问子块的算子翻译问题inplace 规范化与标量类型分析由专用 pass 统一完成符号函数无需感知这些琐碎差异依赖最少化原则能不依赖 tensor 属性就不依赖能靠维度就不靠 shape最大化符号函数对多样输入图的健壮性。对想要深入 ONNX 导出原理、或维护依赖旧版导出路径代码的开发者而言torch/onnx/_internal/torchscript_exporter/ 目录及其 README.md 是一份结构清晰、示例完备的活教材。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表