基於torch.export的ONNX匯出器#

建立日期：2025年6月10日 | 最後更新日期：2025年8月22日

概述
依賴項
一個簡單的示例
使用GUI檢查ONNX模型
轉換失敗時
元資料
API 參考

概述 #

利用torch.export引擎，以提前編譯（AOT）的方式生成一個僅包含函式張量計算的圖。生成的圖（1）在函式式ATen運算元集中產生規範化的運算元（以及任何使用者指定的自定義運算元），（2）消除了所有Python控制流和資料結構（某些例外情況），並且（3）記錄了一組形狀約束，以證明這種規範化和控制流消除對於未來的輸入是合理的，然後最終將其轉換為ONNX圖。

此外，在匯出過程中，記憶體使用量也顯著降低。

一個簡單的示例 #

下面是一個使用簡單多層感知器（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)
      self.fc_combined = nn.Linear(8 + 8 + 8, 8, bias=True)  # Combine all inputs

  def forward(self, tensor_x: torch.Tensor, input_dict: dict, input_list: list):
      """
      Forward method that requires all inputs:
      - tensor_x: A direct tensor input.
      - input_dict: A dictionary containing the tensor under the key 'tensor_x'.
      - input_list: A list where the first element is the tensor.
      """
      # Extract tensors from inputs
      dict_tensor = input_dict['tensor_x']
      list_tensor = input_list[0]

      # Combine all inputs into a single tensor
      combined_tensor = torch.cat([tensor_x, dict_tensor, list_tensor], dim=1)

      # Process the combined tensor through the layers
      combined_tensor = self.fc_combined(combined_tensor)
      combined_tensor = torch.sigmoid(combined_tensor)
      combined_tensor = self.fc0(combined_tensor)
      combined_tensor = torch.sigmoid(combined_tensor)
      combined_tensor = self.fc1(combined_tensor)
      combined_tensor = torch.sigmoid(combined_tensor)
      combined_tensor = self.fc2(combined_tensor)
      combined_tensor = torch.sigmoid(combined_tensor)
      output = self.fc3(combined_tensor)
      return output

model = MLPModel()

# Example inputs
tensor_input = torch.rand((97, 8), dtype=torch.float32)
dict_input = {'tensor_x': torch.rand((97, 8), dtype=torch.float32)}
list_input = [torch.rand((97, 8), dtype=torch.float32)]

# The input_names and output_names are used to identify the inputs and outputs of the ONNX model
input_names = ['tensor_input', 'tensor_x', 'list_input_index_0']
output_names = ['output']

# Exporting the model with all required inputs
onnx_program = torch.onnx.export(model,(tensor_input, dict_input, list_input), dynamic_shapes=({0: "batch_size"},{"tensor_x": {0: "batch_size"}},[{0: "batch_size"}]), input_names=input_names, output_names=output_names, dynamo=True,)

# Check the exported ONNX model is dynamic
assert onnx_program.model.graph.inputs[0].shape == ("batch_size", 8)
assert onnx_program.model.graph.inputs[1].shape == ("batch_size", 8)
assert onnx_program.model.graph.inputs[2].shape == ("batch_size", 8)

如上面的程式碼所示，您只需向torch.onnx.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")

轉換失敗時 #

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

元資料 #

在ONNX匯出過程中，每個ONNX節點都帶有元資料註解，這些註解有助於追溯其在原始PyTorch模型中的來源和上下文。這些元資料對於除錯、模型檢查和理解PyTorch與ONNX圖之間的對映非常有用。

每個ONNX節點都會新增以下元資料欄位：

namespace

一個表示節點分層名稱空間的字串，由模組/方法的堆疊跟蹤組成。

示例： __main__.SimpleAddModel/add: aten.add.Tensor
pkg.torch.onnx.class_hierarchy

一個表示通往該節點的模組層級的類名列表。

示例： ['__main__.SimpleAddModel', 'aten.add.Tensor']
pkg.torch.onnx.fx_node

原始FX節點的字串表示，包括其名稱、使用者數量、目標torch op、引數和關鍵字引數。

示例： %cat : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%tensor_x, %input_dict_tensor_x, %input_list_0], 1), kwargs = {})
pkg.torch.onnx.name_scopes

一個表示PyTorch模型中節點路徑的名稱作用域（方法）列表。

示例： ['', 'add']
pkg.torch.onnx.stack_trace

如果可用，這是建立該節點時原始程式碼的堆疊跟蹤。

