TorchScript-based ONNX Exporter#
Created On: Aug 31, 2017 | Last Updated On: Jun 10, 2025
注意
To export an ONNX model using TorchDynamo instead of TorchScript, please see Learn more about the TorchDynamo-based ONNX Exporter
Example: AlexNet from PyTorch to ONNX#
Here is a simple script which exports a pretrained AlexNet to an ONNX file named alexnet.onnx. The call to torch.onnx.export runs the model once to trace its execution and then exports the traced model to the specified file
import torch
import torchvision
dummy_input = torch.randn(10, 3, 224, 224, device="cuda")
model = torchvision.models.alexnet(pretrained=True).cuda()
# Providing input and output names sets the display names for values
# within the model's graph. Setting these does not change the semantics
# of the graph; it is only for readability.
#
# The inputs to the network consist of the flat list of inputs (i.e.
# the values you would pass to the forward() method) followed by the
# flat list of parameters. You can partially specify names, i.e. provide
# a list here shorter than the number of inputs to the model, and we will
# only set that subset of names, starting from the beginning.
input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]
output_names = [ "output1" ]
torch.onnx.export(model, dummy_input, "alexnet.onnx", verbose=True, input_names=input_names, output_names=output_names)
The resulting alexnet.onnx file contains a binary protocol buffer which contains both the network structure and parameters of the model you exported (in this case, AlexNet). The argument verbose=True causes the exporter to print out a human-readable representation of the model
# These are the inputs and parameters to the network, which have taken on
# the names we specified earlier.
graph(%actual_input_1 : Float(10, 3, 224, 224)
%learned_0 : Float(64, 3, 11, 11)
%learned_1 : Float(64)
%learned_2 : Float(192, 64, 5, 5)
%learned_3 : Float(192)
# ---- omitted for brevity ----
%learned_14 : Float(1000, 4096)
%learned_15 : Float(1000)) {
# Every statement consists of some output tensors (and their types),
# the operator to be run (with its attributes, e.g., kernels, strides,
# etc.), its input tensors (%actual_input_1, %learned_0, %learned_1)
%17 : Float(10, 64, 55, 55) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[11, 11], pads=[2, 2, 2, 2], strides=[4, 4]](%actual_input_1, %learned_0, %learned_1), scope: AlexNet/Sequential[features]/Conv2d[0]
%18 : Float(10, 64, 55, 55) = onnx::Relu(%17), scope: AlexNet/Sequential[features]/ReLU[1]
%19 : Float(10, 64, 27, 27) = onnx::MaxPool[kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[2, 2]](%18), scope: AlexNet/Sequential[features]/MaxPool2d[2]
# ---- omitted for brevity ----
%29 : Float(10, 256, 6, 6) = onnx::MaxPool[kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[2, 2]](%28), scope: AlexNet/Sequential[features]/MaxPool2d[12]
# Dynamic means that the shape is not known. This may be because of a
# limitation of our implementation (which we would like to fix in a
# future release) or shapes which are truly dynamic.
%30 : Dynamic = onnx::Shape(%29), scope: AlexNet
%31 : Dynamic = onnx::Slice[axes=[0], ends=[1], starts=[0]](%30), scope: AlexNet
%32 : Long() = onnx::Squeeze[axes=[0]](%31), scope: AlexNet
%33 : Long() = onnx::Constant[value={9216}](), scope: AlexNet
# ---- omitted for brevity ----
%output1 : Float(10, 1000) = onnx::Gemm[alpha=1, beta=1, broadcast=1, transB=1](%45, %learned_14, %learned_15), scope: AlexNet/Sequential[classifier]/Linear[6]
return (%output1);
}
You can also verify the output using the ONNX library, which you can install using pip
pip install onnx
Then, you can run
import onnx
# Load the ONNX model
model = onnx.load("alexnet.onnx")
# Check that the model is well formed
onnx.checker.check_model(model)
# Print a human readable representation of the graph
print(onnx.helper.printable_graph(model.graph))
You can also run the exported model with one of the many runtimes that support ONNX. For example after installing ONNX Runtime, you can load and run the model
import onnxruntime as ort
import numpy as np
ort_session = ort.InferenceSession("alexnet.onnx")
outputs = ort_session.run(
None,
{"actual_input_1": np.random.randn(10, 3, 224, 224).astype(np.float32)},
)
print(outputs[0])
Here is a more involved tutorial on exporting a model and running it with ONNX Runtime.
Tracing vs Scripting#
Internally, torch.onnx.export() requires a torch.jit.ScriptModule rather than a torch.nn.Module. If the passed-in model is not already a ScriptModule, export() will use tracing to convert it to one
Tracing: If
torch.onnx.export()is called with a Module that is not already aScriptModule, it first does the equivalent oftorch.jit.trace(), which executes the model once with the givenargsand records all operations that happen during that execution. This means that if your model is dynamic, e.g., changes behavior depending on input data, the exported model will not capture this dynamic behavior. We recommend examining the exported model and making sure the operators look reasonable. Tracing will unroll loops and if statements, exporting a static graph that is exactly the same as the traced run. If you want to export your model with dynamic control flow, you will need to use scripting.Scripting: Compiling a model via scripting preserves dynamic control flow and is valid for inputs of different sizes. To use scripting
Use
torch.jit.script()to produce aScriptModule.Call
torch.onnx.export()with theScriptModuleas the model. Theargsare still required, but they will be used internally only to produce example outputs, so that the types and shapes of the outputs can be captured. No tracing will be performed.
See Introduction to TorchScript and TorchScript for more details, including how to compose tracing and scripting to suit the particular requirements of different models.
Avoiding Pitfalls#
Avoid NumPy and built-in Python types#
PyTorch models can be written using NumPy or Python types and functions, but during tracing, any variables of NumPy or Python types (rather than torch.Tensor) are converted to constants, which will produce the wrong result if those values should change depending on the inputs.
For example, rather than using numpy functions on numpy.ndarrays
# Bad! Will be replaced with constants during tracing.
x, y = np.random.rand(1, 2), np.random.rand(1, 2)
np.concatenate((x, y), axis=1)
Use torch operators on torch.Tensors
# Good! Tensor operations will be captured during tracing.
x, y = torch.randn(1, 2), torch.randn(1, 2)
torch.cat((x, y), dim=1)
And rather than use torch.Tensor.item() (which converts a Tensor to a Python built-in number)
# Bad! y.item() will be replaced with a constant during tracing.
def forward(self, x, y):
return x.reshape(y.item(), -1)
Use torch’s support for implicit casting of single-element tensors
# Good! y will be preserved as a variable during tracing.
def forward(self, x, y):
return x.reshape(y, -1)
Avoid Tensor.data#
Using the Tensor.data field can produce an incorrect trace and therefore an incorrect ONNX graph. Use torch.Tensor.detach() instead. (Work is ongoing to remove Tensor.data entirely).
Avoid in-place operations when using tensor.shape in tracing mode#
In tracing mode, shapes obtained from tensor.shape are traced as tensors, and share the same memory. This might cause a mismatch the final output values. As a workaround, avoid the use of inplace operations in these scenarios. For example, in the model
class Model(torch.nn.Module):
def forward(self, states):
batch_size, seq_length = states.shape[:2]
real_seq_length = seq_length
real_seq_length += 2
return real_seq_length + seq_length
real_seq_length and seq_length share the same memory in tracing mode. This could be avoided by rewriting the inplace operation
real_seq_length = real_seq_length + 2
Limitations#
Types#
Only
torch.Tensors, numeric types that can be trivially converted to torch.Tensors (e.g. float, int), and tuples and lists of those types are supported as model inputs or outputs. Dict and str inputs and outputs are accepted in tracing mode, butAny computation that depends on the value of a dict or a str input will be replaced with the constant value seen during the one traced execution.
任何輸出為 dict 的內容都將被靜默替換為其值的展平序列（鍵將被移除）。例如，
{"foo": 1, "bar": 2}將會變成(1, 2)。任何輸出為 str 的內容都將被靜默移除。
在指令碼模式下，由於 ONNX 對巢狀序列的支援有限，因此不支援涉及元組和列表的某些操作。特別是，將元組追加到列表是不支援的。在跟蹤模式下，巢狀序列將在跟蹤期間自動展平。
運算子實現差異#
由於運算子實現上的差異，在不同執行時上執行匯出的模型可能會產生彼此不同或與 PyTorch 不同的結果。通常這些差異在數值上很小，所以這隻應在您的應用程式對這些微小差異敏感時才需要關注。
不支援的張量索引模式#
下面列出了無法匯出的張量索引模式。如果您在匯出模型時遇到問題，但該模型不包含以下任何不受支援的模式，請仔細檢查您是否使用了最新的 opset_version。
讀取 / 獲取#
在索引張量進行讀取時，不支援以下模式
# Tensor indices that includes negative values.
data[torch.tensor([[1, 2], [2, -3]]), torch.tensor([-2, 3])]
# Workarounds: use positive index values.
寫入 / 設定#
在索引張量進行寫入時，不支援以下模式
# Multiple tensor indices if any has rank >= 2
data[torch.tensor([[1, 2], [2, 3]]), torch.tensor([2, 3])] = new_data
# Workarounds: use single tensor index with rank >= 2,
# or multiple consecutive tensor indices with rank == 1.
# Multiple tensor indices that are not consecutive
data[torch.tensor([2, 3]), :, torch.tensor([1, 2])] = new_data
# Workarounds: transpose `data` such that tensor indices are consecutive.
# Tensor indices that includes negative values.
data[torch.tensor([1, -2]), torch.tensor([-2, 3])] = new_data
# Workarounds: use positive index values.
# Implicit broadcasting required for new_data.
data[torch.tensor([[0, 2], [1, 1]]), 1:3] = new_data
# Workarounds: expand new_data explicitly.
# Example:
# data shape: [3, 4, 5]
# new_data shape: [5]
# expected new_data shape after broadcasting: [2, 2, 2, 5]
新增對運算子的支援#
匯出包含不受支援的運算子的模型時，您會看到類似以下的錯誤訊息
RuntimeError: ONNX export failed: Couldn't export operator foo
發生這種情況時，您可以採取一些措施：
更改模型以不使用該運算子。
建立一個符號函式來轉換該運算子，並將其註冊為自定義符號函式。
為 PyTorch 貢獻程式碼，將相同的符號函式新增到
torch.onnx本身。
如果您決定實現一個符號函式（我們希望您能將其貢獻回 PyTorch！），以下是如何開始：
ONNX 匯出器內部機制#
“符號函式”是一個將 PyTorch 運算子分解為一系列 ONNX 運算子的組合的函式。
匯出期間，匯出器會按照拓撲順序訪問 TorchScript 圖中的每個節點（包含一個 PyTorch 運算子）。訪問節點時，匯出器會查詢該運算子已註冊的符號函式。符號函式是用 Python 實現的。名為 foo 的運算子的符號函式大致看起來像這樣：
def foo(
g,
input_0: torch._C.Value,
input_1: torch._C.Value) -> Union[None, torch._C.Value, List[torch._C.Value]]:
"""
Adds the ONNX operations representing this PyTorch function by updating the
graph g with `g.op()` calls.
Args:
g (Graph): graph to write the ONNX representation into.
input_0 (Value): value representing the variables which contain
the first input for this operator.
input_1 (Value): value representing the variables which contain
the second input for this operator.
Returns:
A Value or List of Values specifying the ONNX nodes that compute something
equivalent to the original PyTorch operator with the given inputs.
None if it cannot be converted to ONNX.
"""
...
torch._C 型別是 ir.h 中定義的 C++ 型別之上的 Python 包裝器。
新增符號函式的過程取決於運算子的型別。
ATen 運算子#
ATen 是 PyTorch 的內建張量庫。如果運算子是 ATen 運算子（在 TorchScript 圖中顯示為帶有 aten:: 字首），請確保它尚未受支援。
支援的運算子列表#
訪問自動生成的 支援的 TorchScript 運算子列表，瞭解每個 opset_version 支援的運算子的詳細資訊。
為 ATen 或量化運算子新增支援#
如果運算子不在上面的列表中
在
torch/onnx/symbolic_opset<version>.py中定義符號函式，例如 torch/onnx/symbolic_opset9.py。確保函式名與 ATen 函式名相同，ATen 函式名可能宣告在torch/_C/_VariableFunctions.pyi或torch/nn/functional.pyi中（這些檔案在構建時生成，因此在您簽出程式碼時不會出現，直到您構建 PyTorch）。預設情況下，第一個引數是 ONNX 圖。其他引數名必須與
.pyi檔案中的名稱完全匹配，因為分派是透過關鍵字引數完成的。在符號函式中，如果運算子屬於 ONNX 標準運算子集，我們只需要建立一個節點來表示圖中的 ONNX 運算子。如果不是，我們可以組合幾個具有與 ATen 運算子等效語義的標準運算子。
以下是處理 ELU 運算子缺失符號函式的示例。
如果我們執行以下程式碼：
print(
torch.jit.trace(
torch.nn.ELU(), # module
torch.ones(1) # example input
).graph
)
我們會看到類似這樣的內容：
graph(%self : __torch__.torch.nn.modules.activation.___torch_mangle_0.ELU,
%input : Float(1, strides=[1], requires_grad=0, device=cpu)):
%4 : float = prim::Constant[value=1.]()
%5 : int = prim::Constant[value=1]()
%6 : int = prim::Constant[value=1]()
%7 : Float(1, strides=[1], requires_grad=0, device=cpu) = aten::elu(%input, %4, %5, %6)
return (%7)
由於我們在圖中看到了 aten::elu，我們知道這是一個 ATen 運算子。
我們檢查 ONNX 運算子列表，並確認 Elu 在 ONNX 中已標準化。
我們在 torch/nn/functional.pyi 中找到 elu 的簽名：
def elu(input: Tensor, alpha: float = ..., inplace: bool = ...) -> Tensor: ...
我們將以下行新增到 symbolic_opset9.py：
def elu(g, input: torch.Value, alpha: torch.Value, inplace: bool = False):
return g.op("Elu", input, alpha_f=alpha)
現在 PyTorch 能夠匯出包含 aten::elu 運算子的模型了！
請參閱 torch/onnx/symbolic_opset*.py 檔案以獲取更多示例。
torch.autograd.Functions#
如果運算子是 torch.autograd.Function 的子類，有三種方法可以匯出它。
靜態符號方法#
您可以在函式類中新增一個名為 symbolic 的靜態方法。它應該返回表示函式在 ONNX 中行為的 ONNX 運算子。例如：
class MyRelu(torch.autograd.Function):
@staticmethod
def forward(ctx, input: torch.Tensor) -> torch.Tensor:
ctx.save_for_backward(input)
return input.clamp(min=0)
@staticmethod
def symbolic(g: torch.Graph, input: torch.Value) -> torch.Value:
return g.op("Clip", input, g.op("Constant", value_t=torch.tensor(0, dtype=torch.float)))
內聯 Autograd 函式#
在未為其後續的 torch.autograd.Function 提供靜態符號方法，或者沒有提供將 prim::PythonOp 註冊為自定義符號函式的函式的情況下，torch.onnx.export() 會嘗試內聯與該 torch.autograd.Function 對應的圖，以便該函式被分解為函式內使用的各個運算子。只要這些單個運算子受支援，匯出就應該成功。例如：
class MyLogExp(torch.autograd.Function):
@staticmethod
def forward(ctx, input: torch.Tensor) -> torch.Tensor:
ctx.save_for_backward(input)
h = input.exp()
return h.log().log()
此模型沒有提供靜態符號方法，但它被匯出如下：
graph(%input : Float(1, strides=[1], requires_grad=0, device=cpu)):
%1 : float = onnx::Exp[](%input)
%2 : float = onnx::Log[](%1)
%3 : float = onnx::Log[](%2)
return (%3)
如果您需要避免內聯 torch.autograd.Function，您應該設定 operator_export_type 為 ONNX_FALLTHROUGH 或 ONNX_ATEN_FALLBACK 來匯出模型。
自定義運算子#
您可以匯出包含多個標準 ONNX 運算子組合的模型，或者由自定義 C++ 後端驅動的模型。
ONNX-script 函式#
如果某個運算子不是標準的 ONNX 運算子，但可以由多個現有 ONNX 運算子組成，您可以使用 ONNX-script 來建立一個外部 ONNX 函式來支援該運算子。您可以按照此示例匯出它：
import onnxscript
# There are three opset version needed to be aligned
# This is (1) the opset version in ONNX function
from onnxscript.onnx_opset import opset15 as op
opset_version = 15
x = torch.randn(1, 2, 3, 4, requires_grad=True)
model = torch.nn.SELU()
custom_opset = onnxscript.values.Opset(domain="onnx-script", version=1)
@onnxscript.script(custom_opset)
def Selu(X):
alpha = 1.67326 # auto wrapped as Constants
gamma = 1.0507
alphaX = op.CastLike(alpha, X)
gammaX = op.CastLike(gamma, X)
neg = gammaX * (alphaX * op.Exp(X) - alphaX)
pos = gammaX * X
zero = op.CastLike(0, X)
return op.Where(X <= zero, neg, pos)
# setType API provides shape/type to ONNX shape/type inference
def custom_selu(g: jit_utils.GraphContext, X):
return g.onnxscript_op(Selu, X).setType(X.type())
# Register custom symbolic function
# There are three opset version needed to be aligned
# This is (2) the opset version in registry
torch.onnx.register_custom_op_symbolic(
symbolic_name="aten::selu",
symbolic_fn=custom_selu,
opset_version=opset_version,
)
# There are three opset version needed to be aligned
# This is (2) the opset version in exporter
torch.onnx.export(
model,
x,
"model.onnx",
opset_version=opset_version,
# only needed if you want to specify an opset version > 1.
custom_opsets={"onnx-script": 2}
)
上面的示例將其匯出一個名為“onnx-script”的 opset 中的自定義運算子。匯出自定義運算子時，您可以透過匯出的 custom_opsets 字典指定自定義 opset 版本。如果未指定，自定義 opset 版本預設為 1。
注意：請務必匹配上述示例中提到的 opset 版本，並確保它們在匯出器步驟中被使用。如何編寫 onnx-script 函式的示例是其活躍開發中的一個 beta 版本。請遵循最新的 ONNX-script。
C++ 運算子#
如果模型使用 C++ 實現的自定義運算子，如 使用自定義 C++ 運算子擴充套件 TorchScript 中所述，您可以按照此示例匯出它：
from torch.onnx import symbolic_helper
# Define custom symbolic function
@symbolic_helper.parse_args("v", "v", "f", "i")
def symbolic_foo_forward(g, input1, input2, attr1, attr2):
return g.op("custom_domain::Foo", input1, input2, attr1_f=attr1, attr2_i=attr2)
# Register custom symbolic function
torch.onnx.register_custom_op_symbolic("custom_ops::foo_forward", symbolic_foo_forward, 9)
class FooModel(torch.nn.Module):
def __init__(self, attr1, attr2):
super().__init__()
self.attr1 = attr1
self.attr2 = attr2
def forward(self, input1, input2):
# Calling custom op
return torch.ops.custom_ops.foo_forward(input1, input2, self.attr1, self.attr2)
model = FooModel(attr1, attr2)
torch.onnx.export(
model,
(example_input1, example_input1),
"model.onnx",
# only needed if you want to specify an opset version > 1.
custom_opsets={"custom_domain": 2}
)
上面的示例將其匯出一個名為“custom_domain”的 opset 中的自定義運算子。匯出自定義運算子時，您可以透過匯出的 custom_opsets 字典指定自定義域版本。如果未指定，自定義 opset 版本預設為 1。
使用該模型的執行時需要支援自定義 op。請參閱 Caffe2 自定義 ops、ONNX Runtime 自定義 ops 或您選擇的執行時的文件。
一次性發現所有不可轉換的 ATen 運算子#
當由於不可轉換的 ATen 運算子導致匯出失敗時，可能不止一個這樣的運算子，但錯誤訊息只提到第一個。要一次性發現所有不可轉換的運算子，您可以：
# prepare model, args, opset_version
...
torch_script_graph, unconvertible_ops = torch.onnx.utils.unconvertible_ops(
model, args, opset_version=opset_version
)
print(set(unconvertible_ops))
該集合是近似的，因為一些運算子在轉換過程中可能會被移除，無需轉換。其他一些運算子可能支援部分功能，在特定輸入下可能導致轉換失敗，但這應該能讓您對不支援的運算子有一個大致的瞭解。請隨時在 GitHub 上提出 op 支援請求。
常見問題#
問：我匯出了我的 LSTM 模型，但它的輸入大小似乎是固定的？
跟蹤器會記錄示例輸入的形狀。如果模型應該接受動態形狀的輸入，請在呼叫
torch.onnx.export()時設定dynamic_axes。
問：如何匯出包含迴圈的模型？
請參閱 跟蹤 vs. 指令碼。
問：如何匯出帶有原始型別輸入（例如 int、float）的模型？
PyTorch 1.9 中添加了對原始數值型別輸入的支援。但是，匯出器不支援帶有 str 輸入的模型。
問：ONNX 是否支援隱式標量資料型別轉換？
ONNX 標準不支援，但匯出器會嘗試處理這部分。標量被匯出為常量張量。匯出器會找出標量的正確資料型別。在極少數情況下，當它無法做到時，您需要手動指定資料型別，例如使用 dtype=torch.float32。如果您看到任何錯誤，請[建立一個 GitHub issue](pytorch/pytorch#issues)。
問：張量列表是否可以匯出到 ONNX？
是的，對於
opset_version>= 11，因為 ONNX 在 opset 11 中引入了 Sequence 型別。
Python API#
函式#
- torch.onnx.export(model, args=(), f=None, *, kwargs=None, export_params=True, verbose=None, input_names=None, output_names=None, opset_version=None, dynamic_axes=None, keep_initializers_as_inputs=False, dynamo=False, 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, 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。使用
draft_export，它會移除資料相關操作中的一些健全性保證，以允許匯出繼續進行。如果匯出器遇到任何不健全的資料相關操作，您將收到警告。使用
torch.jit.trace()來跟蹤模型，然後轉換為 ExportedProgram。這是最不健全的策略，但對於將 TorchScript 模型轉換為 ONNX 可能有用。
- 引數
model (torch.nn.Module | torch.export.ExportedProgram | torch.jit.ScriptModule | torch.jit.ScriptFunction) – 要匯出的模型。
args (tuple[Any, ...]) – 示例位置輸入。任何非張量引數都將被硬編碼到匯出的模型中；任何張量引數都將成為匯出模型的輸入，順序與元組中出現的順序一致。
f (str | os.PathLike | None) – 輸出 ONNX 模型檔案的路徑。例如 “model.onnx”。
export_params (bool) – 如果為 false，則不匯出引數（權重）。
verbose (bool | None) – 是否啟用詳細日誌記錄。
input_names (Sequence[str] | None) – 分配給圖輸入節點的名稱，按順序。
output_names (Sequence[str] | None) – 分配給圖輸出節點的名稱，按順序。
opset_version (int | None) – 要針對的 預設（ai.onnx）opset 的版本。必須 >= 7。
dynamic_axes (Mapping[str, Mapping[int, str]] | Mapping[str, Sequence[int]] | None) –
預設情況下，匯出的模型的所有輸入和輸出張量的形狀都將精確匹配
args中給定的形狀。要將張量的軸指定為動態（即僅在執行時知道），請將dynamic_axes設定為具有以下架構的字典：- KEY (str): 輸入或輸出名稱。每個名稱也必須在
input_names或 output_names.
- KEY (str): 輸入或輸出名稱。每個名稱也必須在
- 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 ...
keep_initializers_as_inputs (bool) –
如果為 True，則匯出圖中的所有初始化器（通常對應於模型權重）也將作為輸入新增到圖中。如果為 False，則初始化器不作為輸入新增到圖中，只新增使用者輸入作為輸入。
如果您打算在執行時提供模型權重，請將其設定為 True。如果您打算將權重設定為靜態以允許後端/執行時進行更好的最佳化（例如常量摺疊），請將其設定為 False。
dynamo (bool) – 是否使用
torch.exportExportedProgram 而不是 TorchScript 來匯出模型。external_data (bool) – 是否將模型權重儲存為外部資料檔案。這對於權重過大（超過 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。
training (_C_onnx.TrainingMode) – 已棄用的選項。而是應在匯出模型之前設定模型的訓練模式。
operator_export_type (_C_onnx.OperatorExportTypes) – 已棄用的選項。僅支援 ONNX。
do_constant_folding (bool) – 已棄用的選項。
custom_opsets (Mapping[str, int] | None) –
已棄用。一個字典：
KEY (str): opset 域名稱。
VALUE (int): opset 版本。
如果
model引用了自定義 opset 但未在此字典中提及，則 opset 版本將設定為 1。只能透過此引數指示自定義 opset 域名稱和版本。export_modules_as_functions (bool | Collection[type[torch.nn.Module]]) –
已棄用的選項。
啟用將所有
nn.Module前向呼叫匯出為 ONNX 中的本地函式的標誌。或者一個集合，用於指定要匯出為 ONNX 中的本地函式的特定型別模組。此功能需要opset_version>= 15，否則匯出將失敗。這是因為opset_version< 15 意味著 IR 版本 < 8，這意味著不支援本地函式。模組變數將匯出為函式屬性。函式屬性有兩種類別：1. 帶有註解的屬性：透過 PEP 526 風格帶有型別註解的類變數將匯出為屬性。帶有註解的屬性不在 ONNX 本地函式的子圖中被使用，因為它們不是由 PyTorch JIT 跟蹤建立的，但消費者可以使用它們來確定是否用特定的融合核心替換該函式。
2. 推斷的屬性：在模組的子圖中使用到的變數。屬性名稱將帶有字首“inferred::”。這是為了與從 Python 模組註解中檢索到的預定義屬性區分開來。推斷的屬性在 ONNX 本地函式的子圖中使用。
False(預設): 將nn.Module前向呼叫匯出為細粒度節點。True: 將所有nn.Module前向呼叫匯出為本地函式節點。- Set of type of nn.Module：匯出
nn.Module的前向呼叫作為區域性函式節點， 僅當
nn.Module的型別在集合中找到時。
- Set of type of nn.Module：匯出
autograd_inlining (bool) – 已棄用。用於控制是否內聯 autograd 函式的標誌。有關更多詳細資訊，請參閱 pytorch/pytorch#74765。
- 返回
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。
- torch.onnx.register_custom_op_symbolic(symbolic_name, symbolic_fn, opset_version)[source]#
為自定義運算子註冊符號函式。
當用戶為自定義/contrib 運算子註冊符號時，強烈建議透過 setType API 為該運算子新增形狀推斷，否則在某些極端情況下匯出的圖可能具有錯誤的形狀推斷。setType 的一個示例是 test_operators.py 中的 test_aten_embedding_2。
有關示例用法，請參閱模組文件中的“自定義運算子”。
- torch.onnx.unregister_custom_op_symbolic(symbolic_name, opset_version)[source]#
登出
symbolic_name。有關示例用法，請參閱模組文件中的“自定義運算子”。