基於 TorchDynamo 的 ONNX 匯出器¶
警告
TorchDynamo 的 ONNX 匯出器是一項快速發展的 Beta 技術。
概述¶
ONNX 匯出器利用 TorchDynamo 引擎鉤入 Python 的框架評估 API,並將其位元組碼動態重寫為 FX 圖。產生的 FX 圖會先經過潤飾,最後才會轉換為 ONNX 圖。
這種方法的主要優點是,FX 圖 是使用位元組碼分析擷取的,這種分析保留了模型的動態特性,而不是使用傳統的靜態追蹤技術。
匯出器設計為模組化和可擴展的。它由以下元件組成
**ONNX 匯出器**:
Exporter主類別,協調匯出程序。**ONNX 匯出選項**:
ExportOptions有一組控制匯出程序的選項。**ONNX 登錄檔**:
OnnxRegistry是 ONNX 運算子和函數的登錄檔。**FX 圖擷取器**:
FXGraphExtractor從 PyTorch 模型中擷取 FX 圖。**模擬模式**:
ONNXFakeContext是一個上下文管理器,可為大型模型啟用模擬模式。**ONNX 程式**:
ONNXProgram是匯出器的輸出,包含匯出的 ONNX 圖和診斷資訊。**ONNX 程式序列化器**:
ONNXProgramSerializer將匯出的模型序列化為檔案。**ONNX 診斷選項**:
DiagnosticOptions有一組控制匯出器發出診斷資訊的選項。
一個簡單的範例¶
請參閱以下使用簡單的多層感知器 (MLP) 作為範例的匯出器 API 示範
import torch
import torch.nn as nn
class MLPModel(nn.Module):
def __init__(self):
super().__init__()
self.fc0 = nn.Linear(8, 8, bias=True)
self.fc1 = nn.Linear(8, 4, bias=True)
self.fc2 = nn.Linear(4, 2, bias=True)
self.fc3 = nn.Linear(2, 2, bias=True)
def forward(self, tensor_x: torch.Tensor):
tensor_x = self.fc0(tensor_x)
tensor_x = torch.sigmoid(tensor_x)
tensor_x = self.fc1(tensor_x)
tensor_x = torch.sigmoid(tensor_x)
tensor_x = self.fc2(tensor_x)
tensor_x = torch.sigmoid(tensor_x)
output = self.fc3(tensor_x)
return output
model = MLPModel()
tensor_x = torch.rand((97, 8), dtype=torch.float32)
onnx_program = torch.onnx.dynamo_export(model, tensor_x)
如上述程式碼所示,您只需要提供 torch.onnx.dynamo_export() 模型的實例及其輸入。然後,匯出器會傳回 torch.onnx.ONNXProgram 的實例,其中包含匯出的 ONNX 圖以及額外資訊。
透過 onnx_program.model_proto 可用的記憶體模型是符合 ONNX IR 規範 的 onnx.ModelProto 物件。然後可以使用 torch.onnx.ONNXProgram.save() API 將 ONNX 模型序列化為 Protobuf 檔案。
onnx_program.save("mlp.onnx")
使用 GUI 檢查 ONNX 模型¶
您可以使用 Netron 檢視匯出的模型。
請注意,每個圖層都以矩形框表示,右上角有一個 *f* 圖示。
展開它就會顯示函數主體。
函數主體是一系列 ONNX 運算子或其他函數。
使用 SARIF 診斷問題¶
ONNX 診斷透過採用 靜態分析結果交換格式 (又稱 SARIF) 來超越常規日誌,以協助使用者使用 GUI(例如 Visual Studio Code 的 SARIF 檢視器)對其模型進行偵錯和改進。
主要優點如下:
診斷資訊以機器可解析的 靜態分析結果交換格式 (SARIF) 發出。
一種全新的、更清晰、結構化的方式來添加新的診斷規則並追蹤現有規則。
作為未來更多改進的基礎,這些改進將使用診斷資訊。
ONNX 診斷 SARIF 規則
- FXE0007:fx-graph-to-onnx
- FXE0008:fx-node-to-onnx
- FXE0010:fx-pass
- FXE0011:no-symbolic-function-for-call-function
- FXE0012:unsupported-fx-node-analysis
- FXE0013:op-level-debugging
- FXE0014:find-opschema-matched-symbolic-function
- FXE0015:fx-node-insert-type-promotion
- FXE0016:find-operator-overloads-in-onnx-registry
API 參考¶
- torch.onnx.dynamo_export(model, /, *model_args, export_options=None, **model_kwargs)¶
將 torch.nn.Module 匯出至 ONNX 圖形。
- 參數
model (Union[Module, Callable, ExportedProgram]) – 要匯出至 ONNX 的 PyTorch 模型。
model_args –
model的位置輸入。model_kwargs –
model的關鍵字輸入。export_options (Optional[ExportOptions]) – 影響匯出至 ONNX 的選項。
- 回傳值
匯出 ONNX 模型的記憶體中表示形式。
- 回傳類型
範例 1 - 最簡單的匯出
class MyModel(torch.nn.Module): def __init__(self) -> None: super().__init__() self.linear = torch.nn.Linear(2, 2) def forward(self, x, bias=None): out = self.linear(x) out = out + bias return out model = MyModel() kwargs = {"bias": 3.} args = (torch.randn(2, 2, 2),) onnx_program = torch.onnx.dynamo_export( model, *args, **kwargs).save("my_simple_model.onnx")
範例 2 - 使用動態形狀匯出
# The previous model can be exported with dynamic shapes export_options = torch.onnx.ExportOptions(dynamic_shapes=True) onnx_program = torch.onnx.dynamo_export( model, *args, **kwargs, export_options=export_options) onnx_program.save("my_dynamic_model.onnx")
透過列印輸入動態維度,我們可以看到輸入形狀不再是 (2,2,2)
>>> print(onnx_program.model_proto.graph.input[0]) name: "arg0" type { tensor_type { elem_type: 1 shape { dim { dim_param: "arg0_dim_0" } dim { dim_param: "arg0_dim_1" } dim { dim_param: "arg0_dim_2" } } } }
- class torch.onnx.ExportOptions(*, dynamic_shapes=None, op_level_debug=None, fake_context=None, onnx_registry=None, diagnostic_options=None)¶
影響 TorchDynamo ONNX 匯出器的選項。
- 變數
dynamic_shapes (Optional[bool]) – 輸入/輸出張量的形狀資訊提示。當為
None時,匯出器會決定最相容的設定。當為True時,所有輸入形狀都被視為動態的。當為False時,所有輸入形狀都被視為靜態的。op_level_debug (Optional[bool]) – 是否匯出具有操作級別偵錯資訊的模型
diagnostic_options (DiagnosticOptions) – 匯出器的診斷選項。
fake_context (Optional[ONNXFakeContext]) – 用於符號追蹤的虛擬上下文。
onnx_registry (Optional[OnnxRegistry]) – 用於將 ATen 運算子註冊到 ONNX 函數的 ONNX 註冊表。
- torch.onnx.enable_fake_mode()¶
在上下文期間啟用虛擬模式。
在內部,它會實例化一個
torch._subclasses.fake_tensor.FakeTensorMode上下文管理器,將使用者輸入和模型參數轉換為torch._subclasses.fake_tensor.FakeTensor。torch._subclasses.fake_tensor.FakeTensor是一個torch.Tensor,它能夠在無需透過分配在meta裝置上的張量實際執行計算的情況下執行 PyTorch 程式碼。由於裝置上沒有實際分配資料,因此此 API 允許匯出大型模型,而無需執行它所需的實際記憶體佔用空間。強烈建議在匯出因太大而無法放入記憶體的模型時啟用虛擬模式。
- 回傳值
一個
ONNXFakeContext物件,必須透過ExportOptions.fake_context參數傳遞給dynamo_export()。
範例
# xdoctest: +REQUIRES(env:TORCH_DOCTEST_ONNX) >>> import torch >>> import torch.onnx >>> class MyModel(torch.nn.Module): # Dummy model ... def __init__(self) -> None: ... super().__init__() ... self.linear = torch.nn.Linear(2, 2) ... def forward(self, x): ... out = self.linear(x) ... return out >>> with torch.onnx.enable_fake_mode() as fake_context: ... my_nn_module = MyModel() ... arg1 = torch.randn(2, 2, 2) # positional input 1 >>> export_options = torch.onnx.ExportOptions(fake_context=fake_context) >>> onnx_program = torch.onnx.dynamo_export( ... my_nn_module, ... arg1, ... export_options=export_options ... ) >>> # Saving model WITHOUT initializers >>> onnx_program.save("my_model_without_initializers.onnx") >>> # Saving model WITH initializers >>> onnx_program.save("my_model_with_initializers.onnx", model_state=MyModel().state_dict())
警告
此 API 為實驗性 API,且不向後相容。
- class torch.onnx.ONNXProgram(model_proto, input_adapter, output_adapter, diagnostic_context, *, fake_context=None, export_exception=None, model_signature=None, model_torch=None)¶
已匯出至 ONNX 的 PyTorch 模型的記憶體中表示形式。
- 參數
model_proto (onnx.ModelProto) – 匯出的 ONNX 模型,格式為
onnx.ModelProto。input_adapter (io_adapter.InputAdapter) – 用於將 PyTorch 輸入轉換為 ONNX 輸入的輸入適配器。
output_adapter (io_adapter.OutputAdapter) – 用於將 PyTorch 輸出轉換為 ONNX 輸出的輸出適配器。
diagnostic_context (diagnostics.DiagnosticContext) – SARIF 診斷系統的上下文物件,負責記錄錯誤和中繼資料。
fake_context (Optional[ONNXFakeContext]) – 用於符號追蹤的虛擬上下文。
export_exception (Optional[Exception]) – 匯出期間發生的例外狀況(如果有)。
model_signature (Optional[torch.export.ExportGraphSignature]) – 匯出 ONNX 圖形的模型簽章。
- adapt_torch_inputs_to_onnx(*model_args, model_with_state_dict=None, **model_kwargs)[原始碼]¶
將 PyTorch 模型輸入轉換為匯出的 ONNX 模型輸入格式。
由於設計差異,PyTorch 模型和匯出的 ONNX 模型之間的輸入/輸出格式通常不同。例如,PyTorch 模型允許 None,但 ONNX 不支援。PyTorch 模型允許張量的巢狀結構,但 ONNX 僅支援扁平化張量,等等。
實際的調整步驟與每個匯出相關聯。這取決於 PyTorch 模型、用於匯出的特定 model_args 和 model_kwargs 集,以及匯出選項。
此方法會重播匯出期間記錄的調整步驟。
- 參數
model_args – PyTorch 模型輸入。
model_with_state_dict (Optional[Union[Module, Callable, ExportedProgram]]) – 要從中獲取額外狀態的 PyTorch 模型。如果未指定,則使用匯出期間使用的模型。當使用
enable_fake_mode()提取 ONNX 圖形所需的實際初始化器時需要。model_kwargs – PyTorch 模型關鍵字輸入。
- 回傳值
從 PyTorch 模型輸入轉換而來的張量序列。
- 回傳類型
範例
# xdoctest: +REQUIRES(env:TORCH_DOCTEST_ONNX) >>> import torch >>> import torch.onnx >>> from typing import Dict, Tuple >>> def func_nested_input( ... x_dict: Dict[str, torch.Tensor], ... y_tuple: Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] ... ): ... if "a" in x_dict: ... x = x_dict["a"] ... elif "b" in x_dict: ... x = x_dict["b"] ... else: ... x = torch.randn(3) ... ... y1, (y2, y3) = y_tuple ... ... return x + y1 + y2 + y3 >>> x_dict = {"a": torch.tensor(1.)} >>> y_tuple = (torch.tensor(2.), (torch.tensor(3.), torch.tensor(4.))) >>> onnx_program = torch.onnx.dynamo_export(func_nested_input, x_dict, y_tuple) >>> print(x_dict, y_tuple) {'a': tensor(1.)} (tensor(2.), (tensor(3.), tensor(4.))) >>> print(onnx_program.adapt_torch_inputs_to_onnx(x_dict, y_tuple, model_with_state_dict=func_nested_input)) (tensor(1.), tensor(2.), tensor(3.), tensor(4.))
警告
此 API 為實驗性 API,且不向後相容。
- adapt_torch_outputs_to_onnx(model_outputs, model_with_state_dict=None)[原始碼]¶
將 PyTorch 模型輸出轉換為匯出的 ONNX 模型輸出格式。
由於設計差異,PyTorch 模型和匯出的 ONNX 模型之間的輸入/輸出格式通常不同。例如,PyTorch 模型允許 None,但 ONNX 不支援。PyTorch 模型允許張量的巢狀結構,但 ONNX 僅支援扁平化張量,等等。
實際的調整步驟與每個匯出相關聯。這取決於 PyTorch 模型、用於匯出的特定 model_args 和 model_kwargs 集,以及匯出選項。
此方法會重播匯出期間記錄的調整步驟。
- 參數
model_outputs (Any) – PyTorch 模型輸出。
model_with_state_dict (Optional[Union[Module, Callable, ExportedProgram]]) – 要從中獲取額外狀態的 PyTorch 模型。如果未指定,則使用匯出期間使用的模型。當使用
enable_fake_mode()提取 ONNX 圖形所需的實際初始化器時需要。
- 回傳值
匯出的 ONNX 模型輸出格式的 PyTorch 模型輸出。
- 回傳類型
範例
# xdoctest: +REQUIRES(env:TORCH_DOCTEST_ONNX) >>> import torch >>> import torch.onnx >>> def func_returning_tuples(x, y, z): ... x = x + y ... y = y + z ... z = x + y ... return (x, (y, z)) >>> x = torch.tensor(1.) >>> y = torch.tensor(2.) >>> z = torch.tensor(3.) >>> onnx_program = torch.onnx.dynamo_export(func_returning_tuples, x, y, z) >>> pt_output = func_returning_tuples(x, y, z) >>> print(pt_output) (tensor(3.), (tensor(5.), tensor(8.))) >>> print(onnx_program.adapt_torch_outputs_to_onnx(pt_output, model_with_state_dict=func_returning_tuples)) [tensor(3.), tensor(5.), tensor(8.)]
警告
此 API 為實驗性 API,且不向後相容。
- property diagnostic_context: diagnostics.DiagnosticContext¶
與匯出相關聯的診斷上下文。
- property model_proto: onnx.ModelProto¶
匯出的 ONNX 模型,格式為
onnx.ModelProto。
- property model_signature: Optional[ExportGraphSignature]¶
匯出的 ONNX 圖形的模型簽章。
此資訊非常重要,因為 ONNX 規範通常與 PyTorch 的規範不同,導致 ONNX 圖形的輸入和輸出結構與實際 PyTorch 模型實作不同。透過使用模型簽章,使用者可以瞭解輸入和輸出的差異,並在 ONNX Runtime 中正確執行模型。
注意:只有當 ONNX 圖形是從
torch.export.ExportedProgram物件匯出時,模型簽章才可用。注意:對模型進行的任何會改變模型簽章的轉換,都必須透過
InputAdaptStep和/或OutputAdaptStep更新此模型簽章。範例
以下模型會產生不同的輸入和輸出集。前 4 個輸入是模型參數(即 conv1.weight、conv2.weight、fc1.weight、fc2.weight),接下來的 2 個輸入是已註冊的緩衝區(即 my_buffer2、my_buffer1),最後 2 個輸入是使用者輸入(即 x 和 b)。第一個輸出是緩衝區變化(即 my_buffer2),最後一個輸出是實際的模型輸出。
>>> import pprint >>> class CustomModule(torch.nn.Module): ... def __init__(self): ... super().__init__() ... self.my_parameter = torch.nn.Parameter(torch.tensor(2.0)) ... self.register_buffer("my_buffer1", torch.tensor(3.0)) ... self.register_buffer("my_buffer2", torch.tensor(4.0)) ... self.conv1 = torch.nn.Conv2d(1, 32, 3, 1, bias=False) ... self.conv2 = torch.nn.Conv2d(32, 64, 3, 1, bias=False) ... self.fc1 = torch.nn.Linear(9216, 128, bias=False) ... self.fc2 = torch.nn.Linear(128, 10, bias=False) ... def forward(self, x, b): ... tensor_x = self.conv1(x) ... tensor_x = torch.nn.functional.sigmoid(tensor_x) ... tensor_x = self.conv2(tensor_x) ... tensor_x = torch.nn.functional.sigmoid(tensor_x) ... tensor_x = torch.nn.functional.max_pool2d(tensor_x, 2) ... tensor_x = torch.flatten(tensor_x, 1) ... tensor_x = self.fc1(tensor_x) ... tensor_x = torch.nn.functional.sigmoid(tensor_x) ... tensor_x = self.fc2(tensor_x) ... output = torch.nn.functional.log_softmax(tensor_x, dim=1) ... ( ... self.my_buffer2.add_(1.0) + self.my_buffer1 ... ) # Mutate buffer through in-place addition ... return output >>> inputs = (torch.rand((64, 1, 28, 28), dtype=torch.float32), torch.randn(3)) >>> exported_program = torch.export.export(CustomModule(), args=inputs).run_decompositions({}) >>> onnx_program = torch.onnx.dynamo_export(exported_program, *inputs) >>> pprint.pprint(onnx_program.model_signature) ExportGraphSignature(input_specs=[InputSpec(kind=<InputKind.PARAMETER: 2>, arg=TensorArgument(name='p_conv1_weight'), target='conv1.weight', persistent=None), InputSpec(kind=<InputKind.PARAMETER: 2>, arg=TensorArgument(name='p_conv2_weight'), target='conv2.weight', persistent=None), InputSpec(kind=<InputKind.PARAMETER: 2>, arg=TensorArgument(name='p_fc1_weight'), target='fc1.weight', persistent=None), InputSpec(kind=<InputKind.PARAMETER: 2>, arg=TensorArgument(name='p_fc2_weight'), target='fc2.weight', persistent=None), InputSpec(kind=<InputKind.BUFFER: 3>, arg=TensorArgument(name='b_my_buffer2'), target='my_buffer2', persistent=True), InputSpec(kind=<InputKind.BUFFER: 3>, arg=TensorArgument(name='b_my_buffer1'), target='my_buffer1', persistent=True), InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='x'), target=None, persistent=None), InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='b'), target=None, persistent=None)], output_specs=[OutputSpec(kind=<OutputKind.BUFFER_MUTATION: 3>, arg=TensorArgument(name='add'), target='my_buffer2'), OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='_log_softmax'), target=None)])
- save(destination, *, include_initializers=True, model_state=None, serializer=None)[原始碼]¶
使用指定的
serializer將記憶體中的 ONNX 模型儲存到destination。- 參數
destination (Union[str, BufferedIOBase]) – 儲存 ONNX 模型的目的地。它可以是字串或檔案物件。與
model_state一起使用時,它必須是具有完整目的地路徑的字串。如果 destination 是字串,除了將 ONNX 模型儲存到檔案之外,模型權重也會儲存在與 ONNX 模型相同的目錄中的個別檔案中。例如,對於 destination=”/path/model.onnx”,初始設定會與「onnx.model」一起儲存在「/path/」資料夾中。include_initializers (bool) – 是否在 ONNX 圖形中包含初始設定作為外部資料。無法與 model_state_dict 結合使用。
model_state (Optional[Union[Dict[str, Any], str]]) – PyTorch 模型的 state_dict,其中包含其上的所有權重。它可以是具有檢查點路徑的字串,也可以是具有實際模型狀態的字典。支援的檔案格式與 torch.load 和 safetensors.safe_open 支援的格式相同。當使用
enable_fake_mode()但在 ONNX 圖形上需要實際初始設定時,這是必需的。serializer (Optional[ONNXProgramSerializer]) – 要使用的序列化器。如果未指定,模型將序列化為 Protobuf。
- save_diagnostics(destination)[原始碼]¶
將匯出診斷資訊以 SARIF 記錄檔儲存到指定的目的地路徑。
- 參數
destination (str) – 儲存診斷資訊 SARIF 記錄檔的目的地。它必須具有 .sarif 副檔名。
- 引發
ValueError – 如果目的地路徑不是以 .sarif 副檔名結尾。
- class torch.onnx.ONNXProgramSerializer(*args, **kwargs)¶
將 ONNX 圖形序列化為特定格式(例如 Protobuf)的協定。請注意,這是一個進階的使用情境。
- serialize(onnx_program, destination)[原始碼]¶
必須為序列化實作的協定方法。
- 參數
onnx_program (ONNXProgram) – 表示記憶體中匯出的 ONNX 模型
destination (BufferedIOBase) – 二進制 IO 串流或預先配置的緩衝區,序列化模型應寫入其中。
範例
一個簡單的序列化器,它以 Protobuf 格式將匯出的
onnx.ModelProto寫入destination# xdoctest: +REQUIRES(env:TORCH_DOCTEST_ONNX) >>> import io >>> import torch >>> import torch.onnx >>> class MyModel(torch.nn.Module): # Dummy model ... def __init__(self) -> None: ... super().__init__() ... self.linear = torch.nn.Linear(2, 2) ... def forward(self, x): ... out = self.linear(x) ... return out >>> class ProtobufONNXProgramSerializer: ... def serialize( ... self, onnx_program: torch.onnx.ONNXProgram, destination: io.BufferedIOBase ... ) -> None: ... destination.write(onnx_program.model_proto.SerializeToString()) >>> model = MyModel() >>> arg1 = torch.randn(2, 2, 2) # positional input 1 >>> torch.onnx.dynamo_export(model, arg1).save( ... destination="exported_model.onnx", ... serializer=ProtobufONNXProgramSerializer(), ... )
- class torch.onnx.ONNXRuntimeOptions(*, session_options=None, execution_providers=None, execution_provider_options=None)¶
透過 ONNX Runtime 影響 ONNX 模型執行的選項。
- class torch.onnx.InvalidExportOptionsError¶
當使用者為
ExportOptions指定了無效的值時引發。
- class torch.onnx.OnnxExporterError(onnx_program, message)¶
當 ONNX 轉換器發生錯誤時引發。
當 ONNX 匯出過程中發生錯誤時,會拋出此異常。它封裝了直到失敗為止生成的
ONNXProgram物件,允許訪問部分匯出結果和相關聯的元數據。
- class torch.onnx.OnnxRegistry¶
ONNX 函數的註冊表。
註冊表維護在固定 opset 版本下從限定名稱到符號函數的映射。它支持註冊自定義 onnx-script 函數,並支持調度器將調用分派到適當的函數。
- get_op_functions(namespace, op_name, overload=None)[原始碼]¶
傳回給定操作的 ONNXFunction 列表:torch.ops.<namespace>.<op_name>.<overload>。
列表按註冊時間排序。自定義運算符應位於列表的後半部分。
- is_registered_op(namespace, op_name, overload=None)[原始碼]¶
傳回給定操作是否已註冊:torch.ops.<namespace>.<op_name>.<overload>。