torch.onnx#

創建於: 2025年6月10日 | 最後更新於: 2025年8月29日

概述#

Open Neural Network eXchange (ONNX) 是一種表示機器學習模型的開放標準格式。torch.onnx 模組捕獲原生 PyTorch torch.nn.Module 模型中的計算圖，並將其轉換為 ONNX 圖。

匯出的模型可以被許多支援 ONNX 的執行時使用，包括 Microsoft 的 ONNX Runtime。

下一個示例展示瞭如何匯出簡單模型。

import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 128, 5)

    def forward(self, x):
        return torch.relu(self.conv1(x))

input_tensor = torch.rand((1, 1, 128, 128), dtype=torch.float32)

model = MyModel()

torch.onnx.export(
    model,                  # model to export
    (input_tensor,),        # inputs of the model,
    "my_model.onnx",        # filename of the ONNX model
    input_names=["input"],  # Rename inputs for the ONNX model
    dynamo=True             # True or False to select the exporter to use
)

基於 torch.export 的 ONNX 匯出器#

基於 torch.export 的 ONNX 匯出器是 PyTorch 2.6 及更新版本的最新匯出器

使用 torch.export 引擎以“提前編譯”（AOT）方式生成一個僅表示函式張量計算的追蹤圖。生成的追蹤圖（1）在功能性的 ATen 運算元集中產生規範化的運算元（以及任何使用者指定的自定義運算元），（2）消除了所有 Python 控制流和資料結構（某些例外情況除外），並且（3）記錄了顯示這種規範化和控制流消除對未來輸入是健全的形狀約束集，然後才最終將其翻譯成 ONNX 圖。

瞭解更多關於基於 torch.export 的 ONNX 匯出器

常見問題#

問：我已經匯出了我的 LLM 模型，但它的輸入大小似乎是固定的？

追蹤器會記錄示例輸入的形狀。如果模型應該接受動態形狀的輸入，請在呼叫 torch.onnx.export() 時設定 dynamic_shapes。

問：如何匯出包含迴圈的模型？

請參閱 torch.cond。

貢獻/開發#

ONNX 匯出器是一個社群專案，我們歡迎貢獻。我們遵循 PyTorch 的貢獻指南，但您可能也對閱讀我們的開發 Wiki 感興趣。

torch.onnx API#

函式#

torch.onnx.export(model, args=(), f=None, *, kwargs=None, verbose=None, input_names=None, output_names=None, opset_version=None, dynamo=True, external_data=True, dynamic_shapes=None, custom_translation_table=None, report=False, optimize=True, verify=False, profile=False, dump_exported_program=False, artifacts_dir='.', fallback=False, export_params=True, keep_initializers_as_inputs=False, dynamic_axes=None, training=<TrainingMode.EVAL: 0>, operator_export_type=<OperatorExportTypes.ONNX: 0>, do_constant_folding=True, custom_opsets=None, export_modules_as_functions=False, autograd_inlining=True)[source]

將模型匯出為 ONNX 格式。

設定 dynamo=True 將啟用新的 ONNX 匯出邏輯，該邏輯基於 torch.export.ExportedProgram 和更現代的翻譯邏輯集。這是將模型匯出為 ONNX 的推薦方法，也是預設方法。

當 dynamo=True 時

匯出器嘗試以下策略來獲取用於轉換為 ONNX 的 ExportedProgram。

如果模型已經是 ExportedProgram，則按原樣使用。
使用 torch.export.export() 並設定 strict=False。
使用 torch.export.export() 並設定 strict=True。

引數

