基于 TorchScript 的 ONNX 导出器¶

注意

要使用 TorchDynamo 而非 TorchScript 导出 ONNX 模型，请参阅了解更多关于基于 TorchDynamo 的 ONNX 导出器

示例：从 PyTorch 到 ONNX 的 AlexNet ¶

这是一个简单的脚本，它将预训练的 AlexNet 导出到名为 alexnet.onnx 的 ONNX 文件。调用 torch.onnx.export 会运行模型一次以追踪其执行，然后将追踪到的模型导出到指定文件。

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)

生成的 alexnet.onnx 文件包含一个二进制 protocol buffer，其中包含您导出的模型（在本例中为 AlexNet）的网络结构和参数。参数 verbose=True 会使导出器打印出模型的可读表示。

# 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);
}

您还可以使用 ONNX 库来验证输出，您可以使用 pip 进行安装。

pip install onnx

然后，您可以运行

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

您还可以使用许多支持 ONNX 的运行时之一来运行导出的模型。例如，安装 ONNX Runtime 后，您可以加载并运行模型。

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])

这是一个更详细的教程，讲解如何导出模型并使用 ONNX Runtime 运行它。

追踪 (Tracing) 与脚本化 (Scripting)¶

在内部，torch.onnx.export() 需要一个 torch.jit.ScriptModule，而不是 torch.nn.Module。如果传入的模型还不是 ScriptModule，export() 将使用追踪将其转换为一个。

追踪：如果使用一个还不是 ScriptModule 的 Module 调用 torch.onnx.export()，它会首先执行等同于 torch.jit.trace() 的操作，该操作会使用给定的 args 并记录该执行期间发生的所有操作。这意味着如果您的模型是动态的，例如，行为取决于输入数据，导出的模型将不捕捉这种动态行为。我们建议检查导出的模型并确保算子看起来合理。追踪将展开循环和条件语句，导出一个与追踪运行完全相同的静态图。如果您想导出带有动态控制流的模型，则需要使用 脚本化。
脚本化：通过脚本化编译模型保留了动态控制流，并且对不同大小的输入有效。要使用脚本化，请执行以下操作：
- 使用 torch.jit.script() 来生成一个 ScriptModule。
- 使用 ScriptModule 作为模型来调用 torch.onnx.export()。 args 仍然是必需的，但它们只会在内部用于生成示例输出，以便捕获输出的类型和形状。不会执行追踪。

请参阅 TorchScript 入门和 TorchScript 以获取更多详细信息，包括如何组合追踪和脚本化以适应不同模型的特定需求。

避免陷阱 ¶

避免使用 NumPy 和内置 Python 类型 ¶

PyTorch 模型可以使用 NumPy 或 Python 类型和函数编写，但在追踪期间，任何 NumPy 或 Python 类型的变量（而不是 torch.Tensor）都会被转换为常量，如果这些值应根据输入而改变，这将产生错误的结果。

例如，与其在 numpy.ndarrays 上使用 numpy 函数

# 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)

不如在 torch.Tensors 上使用 torch 算子

# Good! Tensor operations will be captured during tracing.
x, y = torch.randn(1, 2), torch.randn(1, 2)
torch.cat((x, y), dim=1)

与其使用 torch.Tensor.item()（它将 Tensor 转换为 Python 内置数值）

# Bad! y.item() will be replaced with a constant during tracing.
def forward(self, x, y):
    return x.reshape(y.item(), -1)

不如使用 torch 对单元素张量隐式类型转换的支持

# Good! y will be preserved as a variable during tracing.
def forward(self, x, y):
    return x.reshape(y, -1)

避免使用 Tensor.data ¶

使用 Tensor.data 字段可能会产生不正确的追踪，从而生成不正确的 ONNX 图。请改用 torch.Tensor.detach()。（彻底移除 Tensor.data 的工作正在进行中）。

在追踪模式下使用 tensor.shape 时避免原地操作 ¶