示例
```
File "simpleadd.py", line 7, in forward
    return torch.add(x, y)
```

這些元資料欄位儲存在每個ONNX節點的metadata_props屬性中，可以使用Netron或透過程式設計方式進行檢查。

整個ONNX圖具有以下metadata_props：

pkg.torch.export.ExportedProgram.graph_signature

此屬性包含原始PyTorch ExportedProgram的graph_signature的字串表示。圖簽名描述了模型輸入和輸出的結構以及它們如何對映到ONNX圖。輸入被定義為InputSpec物件，這些物件包括輸入的種類（例如，引數的InputKind.PARAMETER，使用者定義的輸入的InputKind.USER_INPUT）、引數名稱、目標（可以是模型中的特定節點）以及輸入是否是持久的。輸出被定義為OutputSpec物件，這些物件指定輸出的種類（例如，OutputKind.USER_OUTPUT）和引數名稱。

要了解有關圖簽名的更多資訊，請參閱torch.export。
pkg.torch.export.ExportedProgram.range_constraints

此屬性包含原始PyTorch ExportedProgram中存在的任何範圍約束的字串表示。範圍約束指定了模型中符號形狀或值的有效範圍，這對於使用動態形狀或符號維度的模型可能很重要。

示例： s0: VR[2, int_oo]，這表示輸入張量的尺寸必須至少為2。

要了解有關範圍約束的更多資訊，請參閱torch.export。

ONNX圖中的每個輸入值可能具有以下元資料屬性：

pkg.torch.export.graph_signature.InputSpec.kind

輸入的種類，由PyTorch的InputKind列舉定義。

示例值：
- “USER_INPUT”: 使用者提供的模型輸入。
- “PARAMETER”: 模型引數（例如，權重）。
- “BUFFER”: 模型緩衝區（例如，BatchNorm中的執行均值）。
- “CONSTANT_TENSOR”: 常量張量引數。
- “CUSTOM_OBJ”: 自定義物件輸入。
- “TOKEN”: token輸入。
pkg.torch.export.graph_signature.InputSpec.persistent

指示輸入是否持久（即，是否應作為模型狀態的一部分儲存）。

示例值：
- “True”
- “False”

ONNX圖中的每個輸出值可能具有以下元資料屬性：

pkg.torch.export.graph_signature.OutputSpec.kind

輸入的種類，由PyTorch的OutputKind列舉定義。

示例值：
- “USER_OUTPUT”: 使用者可見的輸出。
- “LOSS_OUTPUT”: 損失值輸出。
- “BUFFER_MUTATION”: 表示緩衝區已發生變異。
- “GRADIENT_TO_PARAMETER”: 引數的梯度輸出。
- “GRADIENT_TO_USER_INPUT”: 使用者輸入的梯度輸出。
- “USER_INPUT_MUTATION”: 表示使用者輸入已發生變異。
- “TOKEN”: token輸出。

每個已初始化的值、輸入、輸出都具有以下元資料：

pkg.torch.onnx.original_node_name

在PyTorch FX圖中生成此值的節點的原始名稱，以防該值被重新命名。這有助於將初始化器追溯到其在原始模型中的來源。

示例： fc1.weight

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時有效。啟用回退時，建議即使提供了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。

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

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

變數

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

apply_weights(state_dict)[source]#

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

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

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

call_reference(*args, **kwargs)[source]#

使用參考後端執行ONNX模型。

返回型別: Sequence[Tensor]

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

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

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

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

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

初始化ONNX Runtime推理會話。

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

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

optimize()[source]#

最佳化ONNX模型。

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

release()[source]#

釋放推理會話。

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

save(destination, *, include_initializers=True, keep_initializers_as_inputs=False, external_data=None)[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 (Optional[bool]) – 是否將權重儲存為外部資料檔案。

引發

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

torch.onnx.is_in_onnx_export()[source]#

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

返回型別: 布林值

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

基於torch.export的ONNX匯出器#

概述 #

依賴項 #

一個簡單的示例 #

使用GUI檢查ONNX模型 #

轉換失敗時 #

元資料 #

API參考 #

文件

教程

資源

基於torch.export的ONNX匯出器#

概述#

依賴項#

一個簡單的示例#

使用GUI檢查ONNX模型#

轉換失敗時#

元資料#

API參考#

文件

教程

資源

概述 #

依賴項 #

一個簡單的示例 #

使用GUI檢查ONNX模型 #

轉換失敗時 #

元資料 #

API參考 #