基于 TorchScript 的 ONNX 导出器 ¶

注意

要使用 TorchDynamo 而不是 TorchScript 导出 ONNX 模型，请参阅有关基于 TorchDynamo 的 ONNX 导出器的更多信息

示例：将 PyTorch 的 AlexNet 转换为 ONNX ¶

这里是一个简单的脚本，用于将预训练的 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 文件包含一个二进制协议缓冲区，其中包含您导出的模型（在这种情况下为 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 运行时后，您可以加载并运行该模型：

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 运行时运行它。

追踪与脚本化比较

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

• 跟踪：如果使用 torch.onnx.export() 调用一个尚未成为 ScriptModule 的模块，它首先执行 torch.jit.trace() 的等效操作，即使用给定的 args 执行模型一次，并记录执行过程中发生的所有操作。这意味着如果您的模型是动态的，例如根据输入数据改变行为，则导出的模型将无法捕获这种动态行为。我们建议检查导出的模型，确保操作符合理。跟踪将展开循环和 if 语句，导出与跟踪运行完全相同的静态图。如果您想导出具有动态控制流的模型，则需要使用脚本。

• 脚本：通过脚本编译模型可以保留动态控制流，适用于不同大小的输入。要使用脚本：

  • 使用 torch.jit.script() 生成 ScriptModule 。

  • 使用 torch.onnx.export() 调用模型，其中 ScriptModule 作为模型。 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 的 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)

