基於 TorchDynamo 的 ONNX Exporter¶

概覽
依賴項
一個簡單的例子
使用 GUI 檢查 ONNX 模型
轉換失敗時
API 參考

警告

TorchDynamo 的 ONNX exporter 是一項快速發展的 beta 技術。

概覽 ¶

ONNX exporter 利用 TorchDynamo 引擎 Hook 進入 Python 的幀評估 API，並將其位元組碼動態重寫為 FX Graph。然後對生成的 FX Graph 進行精煉，最後將其轉換為 ONNX 圖。

這種方法的主要優點是 FX 圖是透過位元組碼分析捕獲的，它保留了模型的動態特性，而不是使用傳統的靜態追蹤技術。

此外，在匯出過程中，與啟用 TorchScript 的 exporter 相比，記憶體使用量顯著降低。更多資訊請參閱記憶體使用文件。

一個簡單的例子 ¶

下面是一個使用簡單的多層感知機 (MLP) 作為示例來演示 exporter 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.export(model, (tensor_x,), dynamo=True)

如上面的程式碼所示，您只需向 torch.onnx.export() 提供模型例項及其輸入。Exporter 將返回一個 torch.onnx.ONNXProgram 例項，其中包含匯出的 ONNX 圖以及額外資訊。

onnx_program.optimize() 可以用來最佳化 ONNX 圖，執行常量摺疊和冗餘運算子消除。最佳化是就地完成的。

onnx_program.optimize()

透過 onnx_program.model_proto 獲得的記憶體中的模型是一個符合 ONNX IR 規範的 onnx.ModelProto 物件。然後可以使用 torch.onnx.ONNXProgram.save() API 將 ONNX 模型序列化為 Protobuf 檔案。

onnx_program.save("mlp.onnx")

存在兩個基於 TorchDynamo 引擎將模型匯出為 ONNX 的函式。它們在生成 torch.export.ExportedProgram 的方式上略有不同。torch.onnx.dynamo_export() 在 PyTorch 2.1 中引入，而 torch.onnx.export() 在 PyTorch 2.5 中進行了擴充套件，可以方便地從 TorchScript 切換到 TorchDynamo。要呼叫前者，可以將上一個示例的最後一行替換為以下程式碼。

注意

torch.onnx.dynamo_export() 將在未來棄用。請改用帶有引數 dynamo=True 的 torch.onnx.export()。

onnx_program = torch.onnx.dynamo_export(model, tensor_x)

使用 GUI 檢查 ONNX 模型 ¶

您可以使用 Netron 檢視匯出的模型。

轉換失敗時 ¶

應第二次呼叫函式 torch.onnx.export()，並設定引數 report=True。將生成一份 markdown 報告，以幫助使用者解決問題。

API 參考 ¶

torch.onnx.dynamo_export(model, /, *model_args, export_options=None, **model_kwargs)[source][source]¶

將 torch.nn.Module 匯出為 ONNX 圖。

自版本 2.7 起已棄用: 請改用 torch.onnx.export(..., dynamo=True)。

引數

model (torch.nn.Module | Callable | torch.export.ExportedProgram) – 要匯出到 ONNX 的 PyTorch 模型。
model_args – model 的位置輸入。
model_kwargs – model 的關鍵字輸入。
export_options (ExportOptions | None) – 影響匯出到 ONNX 的選項。

返回值

匯出的 ONNX 模型的記憶體中表示。

返回型別

ONNXProgram

class torch.onnx.ONNXProgram(model, exported_program)¶

一個類，表示可使用 torch 張量呼叫的 ONNX 程式。

apply_weights(state_dict)[source][source]¶

將指定 state dict 中的權重應用於 ONNX 模型。

使用此方法替換 FakeTensors 或其他權重。

引數: state_dict (dict[str, torch.Tensor]) – 包含要應用於 ONNX 模型的權重的 state dict。

compute_values(value_names, args=(), kwargs=None)[source][source]¶

計算 ONNX 模型中指定名稱的值。

此方法用於計算 ONNX 模型中指定名稱的值。值將作為名稱到張量的字典返回。

引數: value_names (Sequence[str]) – 要計算的值的名稱。
返回值: 一個將名稱對映到張量的字典。
返回型別: Sequence[torch.Tensor]

initialize_inference_session(initializer=<function _ort_session_initializer>)[source][source]¶

初始化 ONNX Runtime 推理會話。

引數: initializer (Callable[[str | bytes], ort.InferenceSession]) – 用於使用指定模型初始化 ONNX Runtime 推理會話的函式。預設情況下，它使用 _ort_session_initializer() 函式。

