捷徑

基於 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 有一組控制匯出器發出診斷資訊的選項。

相依性

ONNX 匯出器依賴於額外的 Python 套件

它們可以透過 pip 安裝

pip install --upgrade onnx onnxscript

一個簡單的範例

請參閱以下使用簡單的多層感知器 (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 檢視匯出的模型。

MLP model as viewed using Netron

請注意,每個圖層都以矩形框表示,右上角有一個 *f* 圖示。

ONNX function highlighted on MLP model

展開它就會顯示函數主體。

ONNX function body

函數主體是一系列 ONNX 運算子或其他函數。

使用 SARIF 診斷問題

ONNX 診斷透過採用 靜態分析結果交換格式 (又稱 SARIF) 來超越常規日誌,以協助使用者使用 GUI(例如 Visual Studio Code 的 SARIF 檢視器)對其模型進行偵錯和改進。

主要優點如下:

  • 診斷資訊以機器可解析的 靜態分析結果交換格式 (SARIF) 發出。

  • 一種全新的、更清晰、結構化的方式來添加新的診斷規則並追蹤現有規則。

  • 作為未來更多改進的基礎,這些改進將使用診斷資訊。

API 參考

torch.onnx.dynamo_export(model, /, *model_args, export_options=None, **model_kwargs)

將 torch.nn.Module 匯出至 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.}
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 模型輸入轉換而來的張量序列。

回傳類型

Sequence[Union[Tensor, int, float, bool, dtype]]

範例

# 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 模型輸出。

回傳類型

Sequence[Union[Tensor, int, float, bool]]

範例

# 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 fake_context: Optional[ONNXFakeContext]

與匯出相關聯的虛擬上下文。

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

變數
  • session_options (Optional[Sequence['onnxruntime.SessionOptions']]) – ONNX Runtime 工作階段選項。

  • execution_providers (Optional[Sequence[Union[str, Tuple[str, Dict[Any, Any]]]]]) – 模型執行期間要使用的 ONNX Runtime 執行提供者。

  • execution_provider_options (Optional[Sequence[Dict[Any, Any]]]) – ONNX Runtime 執行提供者選項。

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

列表按註冊時間排序。自定義運算符應位於列表的後半部分。

參數
  • namespace (str) – 要獲取的運算符的命名空間。

  • op_name (str) – 要獲取的運算符的名稱。

  • overload (Optional[str]) – 要獲取的運算符的重載。如果是默認重載,則將其保留為 None。

回傳值

與給定名稱對應的 ONNXFunction 列表,如果名稱不在註冊表中,則為 None。

回傳類型

Optional[List[ONNXFunction]]

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

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

參數
  • namespace (str) – 要檢查的運算符的命名空間。

  • op_name (str) – 要檢查的運算符的名稱。

  • overload (Optional[str]) – 要檢查的運算符的重載。如果是默認重載,則將其保留為 None。

回傳值

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

回傳類型

bool

property opset_version: int

匯出器應定位的 ONNX opset 版本。默認為最新支持的 ONNX opset 版本:18。默認版本會隨著 ONNX 的不斷發展而增加。

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

註冊一個自定義運算符:torch.ops.<namespace>.<op_name>.<overload>。

參數
  • function (Union[onnxscript.OnnxFunction, onnxscript.TracedOnnxFunction]) – 要註冊的 onnx-sctip 函數。

  • namespace (str) – 要註冊的運算符的命名空間。

  • op_name (str) – 要註冊的運算符的名稱。

  • overload (Optional[str]) – 要註冊的運算符的重載。如果是默認重載,則將其保留為 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,則將警告診斷視為錯誤診斷。

文件

訪問 PyTorch 的完整開發者文檔

查看文檔

教程

獲取針對初學者和高級開發者的深入教程

查看教程

資源

查找開發資源並獲得您的問題的答案

查看資源