在追踪模式下，从 tensor.shape 获取的形状被追踪为张量，并共享相同的内存。这可能会导致最终输出值不匹配。作为一种解决方法，请避免在这些场景中使用原地操作。例如，在模型中

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 和 seq_length 在追踪模式下共享相同的内存。这可以通过重写原地操作来避免

real_seq_length = real_seq_length + 2

限制 ¶

类型 ¶

只有 torch.Tensors、可以轻易转换为 torch.Tensors 的数值类型（例如 float, int）以及这些类型的元组和列表被支持作为模型输入或输出。在追踪模式下接受 Dict 和 str 输入和输出，但
- 任何依赖于 dict 或 str 输入值的计算 将被替换为 单次追踪执行期间看到 的常量值。
- 任何作为 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 类型是 C++ 中 ir.h 文件中定义的类型的 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)))

内联自动微分函数 ¶

在其后续 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”算子集中的自定义算子导出。导出自定义算子时，可以使用导出时的 custom_opsets 字典指定自定义域版本。如果未指定，自定义算子集版本默认为 1。

注意：请务必对齐上述示例中提到的 opset 版本，并确保它们在导出步骤中被使用。关于如何编写 onnx-script 函数的示例用法是 onnx-script 活跃开发中的一个 beta 版本。请遵循最新的 ONNX-script

C++ 算子 ¶

如果模型使用了在使用自定义 C++ 算子扩展 TorchScript 中描述的自定义 C++ 算子，您可以按照此示例导出它

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”算子集中的自定义算子导出。导出自定义算子时，可以使用导出时的 custom_opsets 字典指定自定义域版本。如果未指定，自定义算子集版本默认为 1。

使用该模型的运行时需要支持自定义算子。请参阅 Caffe2 自定义算子、ONNX Runtime 自定义算子，或您选择的运行时的文档。

一次性发现所有不可转换的 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 上提交 Issue。

常见问题 ¶

问：我导出了我的 LSTM 模型，但其输入大小似乎是固定的？

追踪器记录了示例输入的形状。如果模型应该接受动态形状的输入，请在调用 torch.onnx.export() 时设置 dynamic_axes。

问：如何导出包含循环的模型？

请参阅追踪 (Tracing) 与脚本化 (Scripting)。

问：如何导出包含原始类型输入（例如 int, float）的模型？

PyTorch 1.9 中添加了对原始数值类型输入的支持。但是，导出器不支持包含 str 输入的模型。

问：ONNX 是否支持隐式标量数据类型转换？

ONNX 标准本身不支持，但导出器会尝试处理这部分。标量将作为常量张量导出。导出器会为标量确定正确的数据类型。在少数无法确定数据类型的情况下，您需要手动指定，例如使用 dtype=torch.float32。如果看到任何错误，请创建 GitHub Issue。

问：Tensor 列表可以导出为 ONNX 吗？

是的，对于 opset_version >= 11，因为 ONNX 在 opset 11 中引入了 Sequence 类型。

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][source]¶

将模型导出为 ONNX 格式。

参数

model (torch.nn.Module | torch.export.ExportedProgram | torch.jit.ScriptModule | torch.jit.ScriptFunction) – 要导出的模型。
args (tuple[Any, ...]) – 示例位置输入。任何非 Tensor 参数都将硬编码到导出的模型中；任何 Tensor 参数都将成为导出模型的输入，按照它们在元组中出现的顺序排列。
f (str | os.PathLike | None) – 输出 ONNX 模型文件的路径。例如，“model.onnx”。
kwargs (dict[str, Any] | None) – 可选的示例关键字输入。
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 设置为具有以下模式的字典：

键 (str)：输入或输出名称。每个名称也必须在 input_names 或
output_names 中提供。.
值 (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.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 节点中的可调用目标作为键（例如 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) –
已弃用。一个字典
- 键 (str)：opset 域名称
- 值 (int)：opset 版本
如果 model 引用了自定义 opset 但未在此字典中提及，则 opset 版本将设置为 1。仅应通过此参数指定自定义 opset 域名称和版本。
export_modules_as_functions (bool | Collection[type[torch.nn.Module]]) –
已弃用选项。

