PyTorch torch.onnx

原文： PyTorch torch.onnx

示例：从 PyTorch 到 ONNX 的端到端 AlexNet

这是一个简单的脚本，可以将 Torchvision 中定义的经过预训练的 AlexNet 导出到 ONNX 中。它运行一轮推断，然后将生成的跟踪模型保存到alexnet.onnx：

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是二进制 protobuf 文件，其中包含您导出的模型的网络结构和参数(在本例中为 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 库来验证 protobuf。您可以使用 conda 安装ONNX：

conda install -c conda-forge onnx

然后，您可以运行：

import onnx


## Load the ONNX model
model = onnx.load("alexnet.onnx")


## Check that the IR is well formed
onnx.checker.check_model(model)


## Print a human readable representation of the graph
onnx.helper.printable_graph(model.graph)

要使用 caffe2 运行导出的脚本，您将需要安装 <cite>caffe2</cite> ：如果尚未安装，请按照安装说明进行操作。

一旦安装了这些，就可以将后端用于 Caffe2：

## ...continuing from above
import caffe2.python.onnx.backend as backend
import numpy as np


rep = backend.prepare(model, device="CUDA:0") # or "CPU"
## For the Caffe2 backend:
##     rep.predict_net is the Caffe2 protobuf for the network
##     rep.workspace is the Caffe2 workspace for the network
##       (see the class caffe2.python.onnx.backend.Workspace)
outputs = rep.run(np.random.randn(10, 3, 224, 224).astype(np.float32))
## To run networks with more than one input, pass a tuple
## rather than a single numpy ndarray.
print(outputs[0])

您还可以使用 ONNX Runtime 运行导出的模型，您将需要安装 <cite>ONNX Runtime</cite> ：请按照这些说明进行操作。

一旦安装了这些，就可以将后端用于 ONNX Runtime：

## ...continuing from above
import onnxruntime as ort


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

这是将 SuperResolution 模型导出到 ONNX 的另一本教程。。

将来，其他框架也会有后端。

跟踪与脚本编写

ONNX 导出器可以是基于跟踪的和基于脚本的导出器。

基于跟踪的表示它通过执行一次模型并导出在此运行期间实际运行的运算符进行操作。这意味着如果您的模型是动态的，例如根据输入数据更改行为，则导出将不准确。同样，跟踪可能仅对特定的输入大小才有效(这是我们在跟踪时需要显式输入的原因之一。）我们建议检查模型跟踪并确保所跟踪的运算符看起来合理。如果您的模型包含控制循环(如 for 循环）和 if 条件，则基于基于跟踪的导出器将展开循环以及 if 条件，并导出与此运行完全相同的静态图形。如果要使用动态控制流导出模型，则需要使用基于脚本的导出器。
基于脚本的表示您要导出的模型是 ScriptModule 。 <cite>ScriptModule</cite> 是 <cite>TorchScript</cite> 中的核心数据结构， <cite>TorchScript</cite> 是 Python 语言的子集，可从 PyTorch 代码创建可序列化和可优化的模型。

我们允许混合跟踪和脚本编写。您可以组合跟踪和脚本以适合模型部分的特定要求。看看这个例子：

import torch


## Trace-based only


class LoopModel(torch.nn.Module):
    def forward(self, x, y):
        for i in range(y):
            x = x + i
        return x


model = LoopModel()
dummy_input = torch.ones(2, 3, dtype=torch.long)
loop_count = torch.tensor(5, dtype=torch.long)


torch.onnx.export(model, (dummy_input, loop_count), 'loop.onnx', verbose=True)

使用基于跟踪的导出器，我们得到结果 ONNX 图，该图展开了 for 循环：

graph(%0 : Long(2, 3),
      %1 : Long()):
  %2 : Tensor = onnx::Constant[value={1}]()
  %3 : Tensor = onnx::Add(%0, %2)
  %4 : Tensor = onnx::Constant[value={2}]()
  %5 : Tensor = onnx::Add(%3, %4)
  %6 : Tensor = onnx::Constant[value={3}]()
  %7 : Tensor = onnx::Add(%5, %6)
  %8 : Tensor = onnx::Constant[value={4}]()
  %9 : Tensor = onnx::Add(%7, %8)
  return (%9)

为了利用基于脚本的导出器捕获动态循环，我们可以在脚本中编写循环，然后从常规 nn.Module 中调用它：

## Mixing tracing and scripting


@torch.jit.script
def loop(x, y):
    for i in range(int(y)):
        x = x + i
    return x


class LoopModel2(torch.nn.Module):
    def forward(self, x, y):
        return loop(x, y)


model = LoopModel2()
dummy_input = torch.ones(2, 3, dtype=torch.long)
loop_count = torch.tensor(5, dtype=torch.long)
torch.onnx.export(model, (dummy_input, loop_count), 'loop.onnx', verbose=True,
                  input_names=['input_data', 'loop_range'])

现在，导出的 ONNX 图变为：

graph(%input_data : Long(2, 3),
      %loop_range : Long()):
  %2 : Long() = onnx::Constant[value={1}](), scope: LoopModel2/loop
  %3 : Tensor = onnx::Cast[to=9](%2)
  %4 : Long(2, 3) = onnx::Loop(%loop_range, %3, %input_data), scope: LoopModel2/loop # custom_loop.py:240:5
    block0(%i.1 : Long(), %cond : bool, %x.6 : Long(2, 3)):
      %8 : Long(2, 3) = onnx::Add(%x.6, %i.1), scope: LoopModel2/loop # custom_loop.py:241:13
      %9 : Tensor = onnx::Cast[to=9](%2)
      -> (%9, %8)
  return (%4)

动态控制流已正确捕获。我们可以在具有不同循环范围的后端进行验证。

import caffe2.python.onnx.backend as backend
import numpy as np
import onnx
model = onnx.load('loop.onnx')


rep = backend.prepare(model)
outputs = rep.run((dummy_input.numpy(), np.array(9).astype(np.int64)))
print(outputs[0])
#[[37 37 37]
## [37 37 37]]


import onnxruntime as ort
ort_sess = ort.InferenceSession('loop.onnx')
outputs = ort_sess.run(None, {'input_data': dummy_input.numpy(),
                              'loop_range': np.array(9).astype(np.int64)})
print(outputs)
#[array([[37, 37, 37],
##       [37, 37, 37]], dtype=int64)]

局限性

导出中目前不支持张量就地索引分配，例如 <cite>data [index] = new_data</cite> 。解决此类问题的一种方法是使用运算符<cite>散布</cite>，显式更新原始张量。

  data = torch.zeros(3, 4)
  index = torch.tensor(1)
  new_data = torch.arange(4).to(torch.float32)

  
  # Assigning to left hand side indexing is not supported in exporting.
  # class InPlaceIndexedAssignment(torch.nn.Module):
  # def forward(self, data, index, new_data):
  #     data[index] = new_data
  #     return data

  
  class InPlaceIndexedAssignmentONNX(torch.nn.Module):
      def forward(self, data, index, new_data):
          new_data = new_data.unsqueeze(0)
          index = index.expand(1, new_data.size(1))
          data.scatter_(0, index, new_data)
          return data

  
  out = InPlaceIndexedAssignmentONNX()(data, index, new_data)

  
  torch.onnx.export(InPlaceIndexedAssignmentONNX(), (data, index, new_data), 'inplace_assign.onnx')

  
  # caffe2
  import caffe2.python.onnx.backend as backend
  import onnx

  
  onnx_model = onnx.load('inplace_assign.onnx')
  rep = backend.prepare(onnx_model)
  out_caffe2 = rep.run((torch.zeros(3, 4).numpy(), index.numpy(), new_data.numpy()))

  
  assert torch.all(torch.eq(out, torch.tensor(out_caffe2)))

  
  # onnxruntime
  import onnxruntime
  sess = onnxruntime.InferenceSession('inplace_assign.onnx')
  out_ort = sess.run(None, {
      sess.get_inputs()[0].name: torch.zeros(3, 4).numpy(),
      sess.get_inputs()[1].name: index.numpy(),
      sess.get_inputs()[2].name: new_data.numpy(),
  })

  
  assert torch.all(torch.eq(out, torch.tensor(out_ort)))

ONNX 中没有张量列表的概念。没有这个概念，很难导出消耗或产生张量列表的运算符，尤其是在导出时不知道张量列表的长度的情况下。

  x = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])

  
  # This is not exportable
  class Model(torch.nn.Module):
      def forward(self, x):
          return x.unbind(0)

  
  # This is exportable.
  # Note that in this example we know the split operator will always produce exactly three outputs,
  # Thus we can export to ONNX without using tensor list.
  class AnotherModel(torch.nn.Module):
      def forward(self, x):
          return [torch.squeeze(out, 0) for out in torch.split(x, [1,1,1], dim=0)]

仅将元组，列表和变量作为 JIT 输入/输出支持。也接受字典和字符串，但不建议使用它们。用户需要仔细验证自己的字典输入，并记住动态查询不可用。

PyTorch 和 ONNX 后端(Caffe2，ONNX 运行时等）通常具有一些数字差异的运算符实现。根据模型结构的不同，这些差异可能可以忽略不计，但是它们也可能导致行为上的重大差异(尤其是在未经训练的模型上。）我们允许 Caffe2 直接调用运算符的 Torch 实现，以在精度很重要时帮助您消除这些差异。，并记录这些差异。

支持的运营商

支持以下运算符：

批量标准
ConstantPadNd
转换
退出
嵌入(不支持可选参数）
FeatureDropout(不支持训练模式）
指数
MaxPool1d
MaxPool2d
MaxPool3d
RNN
腹肌
阿科斯
adaptive_avg_pool1d
adaptive_avg_pool2d
adaptive_avg_pool3d
adaptive_max_pool1d
adaptive_max_pool2d
adaptive_max_pool3d
添加(不支持非零 Alpha）
addmm
和
范围
argmax
精氨酸
阿辛
晒黑
avg_pool1d
avg_pool2d
avg_pool2d
avg_pool3d
Baddbmm
猫
细胞
钳
最大钳位
最小钳位
康卡特
cos
cumsum
暗淡的
div
退出
lu
空的
空的喜欢
当量
埃尔夫
经验值
扩大
expand_as
展平
地板
frobenius_norm
充分
满喜欢
收集
ge
格鲁
谷氨酸
gt
Hardtanh
index_
index_fill
index_select
instance_norm
插
伊斯南
layer_norm
乐
leaky_relu
日志
log1p
日志 2
log_sigmoid
log_softmax
对数表达式
lt
masked_fill
最大值
意思
分
毫米
多
多项式
狭窄
NE
负数
非零
规范
那些
喜欢
要么
置换
pixel_shuffle
战俘
prelu(不支持输入通道之间共享的单个权重）
产品
兰德
兰德
randn_like
倒数
Reflection_pad
露露
重复
复制垫
重塑
reshape_as
回合
雷雷鲁
rsqrt
订阅
分散
scatter_add
选择
塞卢
乙状结肠
标志
罪
尺寸
切片
软最大
软加
分类
分裂
sqrt
挤
堆
性病
子(不支持非零 Alpha）
和
Ť
棕褐色
谭
阈值(不支持非零阈值/非零值）
至
托普
转置
type_as
展开(与 ATen-Caffe2 集成的实验支持）
独特
松开
upsample_nearest1d
upsample_nearest2d
upsample_nearest3d
视图
哪里
零
zeros_like

上面设置的运算符足以导出以下模型：

亚历克斯网
DCGAN
密集网
初始阶段(警告：此模型对操作员实施的更改高度敏感）
ResNet
超分辨率
VGG
word_language_model

添加对运营商的支持

为操作员添加导出支持是的高级用法。为此，开发人员需要触摸 PyTorch 的源代码。请按照说明从源代码安装 PyTorch。如果所需的运算符在 ONNX 中已标准化，则应该容易添加对导出此类运算符的支持(为该运算符添加符号功能）。要确认操作员是否标准化，请检查 ONNX 操作员列表。

ATen 运算符

如果该运算符是 ATen 运算符，则意味着您可以在torch/csrc/autograd/generated/VariableType.h中找到该函数的声明(可在 PyTorch 安装目录的生成代码中找到），您应在torch/onnx/symbolic_opset<version>.py中添加符号函数，并按照以下说明进行操作：

在torch/onnx/symbolic_opset<version>.py中定义符号功能，例如 torch / onnx / symbolic_opset9.py 。确保函数具有与VariableType.h中定义的 ATen 运算符/函数相同的名称。
第一个参数始终是导出的 ONNX 图。参数名称必须与VariableType.h中的名称完全匹配，因为分配是通过关键字参数完成的。
参数排序不一定与VariableType.h中的匹配，张量(输入）始终是第一个，然后是非张量参数。
在符号功能中，如果运算符已经在 ONNX 中进行了标准化，我们只需要创建一个节点即可在图中表示 ONNX 运算符。
如果输入参数是张量，但 ONNX 要求标量，则必须显式进行转换。辅助函数_scalar可以将标量张量转换为 python 标量，_if_scalar_type_as可以将 Python 标量转换为 PyTorch 张量。

非 ATen 运营商

如果该运算符是非 ATen 运算符，则必须在相应的 PyTorch Function 类中添加符号函数。请阅读以下说明：

在相应的 Function 类中创建一个名为symbolic的符号函数。
第一个参数始终是导出的 ONNX 图。
除第一个参数名称外，参数名称必须与forward中的名称完全匹配。
输出元组大小必须与forward的输出匹配。
在符号功能中，如果运算符已经在 ONNX 中进行了标准化，我们只需要创建一个节点即可在图中表示 ONNX 运算符。

符号函数应在 Python 中实现。所有这些功能都与通过 C ++-Python 绑定实现的 Python 方法进行交互，但是直观地讲，它们提供的接口如下所示：

def operator/symbolic(g, *inputs):
  """
  Modifies Graph (e.g., using "op"), adding the ONNX operations representing
  this PyTorch function, and returning a Value or tuple of Values specifying the
  ONNX outputs whose values correspond to the original PyTorch return values
  of the autograd Function (or None if an output is not supported by ONNX).


  Arguments:
    g (Graph): graph to write the ONNX representation into
    inputs (Value...): list of values representing the variables which contain
        the inputs for this function
  """


class Value(object):
  """Represents an intermediate tensor value computed in ONNX."""
  def type(self):
    """Returns the Type of the value."""


class Type(object):
  def sizes(self):
    """Returns a tuple of ints representing the shape of a tensor this describes."""


class Graph(object):
  def op(self, opname, *inputs, **attrs):
    """
    Create an ONNX operator 'opname', taking 'args' as inputs
    and attributes 'kwargs' and add it as a node to the current graph,
    returning the value representing the single output of this
    operator (see the `outputs` keyword argument for multi-return
    nodes).


    The set of operators and the inputs/attributes they take
    is documented at https://github.com/onnx/onnx/blob/master/docs/Operators.md


    Arguments:
        opname (string): The ONNX operator name, e.g., `Abs` or `Add`.
        args (Value...): The inputs to the operator; usually provided
            as arguments to the `symbolic` definition.
        kwargs: The attributes of the ONNX operator, with keys named
            according to the following convention: `alpha_f` indicates
            the `alpha` attribute with type `f`.  The valid type specifiers are
            `f` (float), `i` (int), `s` (string) or `t` (Tensor).  An attribute
            specified with type float accepts either a single float, or a
            list of floats (e.g., you would say `dims_i` for a `dims` attribute
            that takes a list of integers).
        outputs (int, optional):  The number of outputs this operator returns;
            by default an operator is assumed to return a single output.
            If `outputs` is greater than one, this functions returns a tuple
            of output `Value`, representing each output of the ONNX operator
            in positional.
    """

ONNX 图形 C ++定义在torch/csrc/jit/ir.h中。

这是处理elu运算符缺失的符号函数的示例。我们尝试导出模型，并看到如下错误消息：

UserWarning: ONNX export failed on elu because torch.onnx.symbolic_opset9.elu does not exist
RuntimeError: ONNX export failed: Couldn't export operator elu

导出失败，因为 PyTorch 不支持导出elu运算符。我们在VariableType.h中找到virtual Tensor elu(const Tensor & input, Scalar alpha, bool inplace) const override;。这意味着elu是 ATen 运算符。我们检查 ONNX 操作员列表，并确认Elu在 ONNX 中已标准化。我们在symbolic_opset9.py中添加以下行：

def elu(g, input, alpha, inplace=False):
    return g.op("Elu", input, alpha_f=_scalar(alpha))

现在，PyTorch 能够导出elu运算符。

symbolic_opset9.py 和 symbolic_opset10.py 中还有更多示例。

用于指定操作员定义的界面是实验性的；冒险的用户应注意，API 可能会在将来的界面中更改。

定制运算符

按照本教程[使用自定义 C ++运算符扩展 TorchScript ] 之后，您可以在 PyTorch 中创建并注册自己的自定义 ops 实现。将这种模型导出到 ONNX 的方法如下：

## Create custom symbolic function
from torch.onnx.symbolic_helper import parse_args
@parse_args('v', 'v', 'f', 'i')
def symbolic_foo_forward(g, input1, input2, attr1, attr2):
    return g.op("Foo", input1, input2, attr1_f=attr1, attr2_i=attr2)


## Register custom symbolic function
from torch.onnx import register_custom_op_symbolic
register_custom_op_symbolic('custom_ops::foo_forward', symbolic_foo_forward, 9)


class FooModel(torch.nn.Module):
    def __init__(self, attr1, attr2):
        super(FooModule, self).__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, (dummy_input1, dummy_input2), 'model.onnx')

根据自定义运算符的不同，您可以将其导出为现有 ONNX 操作之一或组合。您也可以将其导出为 ONNX 中的自定义操作。在这种情况下，您将需要通过匹配的自定义操作实现来扩展选择的后端，例如 Caffe2 定制操作， ONNX Runtime 定制操作。

常见问题解答

问：我已经导出了我的 lstm 模型，但是它的输入大小似乎是固定的？

跟踪器将示例输入形状记录在图中。如果模型应接受动态形状的输入，则可以在导出 api 中使用参数 <cite&dynamic_axes</cite& 。

layer_count = 4 model = nn.LSTM(10, 20, num_layers=layer_count, bidirectional=True) model.eval() with torch.no_grad(): input = torch.randn(5, 3, 10) h0 = torch.randn(layer_count * 2, 3, 20) c0 = torch.randn(layer_count * 2, 3, 20) output, (hn, cn) = model(input, (h0, c0)) # default export torch.onnx.export(model, (input, (h0, c0)), 'lstm.onnx') onnx_model = onnx.load('lstm.onnx') # input shape [5, 3, 10] print(onnx_model.graph.input[0]) # export with `dynamic_axes` torch.onnx.export(model, (input, (h0, c0)), 'lstm.onnx', input_names=['input', 'h0', 'c0'], output_names=['output', 'hn', 'cn'], dynamic_axes={'input': {0: 'sequence'}, 'output': {0: 'sequence'}}) onnx_model = onnx.load('lstm.onnx') # input shape ['sequence', 3, 10] print(onnx_model.graph.input[0])

问：如何导出带有循环的模型？

请签出跟踪与脚本编写。

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

不，但是出口商将尝试处理该部分。标量在 ONNX 中转换为恒定张量。导出器将尝试找出标量的正确数据类型。但是，对于无法执行此操作的情况，您将需要手动提供数据类型信息。这通常发生在脚本模型中，其中未记录数据类型。我们正在尝试改进数据类型在导出器中的传播，以便将来不再需要手动更改。

class ImplicitCastType(torch.jit.ScriptModule): @torch.jit.script_method def forward(self, x): # Exporter knows x is float32, will export '2' as float32 as well. y = x + 2 # Without type propagation, exporter doesn't know the datatype of y. # Thus '3' is exported as int64 by default. return y + 3 # The following will export correctly. # return y + torch.tensor([3], dtype=torch.float32) x = torch.tensor([1.0], dtype=torch.float32) torch.onnx.export(ImplicitCastType(), x, 'models/implicit_cast.onnx', example_outputs=ImplicitCastType()(x))

功能

torch.onnx.export(model, args, f, export_params=True, verbose=False, training=False, input_names=None, output_names=None, aten=False, export_raw_ir=False, operator_export_type=None, opset_version=None, _retain_param_name=True, do_constant_folding=False, example_outputs=None, strip_doc_string=True, dynamic_axes=None, keep_initializers_as_inputs=None)¶

将模型导出为 ONNX 格式。这个导出器运行一次您的模型，以便跟踪要导出的模型执行情况。目前，它支持一组有限的动态模型(例如 RNN）。

参数

模型 (torch.nn.Module)–要导出的模型。

参数(参数元组）–模型的输入，例如，使得model(*args)是模型的有效调用。任何非 Tensor 参数将被硬编码到导出的模型中；任何 Tensor 参数将按照在 args 中出现的顺序成为导出模型的输入。如果 args 是一个 Tensor，则相当于用该 Tensor 的 1 元元组调用了它。 (注意：当前不支持将关键字参数传递给模型。如果需要，请给我们喊叫。）

f –类似于文件的对象(必须实现返回文件描述符的 fileno）或包含文件名的字符串。二进制 Protobuf 将被写入此文件。

export_params (布尔， 默认为 True )–如果指定，将导出所有参数。如果要导出未经训练的模型，请将其设置为 False。在这种情况下，导出的模型将首先以其所有参数作为参数，顺序由model.state_dict().values()指定

详细 (bool ， 默认为 False )–如果指定，我们将打印出导出跟踪的调试描述。

训练 (bool ， 默认为 False )–以训练模式导出模型。目前，ONNX 仅面向导出模型以进行推理，因此通常不需要将其设置为 True。

input_names (字符串列表 ， 默认空列表）–依次分配给图形输入节点的名称

output_names (字符串列表 ， 默认空列表）–依次分配给图形输出节点的名称

和 (bool ， 默认为 False )– [不推荐使用。使用 operator_export_type]以 aten 模式导出模型。如果使用 aten 模式，则 symbolic_opset <版本> .py 中的函数所导出的所有 ops 原始文件都将作为 ATen ops 导出。

export_raw_ir (布尔， 默认为 False )– [不建议使用。使用 operator_export_type]直接导出内部 IR，而不是将其转换为 ONNX ops。

operator_export_type (枚举， 默认 OperatorExportTypes.ONNX )– OperatorExportTypes.ONNX：所有操作均作为常规 ONNX 操作导出。 OperatorExportTypes.ONNX_ATEN：所有操作均导出为 ATen 操作。 OperatorExportTypes.ONNX_ATEN_FALLBACK：如果缺少符号，请使用 ATen op。 OperatorExportTypes.RAW：导出原始 ir。

opset_version (python：int ， 默认为 9 )–默认情况下，我们将模型导出到 onnx 子模块的 opset 版本。由于 ONNX 的最新 opset 可能会在下一个稳定版本之前发展，因此默认情况下，我们会导出到一个稳定的 opset 版本。目前，受支持的稳定 opset 版本为 9。opset_version 必须为 _onnx_master_opset 或在 torch / onnx / symbolic_helper.py 中定义的 _onnx_stable_opsets 中。

do_constant_folding (bool ， 默认 False )–如果为 True，则在导出期间将恒定折叠优化应用于模型。常量折叠优化将用预先计算的常量节点替换一些具有所有常量输入的操作。

example_outputs (张量元组， 默认无）–导出 ScriptModule 或 TorchScript 函数时必须提供 example_outputs。

strip_doc_string (bool ， 默认 True )–如果为 True，则从导出的模型中删除字段“ doc_string”，有关堆栈跟踪。

example_outputs –正在导出的模型的示例输出。

dynamic_axes (dict <字符串 ， dict < python：int ， 字符串 > > 或 dict <字符串 ，列表 ( python：int ） > ， 默认为空字典）–

一个字典，用于指定输入/输出的动态轴，例如：-KEY：输入和/或输出名称-VALUE：给定键的动态轴的索引，以及可能用于导出动态轴的名称。通常，该值是根据以下方式之一或两者的组合定义的：(1）。指定提供的输入的动态轴的整数列表。在这种情况下，将在导出过程中自动生成名称并将其应用于提供的输入/输出的动态轴。或(2）。一个内部字典，该字典指定从对应的输入/输出中的动态轴的索引到在导出过程中希望在此输入/输出的该轴上应用的名称的映射。

例。如果我们的输入和输出具有以下形状：

  shape(input_1) = ('b', 3, 'w', 'h')
  and shape(input_2) = ('b', 4)
  and shape(output)  = ('b', 'd', 5)

  Then dynamic axes can be defined either as:

  (a). ONLY INDICES:

dynamic_axes = {'input_1'：[0，2，3]，'input_2'：[0]，'output'：[0，1]}

其中将为导出的动态轴生成自动名称

  (b). INDICES WITH CORRESPONDING NAMES:

dynamic_axes = {'input_1'：{0：'batch'，1：'width'，2：'height'}，'input_2'：{0：'batch'}，'output'：{0：'batch'， 1：“检测”}

提供的名称将应用于导出的动态轴

  (c). MIXED MODE OF (a) and (b)

dynamic_axes = {'input_1'：[0，2，3]，'input_2'：{0：'batch'}，'output'：[0,1]}

keep_initializers_as_inputs (bool ， 默认值 None )–如果为 True，则导出的图中的所有初始化程序(通常对应于参数）也将被添加为图形的输入。如果为 False，则不会将初始化程序添加为图形的输入，而仅将非参数输入添加为输入。通过执行这些图形的后端/运行时，这可以允许进行更好的优化(例如恒定折叠等）。如果未指定(默认为“无”），则按以下方式自动选择行为。如果 operator_export_type 为 OperatorExportTypes.ONNX，则该行为等效于将此参数设置为 False。对于 operator_export_type 的其他值，此行为等同于将此参数设置为 True。请注意，对于 ONNX opset 版本< 9，初始化器必须是图形输入的一部分。因此，如果 opset_version 参数设置为 8 或更低，则该参数将被忽略。

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

torch.onnx.operators.shape_as_tensor(x)¶

torch.onnx.set_training(model, mode)¶

上下文管理器将“模型”的训练模式临时设置为“模式”，当我们退出 with 块时将其重置。如果模式为“无”，则为无操作。

torch.onnx.is_in_onnx_export()¶

检查它是否在 ONNX 导出的中间。此函数在 torch.onnx.export(）的中间返回 True。 torch.onnx.export 应该使用单线程执行。

PyTorch torch.onnx

示例：从 PyTorch 到 ONNX 的端到端 AlexNet

跟踪与脚本编写

局限性

支持的运营商

添加对运营商的支持

ATen 运算符

非 ATen 运营商

定制运算符

常见问题解答

功能

推荐文章

推荐教程

推荐课程