model (torch.nn.Module | torch.export.ExportedProgram | torch.jit.ScriptModule | torch.jit.ScriptFunction) – 要匯出的模型。
args (tuple[Any, ...]) – 示例位置輸入。任何非張量引數都將被硬編碼到匯出的模型中；任何張量引數都將成為匯出模型的輸入，順序與元組中的順序相同。
f (str | os.PathLike | None) – 輸出 ONNX 模型檔案的路徑。例如：“model.onnx”。此引數保留用於向後相容。建議將其留空（None），並使用返回的 torch.onnx.ONNXProgram 來代替將模型序列化到檔案。
kwargs (dict[str, Any] | None) – 可選的示例關鍵字輸入。
verbose (bool | None) – 是否啟用詳細日誌記錄。
input_names (Sequence[str] | None) – 要分配給圖中輸入節點的名稱，按順序排列。
output_names (Sequence[str] | None) – 要分配給圖中輸出節點的名稱，按順序排列。
opset_version (int | None) – 要定位的預設（ai.onnx）運算元集的版本。您應該根據您想執行匯出模型的執行時後端或編譯器的支援的運算元集版本設定 opset_version。留空（預設值，None）以使用推薦的版本，或參考 ONNX 運算元文件以獲取更多資訊。
dynamo (bool) – 是否使用 torch.export ExportedProgram 而不是 TorchScript 來匯出模型。
external_data (bool) – 是否將模型權重儲存為外部資料檔案。這對於權重超過 ONNX 檔案大小限制（2GB）的模型是必需的。當設定為 False 時，權重將與模型架構一起儲存在 ONNX 檔案中。
dynamic_shapes (dict[str, Any] | tuple[Any, ...] | list[Any] | None) – 模型輸入的動態形狀的字典或元組。有關更多詳細資訊，請參閱 torch.export.export()。當 dynamo 為 True 時，僅使用（且優先）此引數。請注意，dynamic_shapes 設計用於 dynamo=True 匯出模型時使用，而 dynamic_axes 用於 dynamo=False。
custom_translation_table (dict[Callable, Callable | Sequence[Callable]] | None) – 模型中運算元的自定義分解的字典。字典的鍵應該是 fx Node 中的可呼叫目標（例如 torch.ops.aten.stft.default），值應該是使用 ONNX Script 構建該圖的函式。此選項僅在 dynamo 為 True 時有效。
report (bool) – 是否為匯出過程生成 Markdown 報告。此選項僅在 dynamo 為 True 時有效。
optimize (bool) – 是否最佳化匯出的模型。此選項僅在 dynamo 為 True 時有效。預設為 True。
verify (bool) – 是否使用 ONNX Runtime 驗證匯出的模型。此選項僅在 dynamo 為 True 時有效。
profile (bool) – 是否分析匯出過程。此選項僅在 dynamo 為 True 時有效。
dump_exported_program (bool) – 是否將 torch.export.ExportedProgram 轉儲到檔案。這對於除錯匯出器很有用。此選項僅在 dynamo 為 True 時有效。
artifacts_dir (str | os.PathLike) – 用於儲存除錯偽像（如報告和序列化的匯出程式）的目錄。此選項僅在 dynamo 為 True 時有效。
fallback (bool) – 如果 dynamo 匯出器失敗，是否回退到 TorchScript 匯出器。此選項僅在 dynamo 為 True 時有效。啟用 fallback 時，建議即使提供了 dynamic_shapes，也設定 dynamic_axes。
export_params (bool) –
當指定 ``f`` 時：如果為 False，則不會匯出引數（權重）。

您也可以不指定它，並使用返回的 torch.onnx.ONNXProgram 來控制序列化模型時如何處理初始值設定項。
keep_initializers_as_inputs (bool) –
當指定 ``f`` 時：如果為 True，則匯出的圖中的所有初始值設定項（通常對應於模型權重）也將作為圖的輸入新增。如果為 False，則初始值設定項不會被新增為圖的輸入，只有使用者輸入會被新增為輸入。

如果打算在執行時提供模型權重，請將其設定為 True。如果權重是靜態的，請將其設定為 False，以便後端/執行時進行更好的最佳化（例如，常量摺疊）。

您也可以不指定它，並使用返回的 torch.onnx.ONNXProgram 來控制序列化模型時如何處理初始值設定項。

dynamic_axes (Mapping[str, Mapping[int, str]] | Mapping[str, Sequence[int]] | None) –

當 dynamo=True 且 fallback 未啟用時，優先指定 dynamic_shapes。

預設情況下，匯出的模型的所有輸入和輸出張量的形狀都將精確匹配 args 中給出的形狀。要將張量的軸指定為動態（即僅在執行時才知道），請將 dynamic_axes 設定為具有以下架構的字典：

KEY (str)：輸入或輸出名稱。每個名稱還必須在 input_names 或
output_names 中提供。.
VALUE (dict 或 list)：如果是 dict，則鍵是軸索引，值是軸名稱。如果是
list，則每個元素都是一個軸索引。

例如

class SumModule(torch.nn.Module):
    def forward(self, x):
        return torch.sum(x, dim=1)