而不是使用 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 输入值的计算都将被替换为在单次跟踪执行期间看到的常量值。

  • 任何输出为字典的都将静默地替换为其值的扁平序列（键将被移除）。例如， {"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.

写入/设置

当对 Tensor 进行索引写入时，以下模式不受支持：

# 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

当这种情况发生时，您可以尝试以下几种方法：

1. 将模型修改为不使用该运算符。

2. 创建一个符号函数来转换该运算符，并将其注册为自定义符号函数。

3. 为 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 函数相同，该函数可能声明在 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 算子。例如：

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”操作集的自定义操作符。当导出自定义操作符时，您可以使用导出时的 custom_opsets 字典指定自定义域版本。如果没有指定，自定义操作集版本默认为 1。

注意：请仔细对齐上述示例中提到的操作集版本，并确保它们在导出步骤中被消费。关于如何编写 onnx-script 函数的示例使用目前处于 onnx-script 活跃开发阶段的 beta 版本。请遵循最新的 ONNX-script

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

消费模型的运行时需要支持自定义运算符。请参阅 Caffe2 自定义运算符、ONNX 运行时自定义运算符或您选择的运行时的文档。

一次性发现所有不可转换的 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 上提交操作支持请求的问题。

常见问题解答 ¶

问：我已经导出了我的 LSTM 模型，但它的输入大小看起来是固定的吗？

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

Q: 如何导出包含循环的模型？

查看跟踪与脚本的区别。

Q: 如何导出具有原始类型输入（例如 int、float）的模型？

PyTorch 1.9 版本中增加了对原始数值类型输入的支持。然而，导出器不支持具有 str 输入的模型。

Q: ONNX 支持隐式标量数据类型转换吗？

ONNX 标准不支持，但导出器会尝试处理这部分。标量以常量张量的形式导出。导出器会确定标量的正确数据类型。在极少数无法做到的情况下，您需要手动指定数据类型，例如使用 dtype=torch.float32。如果您看到任何错误，请[在 GitHub 上创建问题](https://github.com/pytorch/pytorch/issues)。

Q: 张量列表可导出到 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)

将模型导出为 ONNX 格式。

参数:

• 导出模型（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 设置为具有以下模式的字典：

  • KEY (str): 输入或输出名称。每个名称也必须在 input_names 中提供。

    output_names.

  • VALUE (字典或列表)：如果是一个字典，键是轴索引，值是轴名称。如果是列表，

    则每个元素是一个轴索引。

  例如：

  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 (布尔值) –

  如果为 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 时有效。

• 报告（布尔值）- 是否为导出过程生成 Markdown 报告。此选项仅在 dynamo 为 True 时有效。

• 优化（布尔值）- 是否优化导出的模型。此选项仅在 dynamo 为 True 时有效。默认为 True。

• 验证（布尔值）- 是否使用 ONNX Runtime 验证导出的模型。此选项仅在 dynamo 为 True 时有效。

• 性能分析（布尔值）- 是否对导出过程进行性能分析。此选项仅在 dynamo 为 True 时有效。

• 是否将导出的程序输出到文件（布尔值）- 是否将 torch.export.ExportedProgram 输出到文件。这对于调试导出器很有用。此选项仅在 dynamo 为 True 时有效。

• 艺术品目录（字符串 | os.PathLike）- 保存调试艺术品（如报告和序列化的导出程序）的目录。此选项仅在 dynamo 为 True 时有效。

• 回退（布尔值）- 当 dynamo 导出器失败时是否回退到 TorchScript 导出器。此选项仅在 dynamo 为 True 时有效。当启用回退时，即使提供了 dynamic_shapes，也建议设置 dynamic_axes。

• 训练（_C_onnx.TrainingMode）- 已弃用选项。导出之前，请设置模型的训练模式。

• operator_export_type (_C_onnx.OperatorExportTypes) – 已废弃选项。仅支持 ONNX。

• do_constant_folding (bool) – 已废弃选项。

• custom_opsets (Mapping[str, int] | None) –

  已废弃。一个字典：

  • KEY (字符串): opset 域名

  • VALUE (整数): opset 版本

  如果自定义 opset 通过 model 引用但未在此字典中提及，则将 opset 版本设置为 1。仅通过此参数指示自定义 opset 域名和版本。

• export_modules_as_functions (布尔值 | torch.nn.Module 类型集合) –

  已弃用选项。

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

  1. 注解属性：通过 PEP 526 风格的类型注解的类变量将被导出为属性。注解属性在 ONNX 本地函数的子图中不被使用，因为它们不是由 PyTorch JIT 跟踪创建的，但消费者可以使用它们来确定是否用特定的融合内核替换函数。

  2. 推断属性：模块内部操作使用的变量。属性名称将带有前缀“inferred::”。这是为了区分从 Python 模块注解中检索到的预定义属性。推断属性在 ONNX 本地函数的子图中使用。

  • 默认导出（ False ）：将调用导出为细粒度节点。

  • 导出所有（ True ）：将调用导出为本地函数节点。

  • nn.Module 类型集合：导出（ nn.Module ）调用为本地函数节点，

    仅当找到类型集合中的 nn.Module 类型时。

• autograd_inlining (bool) – 已弃用。用于控制是否内联 autograd 函数的标志。有关更多详细信息，请参阅 https://github.com/pytorch/pytorch/pull/74765。

返回:

torch.onnx.ONNXProgram 如果 dynamo 为 True，否则为 None。

返回类型:

ONNXProgram | None

版本 2.6 中的变更：现在训练已弃用。导出前，请设置模型的训练模式。operator_export_type 已弃用。现在仅支持 ONNX。do_constant_folding 已弃用。它始终启用。export_modules_as_functions 已弃用。autograd_inlining 已弃用。

版本 2.7 变更：优化现在默认为 True。

torch.onnx.register_custom_op_symbolic(symbolic_name, symbolic_fn, opset_version)

注册自定义算子的符号函数。

当用户为自定义/贡献操作注册符号时，强烈建议通过 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)

取消注册 symbolic_name .

请参阅模块文档中的“自定义算子”部分以获取示例用法。

参数:

• symbolic_name (str) – 自定义算子的名称，格式为“::”。

• opset_version (int) – 要注销的 ONNX opset 版本。

torch.onnx.select_model_mode_for_export(model, mode)

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

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

参数:

• 模型 – 与 model 参数到 export() 的类型和意义相同。

• 模式（训练模式） – 与 training 参数到 export() 的类型和意义相同。

torch.onnx.is_in_onnx_export()

返回是否处于 ONNX 导出过程中。

返回类型:

布尔型

班级

JitScalarType

在 torch 中定义的标量类型。