基於 TorchDynamo 的 ONNX 匯出器¶

概觀
依賴關係
一個簡單的範例
使用 GUI 檢查 ONNX 模型
轉換失敗時
API 參考

警告

用於 TorchDynamo 的 ONNX 匯出器是一項快速發展的 beta 技術。

概觀 ¶

ONNX 匯出器利用 TorchDynamo 引擎來連接 Python 的 frame evaluation API，並動態地將其 bytecode 重寫為 FX Graph。產生的 FX Graph 會先經過潤飾，最後才會轉換成 ONNX 圖形。

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

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

一個簡單的例子 ¶

以下是一個簡單的多層感知器 (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.export(model, (tensor_x,), dynamo=True)

如上面的程式碼所示，您只需要提供 torch.onnx.export() 模型和其輸入的實例。然後，匯出器將回傳一個 torch.onnx.ONNXProgram 的實例，其中包含匯出的 ONNX 圖形以及額外的資訊。

可以呼叫 onnx_program.optimize() 來使用常數摺疊 (constant folding) 和消除冗餘運算子來優化 ONNX 圖形。最佳化是就地 (in-place) 完成的。

onnx_program.optimize()

通過 onnx_program.model_proto 提供的記憶體模型是一個符合 ONNX IR spec 的 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() 將在未來被棄用。請改用 torch.onnx.export() 並帶有參數 dynamo=True。

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 圖形。

參數

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

回傳

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

回傳類型

ONNXProgram

範例 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.0}
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")

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

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

apply_weights(state_dict)[source][source]¶

將指定狀態字典中的權重應用於 ONNX 模型。

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

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

initialize_inference_session(initializer=<function _ort_session_initializer>)[原始碼][原始碼]¶

初始化 ONNX Runtime 推論階段。

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

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

optimize()[原始碼][原始碼]¶

最佳化 ONNX 模型。

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

release()[原始碼][原始碼]¶

釋放推論階段。

您可以呼叫此方法來釋放推論階段使用的資源。

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

將 ONNX 模型儲存到指定的目的地。

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

Initializer (模型權重) 序列化行為： * include_initializers=True, keep_initializers_as_inputs=False (預設)：initializer 包含在已儲存的模型中。 * include_initializers=True, keep_initializers_as_inputs=True：initializer 包含在已儲存的模型中，並保留為模型輸入。如果您希望在推論期間能夠覆寫模型權重，請選擇此選項。 * include_initializers=False, keep_initializers_as_inputs=False：initializer 不包含在已儲存的模型中，也不會列為模型輸入。如果您希望在單獨的後處理步驟中將 initializer 附加到 ONNX 模型，請選擇此選項。 * include_initializers=False, keep_initializers_as_inputs=True：initializer 不包含在已儲存的模型中，但會列為模型輸入。如果您希望在推論期間提供 initializer，並且想要最小化已儲存模型的大小，請選擇此選項。

參數

destination (str | os.PathLike) – 儲存 ONNX 模型的路徑。
include_initializers (bool) – 是否在儲存的模型中包含 initializer。
keep_initializers_as_inputs (bool) – 是否在儲存的模型中將 initializer 保留為輸入。如果為 True，則將 initializer 作為輸入新增到模型，這表示可以覆寫它們。透過提供 initializer 作為模型輸入。
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 匯出器的選項。

變數

dynamic_shapes (bool | None) – 輸入/輸出張量的形狀資訊提示。當 None 時，匯出器會決定最相容的設定。當 True 時，所有輸入形狀都被視為動態。當 False 時，所有輸入形狀都被視為靜態。
diagnostic_options (DiagnosticOptions) – 匯出器的診斷選項。
fake_context (ONNXFakeContext | None) – 用於符號追蹤的虛擬上下文。
onnx_registry (OnnxRegistry | None) – 用於將 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，具有運行 PyTorch 程式碼的能力，而無需透過在 meta 裝置上分配的張量來實際進行計算。由於沒有在裝置上分配實際資料，因此此 API 允許初始化和匯出大型模型，而無需執行它所需的實際記憶體佔用空間。

強烈建議在匯出過大而無法放入記憶體的模型時，在虛擬模式下初始化模型。

範例

# 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)
>>> # 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 模型執行的選項。

變數

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 函式的註冊表。

該註冊表維護從限定名稱到固定 opset 版本下的符號函式的映射。它支援註冊自訂 onnx-script 函式，並支援 dispatcher 將呼叫分派到適當的函式。

get_op_functions(namespace, op_name, overload=None)[原始碼][原始碼]¶

傳回給定運算子的 ONNXFunction 清單：torch.ops.<namespace>.<op_name>.<overload>。

此清單依註冊時間排序。自訂運算子應位於清單的後半部分。

參數

namespace (str) – 要取得的運算子的命名空間。
op_name (str) – 要取得的運算子的名稱。
overload (str | None) – 要取得的運算子的重載。如果它是預設重載，請將其保留為 None。

回傳

與給定名稱對應的 ONNXFunction 清單，如果該名稱不在註冊表中，則為 None。

回傳類型

list[registration.ONNXFunction] | None

is_registered_op(namespace, op_name, overload=None)[原始碼][原始碼]¶

傳回是否已註冊給定的運算子：torch.ops.<namespace>.<op_name>.<overload>。

參數

namespace (str) – 要檢查的運算子的命名空間。
op_name (str) – 要檢查的運算子的名稱。
overload (str | None) – 要檢查的運算子多載 (overload)。如果是預設多載，請保留為 None。

回傳

如果給定的 op 已註冊，則為 True，否則為 False。

回傳類型

bool

property opset_version: int¶: 匯出工具應針對的 ONNX opset 版本。

register_op(function, namespace, op_name, overload=None, is_complex=False)[原始碼][原始碼]¶

註冊自訂運算子：torch.ops.<namespace>.<op_name>.<overload>。

參數

function (onnxscript.OnnxFunction | onnxscript.TracedOnnxFunction) – 要註冊的 onnx-sctip 函式。
namespace (str) – 要註冊的運算子的命名空間。
op_name (str) – 要註冊的運算子的名稱。
overload (str | None) – 要註冊的運算子的多載 (overload)。如果是預設多載，請保留為 None。
is_complex (bool) – 該函式是否為處理複數值輸入的函式。

引發

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

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

用於診斷內容的選項。

變數

verbosity_level (int) – 設定每個診斷記錄的資訊量，相當於 Python logging 模組中的 'level'。
warnings_as_errors (bool) – 如果為 True，則警告診斷會被視為錯誤診斷。

基於 TorchDynamo 的 ONNX 匯出器¶

概觀 ¶

相依性 ¶

一個簡單的例子 ¶

使用 GUI 檢查 ONNX 模型 ¶

當轉換失敗時 ¶

API 參考 ¶

文件

教學課程

資源

基於 TorchDynamo 的 ONNX 匯出器¶

概觀¶

相依性¶

一個簡單的例子¶

使用 GUI 檢查 ONNX 模型¶

當轉換失敗時¶

API 參考¶

文件

教學課程

資源

概觀 ¶

相依性 ¶

一個簡單的例子 ¶

使用 GUI 檢查 ONNX 模型 ¶

當轉換失敗時 ¶

API 參考 ¶