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