标志，用于启用将所有 nn.Module 的 forward 调用导出为 ONNX 中的局部函数。或者是一个集合，用于指示要导出为 ONNX 中的局部函数的特定模块类型。此功能需要 opset_version >= 15，否则导出将失败。这是因为 opset_version < 15 表示 IR 版本 < 8，这意味着不支持局部函数。模块变量将作为函数属性导出。函数属性有两种类别。

1. 注释属性：通过PEP 526 风格进行类型注释的类变量将作为属性导出。注释属性不用于 ONNX 局部函数的子图中，因为它们不是由 PyTorch JIT tracing 创建的，但消费者可以使用它们来确定是否用特定的融合核替换该函数。

2. 推断属性：模块内算子使用的变量。属性名称将带有前缀 “inferred::”。这与从 python 模块注释中检索的预定义属性区分开来。推断属性用于 ONNX 局部函数的子图中。
- False (默认)：将 nn.Module 的 forward 调用导出为细粒度节点。
- True：将所有 nn.Module 的 forward 调用导出为局部函数节点。
- nn.Module 类型集合：将 nn.Module 的 forward 调用导出为局部函数节点，
  仅当 nn.Module 的类型在此集合中时。
autograd_inlining (bool) – 已弃用。标志，用于控制是否内联 autograd 函数。有关更多详细信息，请参阅 https://github.com/pytorch/pytorch/pull/74765。

返回值

如果 dynamo 为 True，则返回 torch.onnx.ONNXProgram，否则返回 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][source]¶

为自定义算子注册一个符号函数。

当用户为 custom/contrib 算子注册符号时，强烈建议通过 setType API 为该算子添加形状推断，否则在某些极端情况下导出的图可能具有不正确的形状推断。setType 的一个示例是 test_operators.py 中的 test_aten_embedding_2。

有关用法示例，请参阅模块文档中的“自定义算子”。

参数

symbolic_name (str) – 自定义算子的名称，格式为“<域>::<算子>”。
symbolic_fn (Callable) – 一个函数，接受 ONNX 图和当前算子的输入参数，并返回要添加到图中的新算子节点。
opset_version (int) – 要注册的 ONNX opset 版本。

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

注销 symbolic_name。

有关用法示例，请参阅模块文档中的“自定义算子”。

参数

symbolic_name (str) – 自定义算子的名称，格式为“<域>::<算子>”。
opset_version (int) – 要注销的 ONNX opset 版本。

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

一个上下文管理器，用于暂时将 model 的训练模式设置为 mode，并在退出 with 代码块时将其重置。

自版本 2.7 起弃用：请在导出模型之前设置训练模式。

参数

model – 类型和含义与 export() 的 model 参数相同。
mode (TrainingMode) – 类型和含义与 export() 的 training 参数相同。

torch.onnx.is_in_onnx_export()[source][source]¶

返回当前是否正在进行 ONNX 导出。

返回类型: bool

类 ¶

JitScalarType

在 torch 中定义的标量类型。

基于 TorchScript 的 ONNX 导出器¶

示例：从 PyTorch 到 ONNX 的 AlexNet ¶

追踪 (Tracing) 与脚本化 (Scripting)¶

避免陷阱 ¶

避免使用 NumPy 和内置 Python 类型 ¶

避免使用 Tensor.data ¶

在追踪模式下使用 tensor.shape 时避免原地操作 ¶

限制 ¶

类型 ¶

算子实现差异 ¶

不支持的张量索引模式 ¶

读取 / 获取 ¶

写入 / 设置 ¶

添加算子支持 ¶

ONNX 导出器内部机制 ¶

ATen 算子 ¶

支持的算子列表 ¶

添加对 aten 或量化算子的支持 ¶

torch.autograd.Functions ¶

静态符号方法 ¶

内联自动微分函数 ¶

自定义算子 ¶

ONNX-script 函数 ¶

C++ 算子 ¶

一次性发现所有不可转换的 ATen 算子 ¶

常见问题 ¶

Python API ¶

函数 ¶

类 ¶

文档

教程

资源