torch.onnx.export(
    SumModule(),
    (torch.ones(2, 2),),
    "onnx.pb",
    input_names=["x"],
    output_names=["sum"],
)

生成

input {
  name: "x"
  ...
      shape {
        dim {
          dim_value: 2  # axis 0
        }
        dim {
          dim_value: 2  # axis 1
...
output {
  name: "sum"
  ...
      shape {
        dim {
          dim_value: 2  # axis 0
...

雖然

torch.onnx.export(
    SumModule(),
    (torch.ones(2, 2),),
    "onnx.pb",
    input_names=["x"],
    output_names=["sum"],
    dynamic_axes={
        # dict value: manually named axes
        "x": {0: "my_custom_axis_name"},
        # list value: automatic names
        "sum": [0],
    },
)

生成

input {
  name: "x"
  ...
      shape {
        dim {
          dim_param: "my_custom_axis_name"  # axis 0
        }
        dim {
          dim_value: 2  # axis 1
...
output {
  name: "sum"
  ...
      shape {
        dim {
          dim_param: "sum_dynamic_axes_1"  # axis 0
...

training (_C_onnx.TrainingMode) – 已棄用的選項。而是，在匯出模型之前設定模型的訓練模式。
operator_export_type (_C_onnx.OperatorExportTypes) – 已棄用的選項。僅支援 ONNX。
do_constant_folding (bool) – 已棄用的選項。
custom_opsets (Mapping[str, int] | None) – 已棄用的選項。
export_modules_as_functions (bool | Collection[type[torch.nn.Module]]) – 已棄用的選項。
autograd_inlining (bool) – 已棄用的選項。

返回

torch.onnx.ONNXProgram 如果 dynamo 為 True，否則為 None。

返回型別

ONNXProgram | None

版本 2.6 中已更改： training 已棄用。而是，在匯出模型之前設定模型的訓練模式。operator_export_type 已棄用。僅支援 ONNX。do_constant_folding 已棄用。它始終啟用。export_modules_as_functions 已棄用。autograd_inlining 已棄用。

版本 2.7 中已更改： optimize 現在預設為 True。

版本 2.9 中已更改： dynamo 現在預設為 True。

torch.onnx.is_in_onnx_export()[source]

返回當前是否處於 ONNX 匯出過程中。

返回型別: 布林值

類#

class torch.onnx.ONNXProgram(model, exported_program)

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

變數

model – ONNX 模型，作為 ONNX IR 模型物件。
exported_program – 生成 ONNX 模型的匯出程式。

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

已棄用的 API#

版本 2.6 已棄用：這些函式已被棄用，將在未來版本中刪除。

torch.onnx.register_custom_op_symbolic(symbolic_name, symbolic_fn, opset_version)[source]#

為自定義運算子註冊符號函式。

當用戶為自定義/貢獻運算子註冊符號時，強烈建議透過 setType API 為該運算子新增形狀推斷，否則匯出的圖在某些極端情況下可能具有不正確的形狀推斷。setType 的一個示例是 test_operators.py 中的 test_aten_embedding_2。

有關示例用法，請參閱模組文件中的“自定義運算子”。

引數

symbolic_name (str) – 自定義運算子的名稱，格式為“<domain>::<op>”。
symbolic_fn (Callable) – 一個函式，它接收 ONNX 圖和當前運算子的輸入引數，並返回要新增到圖中的新運算子節點。
opset_version (int) – 要註冊的 ONNX 運算元集版本。

torch.onnx.unregister_custom_op_symbolic(symbolic_name, opset_version)[source]#

登出 symbolic_name。

有關示例用法，請參閱模組文件中的“自定義運算子”。

引數

symbolic_name (str) – 自定義運算子的名稱，格式為“<domain>::<op>”。
opset_version (int) – 要登出的 ONNX 運算元集版本。

torch.onnx.select_model_mode_for_export(model, mode)[source]#

一個上下文管理器，用於臨時將 model 的訓練模式設定為 mode，並在退出 with 塊時將其重置。

版本 2.7 已棄用：請在匯出模型之前設定訓練模式。

引數

model – 與 export() 的 model 引數型別和含義相同。
mode (TrainingMode) – 與 export() 的 training 引數型別和含義相同。

torch.onnx#

概述#

基於 torch.export 的 ONNX 匯出器#

常見問題#

貢獻/開發#

torch.onnx API#

函式#

類#

已棄用的 API#

文件

教程

資源