property model_proto: ModelProto¶: 返回 ONNX ModelProto 物件。

optimize()[source][source]¶

最佳化 ONNX 模型。

此方法透過執行常量摺疊和消除圖中的冗餘來最佳化 ONNX 模型。最佳化是就地完成的。

release()[source][source]¶

釋放推理會話。

您可以呼叫此方法來釋放推理會話使用的資源。

save(destination, *, include_initializers=True, keep_initializers_as_inputs=False, external_data=None)[source][source]¶

將 ONNX 模型儲存到指定的目標位置。

當 external_data 為 True 或模型大於 2GB 時，權重將作為外部資料儲存到單獨的檔案中。

初始化器（模型權重）序列化行為：* include_initializers=True, keep_initializers_as_inputs=False (預設)：初始化器包含在儲存的模型中。* include_initializers=True, keep_initializers_as_inputs=True：初始化器包含在儲存的模型中，並保留為模型輸入。如果您希望在推理期間能夠覆蓋模型權重，請選擇此選項。* include_initializers=False, keep_initializers_as_inputs=False：初始化器不包含在儲存的模型中，也不列為模型輸入。如果您希望在單獨的後處理步驟中將初始化器附加到 ONNX 模型，請選擇此選項。* include_initializers=False, keep_initializers_as_inputs=True：初始化器不包含在儲存的模型中，但列為模型輸入。如果您希望在推理期間提供初始化器並最大程度地減小儲存模型的大小，請選擇此選項。

引數

destination (str | os.PathLike) – 儲存 ONNX 模型的路徑。
include_initializers (bool) – 是否在儲存的模型中包含初始化器。
keep_initializers_as_inputs (bool) – 是否在儲存的模型中將初始化器保留為輸入。如果為 True，初始化器將作為輸入新增到模型中，這意味著可以透過提供初始化器作為模型輸入來覆蓋它們。
external_data (bool | None) – 是否將權重儲存為單獨檔案中的外部資料。

引發

TypeError – 如果 external_data 為 True 且 destination 不是檔案路徑。

class torch.onnx.ExportOptions(*, dynamic_shapes=None, fake_context=None, onnx_registry=None, diagnostic_options=None)¶

影響 TorchDynamo ONNX exporter 的選項。

自版本 2.7 起已棄用: 請改用 torch.onnx.export(..., dynamo=True)。

變數

dynamic_shapes (bool | None) – 輸入/輸出張量的形狀資訊提示。當為 None 時，exporter 會確定最相容的設定。當為 True 時，所有輸入形狀都被視為動態。當為 False 時，所有輸入形狀都被視為靜態。
diagnostic_options (DiagnosticOptions) – exporter 的診斷選項。
fake_context (ONNXFakeContext | None) – 用於符號追蹤的 fake 上下文。
onnx_registry (OnnxRegistry | None) – 用於將 ATen 運算子註冊到 ONNX 函式的 ONNX 登錄檔。

torch.onnx.enable_fake_mode()[source]¶

在上下文期間啟用 fake 模式。

它內部例項化一個 torch._subclasses.fake_tensor.FakeTensorMode 上下文管理器，將使用者輸入和模型引數轉換為 torch._subclasses.fake_tensor.FakeTensor。

一個 torch._subclasses.fake_tensor.FakeTensor 是一個 torch.Tensor，它能夠在無需實際透過 meta 裝置上分配的張量進行計算的情況下執行 PyTorch 程式碼。由於裝置上沒有分配實際資料，此 API 允許初始化和匯出大型模型，而無需執行它所需的實際記憶體佔用。

強烈建議在匯出過大而無法完全載入到記憶體中的模型時，在 fake 模式下初始化模型。

注意

此函式不支援 torch.onnx.export(…, dynamo=True, optimize=True)。請在模型匯出後在函式外部呼叫 ONNXProgram.optimize()。

示例

# xdoctest: +REQUIRES(env:TORCH_DOCTEST_ONNX)
>>> import torch
>>> class MyModel(torch.nn.Module):  # Model with a parameter
...     def __init__(self) -> None:
...         super().__init__()
...         self.weight = torch.nn.Parameter(torch.tensor(42.0))
...     def forward(self, x):
...         return self.weight + x
>>> with torch.onnx.enable_fake_mode():
...     # When initialized in fake mode, the model's parameters are fake tensors
...     # They do not take up memory so we can initialize large models
...     my_nn_module = MyModel()
...     arg1 = torch.randn(2, 2, 2)
>>> onnx_program = torch.onnx.export(my_nn_module, (arg1,), dynamo=True, optimize=False)
>>> # Saving model WITHOUT initializers (only the architecture)
>>> onnx_program.save(
...     "my_model_without_initializers.onnx",
...     include_initializers=False,
...     keep_initializers_as_inputs=True,
... )
>>> # Saving model WITH initializers after applying concrete weights
>>> onnx_program.apply_weights({"weight": torch.tensor(42.0)})
>>> onnx_program.save("my_model_with_initializers.onnx")

