范詳解：ExportedProgram、Graph、Node 與元數(shù)據(jù)契約)
PyTorch torch.export IR 規(guī)范詳解ExportedProgram、Graph、Node 與元數(shù)據(jù)契約【免費下載鏈接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration項目地址: https://gitcode.com/GitHub_Trending/py/pytorchtorch.export是 PyTorch 用于將模型以 AOTAhead-of-Time方式捕獲為編譯期中間表示IR的核心入口。本文以倉庫內(nèi) docs/source/user_guide/torch_compiler/export/ir_spec.md 為骨架系統(tǒng)講解 Export IR 的ExportedProgram、Graph、Node、SymInt、FakeTensor與 Pytree-able 類型體系并結(jié)合torch/export/exported_program.py、torch/export/graph_signature.py、torch/export/__init__.py等源碼幫助你從會調(diào)用torch.export.export()進階到能讀懂、能校驗、能為其編寫后端的層次。Export IR 是什么Export IR 是一種面向編譯器的圖中間表示設(shè)計目標與 MLIR、TorchScript 類似但專門用于表達 PyTorch 程序的語義。它把計算表達為一條精簡的操作列表僅對控制流等動態(tài)行為提供有限支持例如通過cond、map等高階算子表達條件與循環(huán)。它的產(chǎn)生流程是前端通過 trace-specializing跟蹤特化機制可靠地soundly捕獲一個 PyTorch 程序得到的 Export IR 圖隨后可被后端優(yōu)化與執(zhí)行。當前倉庫中這一過程由 torch.export.export() 完成。其 docstring 明確給出了三條關(guān)鍵承諾生成的功能化 ATen 算子以及用戶自定義算子是歸一化的、Python 控制流與數(shù)據(jù)結(jié)構(gòu)在特定例外下被消除、并記錄了一組形狀約束以保證這種歸一化與控制流消除對未來輸入是可靠的soundness guarantee。與 torch.fx 的關(guān)系更強的嚴格性Export IR 是實現(xiàn)在torch.fx.Graph之上的——所有 Export IR 圖同時也是合法的 FX 圖若按標準 FX 語義解釋Export IR 可被可靠地解釋。一個直接推論是導出的圖可以經(jīng)由標準 FX codegen 轉(zhuǎn)換回合法的 Python 程序。因此文檔刻意把篇幅集中在Export IR 與 FX 的差異點更嚴格的約束上而不是重復 FX 已有內(nèi)容。理解這一前提是讀懂本文后續(xù)所有差異條款的基礎(chǔ)。ExportedProgramIR 的頂層容器Export IR 的頂層構(gòu)造是torch.export.ExportedProgram類它把模型的計算圖與模型消費的參數(shù)/權(quán)重打包在一起。倉庫中類定義見 torch/export/exported_program.py#L1070-L1153其 docstring 概括為包含torch.fx.Graph表示的 Tensor 計算、包含所有被提升參數(shù)與 buffer 數(shù)值的 state_dict以及各類元數(shù)據(jù)并且可以像原始可調(diào)用對象一樣被調(diào)用。關(guān)鍵屬性屬性類型說明graph_moduletorch.fx.GraphModule存放拍平后的計算圖圖可直接通過ExportedProgram.graph訪問graph_signaturetorch.export.ExportGraphSignature記錄圖中使用與改寫的參數(shù)、buffer 名稱。參數(shù)與 buffer 不再作為圖模塊的屬性存在而是被提升為圖的輸入signature 負責跟蹤這些參數(shù)的附加信息state_dictDict[str, Union[torch.Tensor, torch.nn.Parameter]]存放參數(shù)與 buffer 數(shù)值的數(shù)據(jù)結(jié)構(gòu)range_constraintsDict[sympy.Symbol, RangeConstraint]對帶有數(shù)據(jù)依賴行為導出的程序節(jié)點元數(shù)據(jù)含符號形狀形如s0、i0該屬性把符號映射到其上下界從源碼看ExportedProgram的構(gòu)造函數(shù)在 exported_program.py#L1112-L1153 中依次完成剝離 graph 上的 codegen 并創(chuàng)建扁平圖模塊、執(zhí)行_common_getitem_elimination_pass公共 getitem 消除、保存 signature/state_dict/range_constraints/module_call_graph、默認裝配Verifier校驗器并以self.validate()收尾——校驗始終是構(gòu)造函數(shù)的最后一步。同時graph_module、graph、graph_signature、state_dict等屬性均為只讀setter 直接拋RuntimeError保證導出的程序不會被隨意改寫。graph_signature參數(shù)與 buffer 的戶口本ExportGraphSignature定義于 torch/export/graph_signature.py#L165-L275。它建模了 Export Graph 的輸入/輸出簽名核心不變量是Export Graph 是功能化的圖內(nèi)不通過getattr節(jié)點訪問參數(shù)、buffer 等狀態(tài)export保證參數(shù)、buffer、常量張量被提升為圖的輸入對 buffer 的改寫也不直接出現(xiàn)在圖內(nèi)而是把改寫后的 buffer 值建模為圖額外的輸出。輸入與輸出的排序規(guī)則為Inputs [*parameters_buffers_constant_tensors, *flattened_user_inputs] Outputs [*mutated_inputs, *flattened_user_outputs]InputKindPARAMETER、BUFFER、CONSTANT_TENSOR、USER_INPUT、CUSTOM_OBJ、TOKEN等與OutputKindUSER_OUTPUT、BUFFER_MUTATION、PARAMETER_MUTATION、USER_INPUT_MUTATION、GRADIENT_TO_PARAMETER等分別枚舉輸入、輸出的類別parameters、buffers、non_persistent_buffers、lifted_tensor_constants、user_inputs、user_outputs等屬性則從input_specs/output_specs中按類別抽取對應(yīng)名稱集合見 graph_signature.py#L277-L374。例如一個同時含參數(shù)、兩個 buffer 并在 forward 中原地改寫 buffer 的模塊其 graph_signature 會呈現(xiàn)為# inputs p_my_parameter: PARAMETER targetmy_parameter b_my_buffer1: BUFFER targetmy_buffer1 persistentTrue b_my_buffer2: BUFFER targetmy_buffer2 persistentTrue x1: USER_INPUT x2: USER_INPUT # outputs add_2: BUFFER_MUTATION targetmy_buffer2 add_1: USER_OUTPUT該示例完整代碼與圖對照見 graph_signature.py#L183-L271其中run_decompositions()會把非功能化圖含aten.add_.Tensor原地算子轉(zhuǎn)換為功能化圖buffer 改寫變?yōu)轭~外輸出add_2。這正是理解Export IR 為什么能放心交給后端執(zhí)行的關(guān)鍵圖內(nèi)不再有隱藏狀態(tài)。GraphDAG 與節(jié)點構(gòu)成一個 Export IR Graph 是 DAG有向無環(huán)圖形式的 PyTorch 程序每個節(jié)點代表一次計算/操作邊由節(jié)點間的引用構(gòu)成。其 schema 極簡class Graph: nodes: List[Node]實踐中 Export IR 的圖即torch.fx.GraphPython 類。一個合法的 Export IR 圖必須包含0 個或多個 op 類型為placeholder的節(jié)點0 個或多個 op 類型為call_function的節(jié)點恰好 1 個op 類型為output的節(jié)點。推論最小的合法 Graph 只有一個節(jié)點即nodes永不為空。定義placeholder節(jié)點集合代表 Graph/GraphModule 的輸入output節(jié)點代表其輸出。一個最簡可運行示例對應(yīng)文檔原始示例import torch from torch import nn class MyModule(nn.Module): def forward(self, x, y): return x y example_args (torch.randn(1), torch.randn(1)) mod torch.export.export(MyModule(), example_args) print(mod.graph)輸出Graph 的文本表示每行一個節(jié)點graph(): %x : [num_users1] placeholder[targetx] %y : [num_users1] placeholder[targety] %add : [num_users1] call_functiontargettorch.ops.aten.add.Tensor, kwargs {}) return (add,)Node節(jié)點 Schema 與 FX 文本格式Node 表示一次具體計算/操作Python 層面用torch.fx.Node類表示。節(jié)點間的邊通過Node.args屬性中對其他節(jié)點的直接引用實現(xiàn)借助 FX 機制可表達計算圖所需的操作符調(diào)用、占位符輸入、條件分支與循環(huán)等。Node 的 schemaclass Node: name: str # name of node op_name: str # type of operation # interpretation of the fields below depends on op_name target: [str|Callable] args: List[object] kwargs: Dict[str, object] meta: Dict[str, object]FX 文本格式如上例所示文本格式的每行遵循%name:[...] op_nametargettarget), kwargs {keyword: arg5})該格式以緊湊形式囊括了 Node 類的全部字段meta除外。具體地name即node.nameop_name即node.op字段必須是call_function、placeholder、get_attr、output之一target即node.target其含義取決于op_nameargs即node.args元組中的內(nèi)容若列表中某值是torch.fx.Node會特別以%前綴標識。例如一個 add 算子的調(diào)用顯示為%add1 call_functiontarget torch.op.aten.add.Tensor, kwargs {})其中%x、%y是兩個名字為 x、y 的 Node。值得注意字符串torch.op.aten.add.Tensor表示的是target 字段中實際存儲的可調(diào)用對象而非僅僅是它的字符串名。文本格式的最后一行return [add]對應(yīng)op_name output的節(jié)點表示返回這一元素。call_function算子調(diào)用節(jié)點call_function節(jié)點表示對某個算子的調(diào)用。兩個關(guān)鍵定義Functional功能化一個可調(diào)用對象滿足以下全部條件即為 functional——非改寫不修改輸入的值對張量而言包括元數(shù)據(jù)與數(shù)據(jù)且無副作用不改寫外部可見的狀態(tài)如模塊參數(shù)值Operator算子具有預定義 schema 的 functional 可調(diào)用對象典型如功能化 ATen 算子。FX 表示%name call_functiontarget operator, kwargs {})與 vanilla FXcall_function的差異FX 圖中call_function可指向任意可調(diào)用對象Export IR 則限定為 ATen 算子的選定子集、自定義算子與控制流算子Export IR 中常量參數(shù)會被內(nèi)嵌進圖內(nèi)FX 圖中g(shù)et_attr節(jié)點可讀取圖模塊中存儲的任何屬性Export IR 中它被限定為只能讀取子模塊因為所有參數(shù)/buffer 都以輸入形式傳入圖模塊。call_function 的元數(shù)據(jù)契約Node.meta是掛在每個 FX 節(jié)點上的 dict但 FX 規(guī)范并不規(guī)定其中會有哪些元數(shù)據(jù)。Export IR 提供了更強的契約所有call_function節(jié)點保證恰好包含以下字段node.meta[stack_trace]字符串引用原始 Python 源碼的調(diào)用棧例如File my_module.py, line 19, in forward return x dummy_helper(y) File helper_utility.py, line 89, in dummy_helper return y 1node.meta[val]描述該操作運行后的輸出類型可為symint、FakeTensor、List[Union[FakeTensor, SymInt]]或Nonenode.meta[nn_module_stack]該節(jié)點來源的torch.nn.Module的調(diào)用棧。例如addmm算子出自nn.Sequential內(nèi)的nn.Linear則表現(xiàn)為{self_linear: (self.linear, class torch.nn.Linear), self_sequential: (self.sequential, class torch.nn.Sequential)}node.meta[source_fn_stack]分解decomposition之前該節(jié)點來自的 torch 函數(shù)或葉子torch.nn.Module類。例如addmm來自nn.Linear模塊調(diào)用則source_fn為torch.nn.Linear來自torch.nn.functional.linear函數(shù)調(diào)用則包含torch.nn.functional.Linear。這些元數(shù)據(jù)使后端能夠在編譯期獲得源碼定位、輸出類型形狀、模塊層級與原始算子來源——這正是 Export IR 能支撐高質(zhì)量錯誤報告與算子級優(yōu)化的基礎(chǔ)。placeholder圖輸入節(jié)點placeholder表示圖的輸入語義與 FX 完全一致。占位節(jié)點必須是圖 nodes 列表中最靠前的 N 個節(jié)點N 可以為 0。FX 表示%name placeholdertarget name)target字段是輸入名稱的字符串a(chǎn)rgs若非空則大小為 1表示該輸入的默認值。元數(shù)據(jù)placeholder 節(jié)點同樣有meta[val]與call_function一致此時val表示圖期望為該輸入?yún)?shù)接收的輸入 shape/dtype。output圖出口節(jié)點output調(diào)用對應(yīng)函數(shù)中的 return 語句因此它終止當前圖。圖中有且僅有一個output 節(jié)點且它永遠是圖的最后一個節(jié)點。FX 表示output[](args (%something, …))語義與torch.fx完全一致args表示要返回的節(jié)點。output 節(jié)點的元數(shù)據(jù)與call_function節(jié)點相同。get_attr子模塊讀取節(jié)點get_attr節(jié)點表示從封裝它的torch.fx.GraphModule中讀取一個子模塊。與torch.fx.symbolic_trace得到的 vanilla FX 圖不同——后者用get_attr從頂層 GraphModule 讀取參數(shù)、buffer 等屬性——Export IR 中參數(shù)與 buffer 均作為輸入傳入圖模塊并存儲在頂層ExportedProgram中因此get_attr在這里僅用于讀取子模塊典型場景是控制流分支的嵌套子圖。FX 表示%name get_attrtarget name)示例使用functorch.experimental.control_flow.cond構(gòu)造條件分支from functorch.experimental.control_flow import cond def true_fn(x): return x.sin() def false_fn(x): return x.cos() def f(x, y): return cond(y, true_fn, false_fn, [x])對應(yīng)圖graph(): %x_1 : [num_users1] placeholder[targetx_1] %y_1 : [num_users1] placeholder[targety_1] %true_graph_0 : [num_users1] get_attr[targettrue_graph_0] %false_graph_0 : [num_users1] get_attr[targetfalse_graph_0] %conditional : [num_users1] call_functiontargettorch.ops.higher_order.cond, kwargs {}) return conditional其中%true_graph_0 get_attr[targettrue_graph_0]讀取包含sin算子的子模塊true_graph_0而torch.ops.higher_order.cond即為文檔所提及的控制流算子之一。引用類型SymInt、FakeTensor 與 Pytree-able 類型SymIntSymInt要么是字面整數(shù)要么是表示某個整數(shù)的符號Python 中由sympy.Symbol類表示。當 SymInt 是符號時它描述的是一個編譯期未知、僅運行時可知的整型變量——這正是動態(tài) shape 在 Export IR 中的載體。FakeTensorFakeTensor是包含張量元數(shù)據(jù)的對象可視為class FakeTensor: size: List[SymInt] dtype: torch.dtype device: torch.device dim_order: List[int] # 尚不存在size是整數(shù)或 SymInt 的列表若含 SymInt 表示該張量具有動態(tài)形狀若全是整數(shù)則假定張量具有該精確靜態(tài)形狀。張量元數(shù)據(jù)的秩rank永不是動態(tài)的dtype表示該節(jié)點輸出的數(shù)據(jù)類型Export IR 中不存在隱式類型提升no implicit type promotionsFakeTensor 中沒有 strides。換言之node.meta[val]與node.target算子的返回類型一一對應(yīng)算子返回 Tensor →val為描述該張量的 FakeTensor算子返回 n 元組張量 →val為 n 元 FakeTensor算子返回編譯期已知的 int/float/scalar →val為None算子返回編譯期未知的 int/float/scalar →val為 SymInt。例如aten::add返回 Tensor其 spec 即帶 dtype 與 size 的 FakeTensoraten::sym_size返回整數(shù)其 val 是 SymInt值僅運行時可得max_pool2d_with_indexes返回(Tensor, Tensor)元組spec 就是 2 元 FakeTensor分別描述返回值的兩個元素。Python 代碼示例def add_one(x): return torch.ops.aten(x, 1)對應(yīng)圖graph(): %ph_0 : [#users1] placeholder[targetph_0] %add_tensor : [#users1] call_functiontargettorch.ops.aten.add.Tensor, kwargs {}) return [add_tensor]FakeTensorFakeTensor(dtypetorch.int, size[2,], deviceCPU)Pytree-able 類型一個類型是Pytree-able的當且僅當它是葉子類型或包含其他 Pytree-able 類型的容器類型pytree 概念與 JAX 的 pytrees 文檔一致。葉子類型leaf type類型定義Tensortorch.TensorScalarPython 中任意數(shù)值類型包括整型、浮點型與零維張量intPython intC 中綁定為 int64_tfloatPython floatC 中綁定為 doubleboolPython boolstrPython 字符串ScalarTypetorch.dtypeLayouttorch.layoutMemoryFormattorch.memory_formatDevicetorch.device容器類型container type類型定義TuplePython tupleListPython listDict以 Scalar 為鍵的 Python dictNamedTuplePython namedtupleDataclass必須通過register_dataclass注冊實現(xiàn)見 torch/export/init.pyCustom class通過_register_pytree_node定義的任意自定義類實現(xiàn)見 torch/utils/_pytree.py該類型體系決定了 export 邊界上可合法傳遞的數(shù)據(jù)結(jié)構(gòu)用戶在torch.export.export()中傳入的args/kwargs、模型返回的結(jié)果最終都會被按 pytree 規(guī)則拍平為葉子節(jié)點再映射為圖的 placeholder 與 output。總結(jié)Export IR 的核心設(shè)計可以濃縮為三個關(guān)鍵詞功能化——參數(shù)、buffer 被提升為輸入buffer 改寫被建模為輸出圖中不存在隱藏狀態(tài)見 graph_signature.py#L166-L176嚴格元數(shù)據(jù)契約——stack_trace、val、nn_module_stack、source_fn_stack四字段為每個call_function節(jié)點背書編譯期即可獲得完整語義受限的算子集合與顯式類型——call_function僅限 ATen 算子子集、自定義算子與控制流算子常量內(nèi)嵌類型經(jīng) SymInt/FakeTensor 顯式刻畫無隱式提升。無論你是為 Export IR 編寫自定義后端還是調(diào)試torch.export.export()的輸出圖ir_spec.md 與其對應(yīng)源碼exported_program.py、graph_signature.py、torch/export/init.py都是最權(quán)威的第一手資料?！久赓M下載鏈接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration項目地址: https://gitcode.com/GitHub_Trending/py/pytorch創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考 PyTorch torch.export IR 規(guī)范詳解：ExportedProgram、Graph、Node 與元數(shù)據(jù)契約
PyTorch torch.export IR 規(guī)范詳解ExportedProgram、Graph、Node 與元數(shù)據(jù)契約【免費下載鏈接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration項目地址: https://gitcode.com/GitHub_Trending/py/pytorchtorch.export是 PyTorch 用于將模型以 AOTAhead-of-Time方式捕獲為編譯期中間表示IR的核心入口。本文以倉庫內(nèi) docs/source/user_guide/torch_compiler/export/ir_spec.md 為骨架系統(tǒng)講解 Export IR 的ExportedProgram、Graph、Node、SymInt、FakeTensor與 Pytree-able 類型體系并結(jié)合torch/export/exported_program.py、torch/export/graph_signature.py、torch/export/__init__.py等源碼幫助你從會調(diào)用torch.export.export()進階到能讀懂、能校驗、能為其編寫后端的層次。Export IR 是什么Export IR 是一種面向編譯器的圖中間表示設(shè)計目標與 MLIR、TorchScript 類似但專門用于表達 PyTorch 程序的語義。它把計算表達為一條精簡的操作列表僅對控制流等動態(tài)行為提供有限支持例如通過cond、map等高階算子表達條件與循環(huán)。它的產(chǎn)生流程是前端通過 trace-specializing跟蹤特化機制可靠地soundly捕獲一個 PyTorch 程序得到的 Export IR 圖隨后可被后端優(yōu)化與執(zhí)行。當前倉庫中這一過程由 torch.export.export() 完成。其 docstring 明確給出了三條關(guān)鍵承諾生成的功能化 ATen 算子以及用戶自定義算子是歸一化的、Python 控制流與數(shù)據(jù)結(jié)構(gòu)在特定例外下被消除、并記錄了一組形狀約束以保證這種歸一化與控制流消除對未來輸入是可靠的soundness guarantee。與 torch.fx 的關(guān)系更強的嚴格性Export IR 是實現(xiàn)在torch.fx.Graph之上的——所有 Export IR 圖同時也是合法的 FX 圖若按標準 FX 語義解釋Export IR 可被可靠地解釋。一個直接推論是導出的圖可以經(jīng)由標準 FX codegen 轉(zhuǎn)換回合法的 Python 程序。因此文檔刻意把篇幅集中在Export IR 與 FX 的差異點更嚴格的約束上而不是重復 FX 已有內(nèi)容。理解這一前提是讀懂本文后續(xù)所有差異條款的基礎(chǔ)。ExportedProgramIR 的頂層容器Export IR 的頂層構(gòu)造是torch.export.ExportedProgram類它把模型的計算圖與模型消費的參數(shù)/權(quán)重打包在一起。倉庫中類定義見 torch/export/exported_program.py#L1070-L1153其 docstring 概括為包含torch.fx.Graph表示的 Tensor 計算、包含所有被提升參數(shù)與 buffer 數(shù)值的 state_dict以及各類元數(shù)據(jù)并且可以像原始可調(diào)用對象一樣被調(diào)用。關(guān)鍵屬性屬性類型說明graph_moduletorch.fx.GraphModule存放拍平后的計算圖圖可直接通過ExportedProgram.graph訪問graph_signaturetorch.export.ExportGraphSignature記錄圖中使用與改寫的參數(shù)、buffer 名稱。參數(shù)與 buffer 不再作為圖模塊的屬性存在而是被提升為圖的輸入signature 負責跟蹤這些參數(shù)的附加信息state_dictDict[str, Union[torch.Tensor, torch.nn.Parameter]]存放參數(shù)與 buffer 數(shù)值的數(shù)據(jù)結(jié)構(gòu)range_constraintsDict[sympy.Symbol, RangeConstraint]對帶有數(shù)據(jù)依賴行為導出的程序節(jié)點元數(shù)據(jù)含符號形狀形如s0、i0該屬性把符號映射到其上下界從源碼看ExportedProgram的構(gòu)造函數(shù)在 exported_program.py#L1112-L1153 中依次完成剝離 graph 上的 codegen 并創(chuàng)建扁平圖模塊、執(zhí)行_common_getitem_elimination_pass公共 getitem 消除、保存 signature/state_dict/range_constraints/module_call_graph、默認裝配Verifier校驗器并以self.validate()收尾——校驗始終是構(gòu)造函數(shù)的最后一步。同時graph_module、graph、graph_signature、state_dict等屬性均為只讀setter 直接拋RuntimeError保證導出的程序不會被隨意改寫。graph_signature參數(shù)與 buffer 的戶口本ExportGraphSignature定義于 torch/export/graph_signature.py#L165-L275。它建模了 Export Graph 的輸入/輸出簽名核心不變量是Export Graph 是功能化的圖內(nèi)不通過getattr節(jié)點訪問參數(shù)、buffer 等狀態(tài)export保證參數(shù)、buffer、常量張量被提升為圖的輸入對 buffer 的改寫也不直接出現(xiàn)在圖內(nèi)而是把改寫后的 buffer 值建模為圖額外的輸出。輸入與輸出的排序規(guī)則為Inputs [*parameters_buffers_constant_tensors, *flattened_user_inputs] Outputs [*mutated_inputs, *flattened_user_outputs]InputKindPARAMETER、BUFFER、CONSTANT_TENSOR、USER_INPUT、CUSTOM_OBJ、TOKEN等與OutputKindUSER_OUTPUT、BUFFER_MUTATION、PARAMETER_MUTATION、USER_INPUT_MUTATION、GRADIENT_TO_PARAMETER等分別枚舉輸入、輸出的類別parameters、buffers、non_persistent_buffers、lifted_tensor_constants、user_inputs、user_outputs等屬性則從input_specs/output_specs中按類別抽取對應(yīng)名稱集合見 graph_signature.py#L277-L374。例如一個同時含參數(shù)、兩個 buffer 并在 forward 中原地改寫 buffer 的模塊其 graph_signature 會呈現(xiàn)為# inputs p_my_parameter: PARAMETER targetmy_parameter b_my_buffer1: BUFFER targetmy_buffer1 persistentTrue b_my_buffer2: BUFFER targetmy_buffer2 persistentTrue x1: USER_INPUT x2: USER_INPUT # outputs add_2: BUFFER_MUTATION targetmy_buffer2 add_1: USER_OUTPUT該示例完整代碼與圖對照見 graph_signature.py#L183-L271其中run_decompositions()會把非功能化圖含aten.add_.Tensor原地算子轉(zhuǎn)換為功能化圖buffer 改寫變?yōu)轭~外輸出add_2。這正是理解Export IR 為什么能放心交給后端執(zhí)行的關(guān)鍵圖內(nèi)不再有隱藏狀態(tài)。GraphDAG 與節(jié)點構(gòu)成一個 Export IR Graph 是 DAG有向無環(huán)圖形式的 PyTorch 程序每個節(jié)點代表一次計算/操作邊由節(jié)點間的引用構(gòu)成。其 schema 極簡class Graph: nodes: List[Node]實踐中 Export IR 的圖即torch.fx.GraphPython 類。一個合法的 Export IR 圖必須包含0 個或多個 op 類型為placeholder的節(jié)點0 個或多個 op 類型為call_function的節(jié)點恰好 1 個op 類型為output的節(jié)點。推論最小的合法 Graph 只有一個節(jié)點即nodes永不為空。定義placeholder節(jié)點集合代表 Graph/GraphModule 的輸入output節(jié)點代表其輸出。一個最簡可運行示例對應(yīng)文檔原始示例import torch from torch import nn class MyModule(nn.Module): def forward(self, x, y): return x y example_args (torch.randn(1), torch.randn(1)) mod torch.export.export(MyModule(), example_args) print(mod.graph)輸出Graph 的文本表示每行一個節(jié)點graph(): %x : [num_users1] placeholder[targetx] %y : [num_users1] placeholder[targety] %add : [num_users1] call_functiontargettorch.ops.aten.add.Tensor, kwargs {}) return (add,)Node節(jié)點 Schema 與 FX 文本格式Node 表示一次具體計算/操作Python 層面用torch.fx.Node類表示。節(jié)點間的邊通過Node.args屬性中對其他節(jié)點的直接引用實現(xiàn)借助 FX 機制可表達計算圖所需的操作符調(diào)用、占位符輸入、條件分支與循環(huán)等。Node 的 schemaclass Node: name: str # name of node op_name: str # type of operation # interpretation of the fields below depends on op_name target: [str|Callable] args: List[object] kwargs: Dict[str, object] meta: Dict[str, object]FX 文本格式如上例所示文本格式的每行遵循%name:[...] op_nametargettarget), kwargs {keyword: arg5})該格式以緊湊形式囊括了 Node 類的全部字段meta除外。具體地name即node.nameop_name即node.op字段必須是call_function、placeholder、get_attr、output之一target即node.target其含義取決于op_nameargs即node.args元組中的內(nèi)容若列表中某值是torch.fx.Node會特別以%前綴標識。例如一個 add 算子的調(diào)用顯示為%add1 call_functiontarget torch.op.aten.add.Tensor, kwargs {})其中%x、%y是兩個名字為 x、y 的 Node。值得注意字符串torch.op.aten.add.Tensor表示的是target 字段中實際存儲的可調(diào)用對象而非僅僅是它的字符串名。文本格式的最后一行return [add]對應(yīng)op_name output的節(jié)點表示返回這一元素。call_function算子調(diào)用節(jié)點call_function節(jié)點表示對某個算子的調(diào)用。兩個關(guān)鍵定義Functional功能化一個可調(diào)用對象滿足以下全部條件即為 functional——非改寫不修改輸入的值對張量而言包括元數(shù)據(jù)與數(shù)據(jù)且無副作用不改寫外部可見的狀態(tài)如模塊參數(shù)值Operator算子具有預定義 schema 的 functional 可調(diào)用對象典型如功能化 ATen 算子。FX 表示%name call_functiontarget operator, kwargs {})與 vanilla FXcall_function的差異FX 圖中call_function可指向任意可調(diào)用對象Export IR 則限定為 ATen 算子的選定子集、自定義算子與控制流算子Export IR 中常量參數(shù)會被內(nèi)嵌進圖內(nèi)FX 圖中g(shù)et_attr節(jié)點可讀取圖模塊中存儲的任何屬性Export IR 中它被限定為只能讀取子模塊因為所有參數(shù)/buffer 都以輸入形式傳入圖模塊。call_function 的元數(shù)據(jù)契約Node.meta是掛在每個 FX 節(jié)點上的 dict但 FX 規(guī)范并不規(guī)定其中會有哪些元數(shù)據(jù)。Export IR 提供了更強的契約所有call_function節(jié)點保證恰好包含以下字段node.meta[stack_trace]字符串引用原始 Python 源碼的調(diào)用棧例如File my_module.py, line 19, in forward return x dummy_helper(y) File helper_utility.py, line 89, in dummy_helper return y 1node.meta[val]描述該操作運行后的輸出類型可為symint、FakeTensor、List[Union[FakeTensor, SymInt]]或Nonenode.meta[nn_module_stack]該節(jié)點來源的torch.nn.Module的調(diào)用棧。例如addmm算子出自nn.Sequential內(nèi)的nn.Linear則表現(xiàn)為{self_linear: (self.linear, class torch.nn.Linear), self_sequential: (self.sequential, class torch.nn.Sequential)}node.meta[source_fn_stack]分解decomposition之前該節(jié)點來自的 torch 函數(shù)或葉子torch.nn.Module類。例如addmm來自nn.Linear模塊調(diào)用則source_fn為torch.nn.Linear來自torch.nn.functional.linear函數(shù)調(diào)用則包含torch.nn.functional.Linear。這些元數(shù)據(jù)使后端能夠在編譯期獲得源碼定位、輸出類型形狀、模塊層級與原始算子來源——這正是 Export IR 能支撐高質(zhì)量錯誤報告與算子級優(yōu)化的基礎(chǔ)。placeholder圖輸入節(jié)點placeholder表示圖的輸入語義與 FX 完全一致。占位節(jié)點必須是圖 nodes 列表中最靠前的 N 個節(jié)點N 可以為 0。FX 表示%name placeholdertarget name)target字段是輸入名稱的字符串a(chǎn)rgs若非空則大小為 1表示該輸入的默認值。元數(shù)據(jù)placeholder 節(jié)點同樣有meta[val]與call_function一致此時val表示圖期望為該輸入?yún)?shù)接收的輸入 shape/dtype。output圖出口節(jié)點output調(diào)用對應(yīng)函數(shù)中的 return 語句因此它終止當前圖。圖中有且僅有一個output 節(jié)點且它永遠是圖的最后一個節(jié)點。FX 表示output[](args (%something, …))語義與torch.fx完全一致args表示要返回的節(jié)點。output 節(jié)點的元數(shù)據(jù)與call_function節(jié)點相同。get_attr子模塊讀取節(jié)點get_attr節(jié)點表示從封裝它的torch.fx.GraphModule中讀取一個子模塊。與torch.fx.symbolic_trace得到的 vanilla FX 圖不同——后者用get_attr從頂層 GraphModule 讀取參數(shù)、buffer 等屬性——Export IR 中參數(shù)與 buffer 均作為輸入傳入圖模塊并存儲在頂層ExportedProgram中因此get_attr在這里僅用于讀取子模塊典型場景是控制流分支的嵌套子圖。FX 表示%name get_attrtarget name)示例使用functorch.experimental.control_flow.cond構(gòu)造條件分支from functorch.experimental.control_flow import cond def true_fn(x): return x.sin() def false_fn(x): return x.cos() def f(x, y): return cond(y, true_fn, false_fn, [x])對應(yīng)圖graph(): %x_1 : [num_users1] placeholder[targetx_1] %y_1 : [num_users1] placeholder[targety_1] %true_graph_0 : [num_users1] get_attr[targettrue_graph_0] %false_graph_0 : [num_users1] get_attr[targetfalse_graph_0] %conditional : [num_users1] call_functiontargettorch.ops.higher_order.cond, kwargs {}) return conditional其中%true_graph_0 get_attr[targettrue_graph_0]讀取包含sin算子的子模塊true_graph_0而torch.ops.higher_order.cond即為文檔所提及的控制流算子之一。引用類型SymInt、FakeTensor 與 Pytree-able 類型SymIntSymInt要么是字面整數(shù)要么是表示某個整數(shù)的符號Python 中由sympy.Symbol類表示。當 SymInt 是符號時它描述的是一個編譯期未知、僅運行時可知的整型變量——這正是動態(tài) shape 在 Export IR 中的載體。FakeTensorFakeTensor是包含張量元數(shù)據(jù)的對象可視為class FakeTensor: size: List[SymInt] dtype: torch.dtype device: torch.device dim_order: List[int] # 尚不存在size是整數(shù)或 SymInt 的列表若含 SymInt 表示該張量具有動態(tài)形狀若全是整數(shù)則假定張量具有該精確靜態(tài)形狀。張量元數(shù)據(jù)的秩rank永不是動態(tài)的dtype表示該節(jié)點輸出的數(shù)據(jù)類型Export IR 中不存在隱式類型提升no implicit type promotionsFakeTensor 中沒有 strides。換言之node.meta[val]與node.target算子的返回類型一一對應(yīng)算子返回 Tensor →val為描述該張量的 FakeTensor算子返回 n 元組張量 →val為 n 元 FakeTensor算子返回編譯期已知的 int/float/scalar →val為None算子返回編譯期未知的 int/float/scalar →val為 SymInt。例如aten::add返回 Tensor其 spec 即帶 dtype 與 size 的 FakeTensoraten::sym_size返回整數(shù)其 val 是 SymInt值僅運行時可得max_pool2d_with_indexes返回(Tensor, Tensor)元組spec 就是 2 元 FakeTensor分別描述返回值的兩個元素。Python 代碼示例def add_one(x): return torch.ops.aten(x, 1)對應(yīng)圖graph(): %ph_0 : [#users1] placeholder[targetph_0] %add_tensor : [#users1] call_functiontargettorch.ops.aten.add.Tensor, kwargs {}) return [add_tensor]FakeTensorFakeTensor(dtypetorch.int, size[2,], deviceCPU)Pytree-able 類型一個類型是Pytree-able的當且僅當它是葉子類型或包含其他 Pytree-able 類型的容器類型pytree 概念與 JAX 的 pytrees 文檔一致。葉子類型leaf type類型定義Tensortorch.TensorScalarPython 中任意數(shù)值類型包括整型、浮點型與零維張量intPython intC 中綁定為 int64_tfloatPython floatC 中綁定為 doubleboolPython boolstrPython 字符串ScalarTypetorch.dtypeLayouttorch.layoutMemoryFormattorch.memory_formatDevicetorch.device容器類型container type類型定義TuplePython tupleListPython listDict以 Scalar 為鍵的 Python dictNamedTuplePython namedtupleDataclass必須通過register_dataclass注冊實現(xiàn)見 torch/export/init.pyCustom class通過_register_pytree_node定義的任意自定義類實現(xiàn)見 torch/utils/_pytree.py該類型體系決定了 export 邊界上可合法傳遞的數(shù)據(jù)結(jié)構(gòu)用戶在torch.export.export()中傳入的args/kwargs、模型返回的結(jié)果最終都會被按 pytree 規(guī)則拍平為葉子節(jié)點再映射為圖的 placeholder 與 output。總結(jié)Export IR 的核心設(shè)計可以濃縮為三個關(guān)鍵詞功能化——參數(shù)、buffer 被提升為輸入buffer 改寫被建模為輸出圖中不存在隱藏狀態(tài)見 graph_signature.py#L166-L176嚴格元數(shù)據(jù)契約——stack_trace、val、nn_module_stack、source_fn_stack四字段為每個call_function節(jié)點背書編譯期即可獲得完整語義受限的算子集合與顯式類型——call_function僅限 ATen 算子子集、自定義算子與控制流算子常量內(nèi)嵌類型經(jīng) SymInt/FakeTensor 顯式刻畫無隱式提升。無論你是為 Export IR 編寫自定義后端還是調(diào)試torch.export.export()的輸出圖ir_spec.md 與其對應(yīng)源碼exported_program.py、graph_signature.py、torch/export/init.py都是最權(quán)威的第一手資料?！久赓M下載鏈接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration項目地址: https://gitcode.com/GitHub_Trending/py/pytorch創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考