警告

此 API 是實驗性的，且不向後相容。

class torch.onnx.ONNXRuntimeOptions(*, session_options=None, execution_providers=None, execution_provider_options=None)¶

透過 ONNX Runtime 影響 ONNX 模型執行的選項。

自版本 2.7 起已棄用: 請改用 torch.onnx.export(..., dynamo=True)。

變數

session_options (Sequence[onnxruntime.SessionOptions] | None) – ONNX Runtime 會話選項。
execution_providers (Sequence[str | tuple[str, dict[Any, Any]]] | None) – 在模型執行期間使用的 ONNX Runtime 執行提供者。
execution_provider_options (Sequence[dict[Any, Any]] | None) – ONNX Runtime 執行提供者選項。

class torch.onnx.OnnxExporterError¶: ONNX 匯出器引發的錯誤。這是所有匯出器錯誤的基類。

class torch.onnx.OnnxRegistry¶

ONNX 函式的登錄檔。

自版本 2.7 起已棄用: 請改用 torch.onnx.export(..., dynamo=True)。

該登錄檔維護了在固定 opset 版本下，從限定名稱到符號函式的對映。它支援註冊自定義 onnx-script 函式，並支援排程器將呼叫分派到適當的函式。

get_op_functions(namespace, op_name, overload=None)[source][source]¶

返回給定操作 torch.ops.<namespace>.<op_name>.<overload> 的 ONNXFunction 列表。

該列表按註冊時間排序。自定義操作應位於列表的後半部分。

引數

namespace (str) – 要獲取的操作的名稱空間。
op_name (str) – 要獲取的操作的名稱。
overload (str | None) – 要獲取的操作的過載。如果是預設過載，則設為 None。

返回值

與給定名稱對應的 ONNXFunction 列表，如果該名稱不在登錄檔中，則為 None。

返回型別

list[registration.ONNXFunction] | None

is_registered_op(namespace, op_name, overload=None)[source][source]¶

返回給定操作是否已註冊：torch.ops.<namespace>.<op_name>.<overload>。

引數

namespace (str) – 要檢查的操作的名稱空間。
op_name (str) – 要檢查的操作的名稱。
overload (str | None) – 要檢查的操作的過載。如果是預設過載，則設為 None。

返回值

如果給定操作已註冊，則為 True，否則為 False。

返回型別

bool

property opset_version: int¶: 匯出器應面向的 ONNX opset 版本。

register_op(function, namespace, op_name, overload=None, is_complex=False)[source][source]¶

註冊自定義操作：torch.ops.<namespace>.<op_name>.<overload>。

引數

function (onnxscript.OnnxFunction | onnxscript.TracedOnnxFunction) – 要註冊的 onnx-script 函式。
namespace (str) – 要註冊的操作的名稱空間。
op_name (str) – 要註冊的操作的名稱。
overload (str | None) – 要註冊的操作的過載。如果是預設過載，則設為 None。
is_complex (bool) – 函式是否處理複數值輸入。

引發

ValueError – 如果名稱格式不是 'namespace::op'。

class torch.onnx.DiagnosticOptions(verbosity_level=20, warnings_as_errors=False)¶

診斷上下文的選項。

變數

verbosity_level (int) – 設定每個診斷日誌的資訊量，相當於 Python 日誌模組中的“level”。
warnings_as_errors (bool) – 當為 True 時，將警告診斷視為錯誤診斷。

基於 TorchDynamo 的 ONNX Exporter¶

概覽 ¶

依賴項 ¶

一個簡單的例子 ¶

使用 GUI 檢查 ONNX 模型 ¶

轉換失敗時 ¶

API 參考 ¶

文件

教程

資源

基於 TorchDynamo 的 ONNX Exporter¶

概覽¶

依賴項¶

一個簡單的例子¶

使用 GUI 檢查 ONNX 模型¶

轉換失敗時¶

API 參考¶

文件

教程

資源

概覽 ¶

依賴項 ¶

一個簡單的例子 ¶

使用 GUI 檢查 ONNX 模型 ¶

轉換失敗時 ¶

API 參考 ¶