torch.onnx

原文： https://pytorch.org/docs/stable/onnx.html

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

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

  1. import torch
  2. import torchvision
  4. dummy_input = torch.randn(10, 3, 224, 224, device='cuda')
  5. model = torchvision.models.alexnet(pretrained=True).cuda()
  7. # Providing input and output names sets the display names for values
  8. # within the model's graph. Setting these does not change the semantics
  9. # of the graph; it is only for readability.
  10. #
  11. # The inputs to the network consist of the flat list of inputs (i.e.
  12. # the values you would pass to the forward() method) followed by the
  13. # flat list of parameters. You can partially specify names, i.e. provide
  14. # a list here shorter than the number of inputs to the model, and we will
  15. # only set that subset of names, starting from the beginning.
  16. input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]
  17. output_names = [ "output1" ]
  19. torch.onnx.export(model, dummy_input, "alexnet.onnx", verbose=True, input_names=input_names, output_names=output_names)

生成的alexnet.onnx是二进制 protobuf 文件，其中包含您导出的模型的网络结构和参数(在本例中为 AlexNet）。 关键字参数verbose=True使导出程序打印出人类可读的网络表示形式：

  1. # These are the inputs and parameters to the network, which have taken on
  2. # the names we specified earlier.
  3. graph(%actual_input_1 : Float(10, 3, 224, 224)
  4.       %learned_0 : Float(64, 3, 11, 11)
  5.       %learned_1 : Float(64)
  6.       %learned_2 : Float(192, 64, 5, 5)
  7.       %learned_3 : Float(192)
  8.       # ---- omitted for brevity ----
  9.       %learned_14 : Float(1000, 4096)
  10.       %learned_15 : Float(1000)) {
  11.   # Every statement consists of some output tensors (and their types),
  12.   # the operator to be run (with its attributes, e.g., kernels, strides,
  13.   # etc.), its input tensors (%actual_input_1, %learned_0, %learned_1)
  14.   %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]
  15.   %18 : Float(10, 64, 55, 55) = onnx::Relu(%17), scope: AlexNet/Sequential[features]/ReLU[1]
  16.   %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]
  17.   # ---- omitted for brevity ----
  18.   %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]
  19.   # Dynamic means that the shape is not known. This may be because of a
  20.   # limitation of our implementation (which we would like to fix in a
  21.   # future release) or shapes which are truly dynamic.
  22.   %30 : Dynamic = onnx::Shape(%29), scope: AlexNet
  23.   %31 : Dynamic = onnx::Slice[axes=[0], ends=[1], starts=[0]](%30), scope: AlexNet
  24.   %32 : Long() = onnx::Squeeze[axes=[0]](%31), scope: AlexNet
  25.   %33 : Long() = onnx::Constant[value={9216}](), scope: AlexNet
  26.   # ---- omitted for brevity ----
  27.   %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]
  28.   return (%output1);
  29. }

您还可以使用 ONNX 库来验证 protobuf。 您可以使用 conda 安装ONNX：

  1. conda install -c conda-forge onnx

然后，您可以运行：

  1. import onnx
  3. # Load the ONNX model
  4. model = onnx.load("alexnet.onnx")
  6. # Check that the IR is well formed
  7. onnx.checker.check_model(model)
  9. # Print a human readable representation of the graph
  10. onnx.helper.printable_graph(model.graph)

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

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

  1. # ...continuing from above
  2. import caffe2.python.onnx.backend as backend
  3. import numpy as np
  5. rep = backend.prepare(model, device="CUDA:0") # or "CPU"
  6. # For the Caffe2 backend:
  7. #     rep.predict_net is the Caffe2 protobuf for the network
  8. #     rep.workspace is the Caffe2 workspace for the network
  9. #       (see the class caffe2.python.onnx.backend.Workspace)
  10. outputs = rep.run(np.random.randn(10, 3, 224, 224).astype(np.float32))
  11. # To run networks with more than one input, pass a tuple
  12. # rather than a single numpy ndarray.
  13. print(outputs[0])

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

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

  1. # ...continuing from above
  2. import onnxruntime as ort
  4. ort_session = ort.InferenceSession('alexnet.onnx')
  6. outputs = ort_session.run(None, {'actual_input_1': np.random.randn(10, 3, 224, 224).astype(np.float32)})
  8. print(outputs[0])

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

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

跟踪与脚本编写

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

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

  • 基于脚本的表示您要导出的模型是 ScriptModule 。 <cite>ScriptModule</cite> 是 <cite>TorchScript</cite> 中的核心数据结构， <cite>TorchScript</cite> 是 Python 语言的子集，可从 PyTorch 代码创建可序列化和可优化的模型。

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

  1. import torch
  3. # Trace-based only
  5. class LoopModel(torch.nn.Module):
  6.     def forward(self, x, y):
  7.         for i in range(y):
  8.             x = x + i
  9.         return x
  11. model = LoopModel()
  12. dummy_input = torch.ones(2, 3, dtype=torch.long)
  13. loop_count = torch.tensor(5, dtype=torch.long)
  15. torch.onnx.export(model, (dummy_input, loop_count), 'loop.onnx', verbose=True)

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

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

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

  1. # Mixing tracing and scripting
  3. @torch.jit.script
  4. def loop(x, y):
  5.     for i in range(int(y)):
  6.         x = x + i
  7.     return x
  9. class LoopModel2(torch.nn.Module):
  10.     def forward(self, x, y):
  11.         return loop(x, y)
  13. model = LoopModel2()
  14. dummy_input = torch.ones(2, 3, dtype=torch.long)
  15. loop_count = torch.tensor(5, dtype=torch.long)
  16. torch.onnx.export(model, (dummy_input, loop_count), 'loop.onnx', verbose=True,
  17.                   input_names=['input_data', 'loop_range'])

现在，导出的 ONNX 图变为：

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

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

  1. import caffe2.python.onnx.backend as backend
  2. import numpy as np
  3. import onnx
  4. model = onnx.load('loop.onnx')
  6. rep = backend.prepare(model)
  7. outputs = rep.run((dummy_input.numpy(), np.array(9).astype(np.int64)))
  8. print(outputs[0])
  9. #[[37 37 37]
  10. # [37 37 37]]
  12. import onnxruntime as ort
  13. ort_sess = ort.InferenceSession('loop.onnx')
  14. outputs = ort_sess.run(None, {'input_data': dummy_input.numpy(),
  15.                               'loop_range': np.array(9).astype(np.int64)})
  16. print(outputs)
  17. #[array([[37, 37, 37],
  18. #       [37, 37, 37]], dtype=int64)]

局限性

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

    1. data = torch.zeros(3, 4)
    2. index = torch.tensor(1)
    3. new_data = torch.arange(4).to(torch.float32)
    5. # Assigning to left hand side indexing is not supported in exporting.
    6. # class InPlaceIndexedAssignment(torch.nn.Module):
    7. # def forward(self, data, index, new_data):
    8. #     data[index] = new_data
    9. #     return data
    11. class InPlaceIndexedAssignmentONNX(torch.nn.Module):
    12.     def forward(self, data, index, new_data):
    13.         new_data = new_data.unsqueeze(0)
    14.         index = index.expand(1, new_data.size(1))
    15.         data.scatter_(0, index, new_data)
    16.         return data
    18. out = InPlaceIndexedAssignmentONNX()(data, index, new_data)
    20. torch.onnx.export(InPlaceIndexedAssignmentONNX(), (data, index, new_data), 'inplace_assign.onnx')
    22. # caffe2
    23. import caffe2.python.onnx.backend as backend
    24. import onnx
    26. onnx_model = onnx.load('inplace_assign.onnx')
    27. rep = backend.prepare(onnx_model)
    28. out_caffe2 = rep.run((torch.zeros(3, 4).numpy(), index.numpy(), new_data.numpy()))
    30. assert torch.all(torch.eq(out, torch.tensor(out_caffe2)))
    32. # onnxruntime
    33. import onnxruntime
    34. sess = onnxruntime.InferenceSession('inplace_assign.onnx')
    35. out_ort = sess.run(None, {
    36.     sess.get_inputs()[0].name: torch.zeros(3, 4).numpy(),
    37.     sess.get_inputs()[1].name: index.numpy(),
    38.     sess.get_inputs()[2].name: new_data.numpy(),
    39. })
    41. assert torch.all(torch.eq(out, torch.tensor(out_ort)))

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

    1. x = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])
    3. # This is not exportable
    4. class Model(torch.nn.Module):
    5.     def forward(self, x):
    6.         return x.unbind(0)
    8. # This is exportable.
    9. # Note that in this example we know the split operator will always produce exactly three outputs,
    10. # Thus we can export to ONNX without using tensor list.
    11. class AnotherModel(torch.nn.Module):
    12.     def forward(self, x):
    13.         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_copy

  • 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&lt;version&gt;.py中添加符号函数，并按照以下说明进行操作 ：

  • 在torch/onnx/symbolic_opset&lt;version&gt;.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 方法进行交互，但是直观地讲，它们提供的接口如下所示：

  1. def operator/symbolic(g, *inputs):
  2.   """
  3.   Modifies Graph (e.g., using "op"), adding the ONNX operations representing
  4.   this PyTorch function, and returning a Value or tuple of Values specifying the
  5.   ONNX outputs whose values correspond to the original PyTorch return values
  6.   of the autograd Function (or None if an output is not supported by ONNX).
  8.   Arguments:
  9.     g (Graph): graph to write the ONNX representation into
  10.     inputs (Value...): list of values representing the variables which contain
  11.         the inputs for this function
  12.   """
  14. class Value(object):
  15.   """Represents an intermediate tensor value computed in ONNX."""
  16.   def type(self):
  17.     """Returns the Type of the value."""
  19. class Type(object):
  20.   def sizes(self):
  21.     """Returns a tuple of ints representing the shape of a tensor this describes."""
  23. class Graph(object):
  24.   def op(self, opname, *inputs, **attrs):
  25.     """
  26.     Create an ONNX operator 'opname', taking 'args' as inputs
  27.     and attributes 'kwargs' and add it as a node to the current graph,
  28.     returning the value representing the single output of this
  29.     operator (see the `outputs` keyword argument for multi-return
  30.     nodes).
  32.     The set of operators and the inputs/attributes they take
  33.     is documented at https://github.com/onnx/onnx/blob/master/docs/Operators.md
  35.     Arguments:
  36.         opname (string): The ONNX operator name, e.g., `Abs` or `Add`.
  37.         args (Value...): The inputs to the operator; usually provided
  38.             as arguments to the `symbolic` definition.
  39.         kwargs: The attributes of the ONNX operator, with keys named
  40.             according to the following convention: `alpha_f` indicates
  41.             the `alpha` attribute with type `f`.  The valid type specifiers are
  42.             `f` (float), `i` (int), `s` (string) or `t` (Tensor).  An attribute
  43.             specified with type float accepts either a single float, or a
  44.             list of floats (e.g., you would say `dims_i` for a `dims` attribute
  45.             that takes a list of integers).
  46.         outputs (int, optional):  The number of outputs this operator returns;
  47.             by default an operator is assumed to return a single output.
  48.             If `outputs` is greater than one, this functions returns a tuple
  49.             of output `Value`, representing each output of the ONNX operator
  50.             in positional.
  51.     """

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

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

  1. UserWarning: ONNX export failed on elu because torch.onnx.symbolic_opset9.elu does not exist
  2. 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中添加以下行：

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

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

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

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

定制运算符

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

  1. # Create custom symbolic function
  2. from torch.onnx.symbolic_helper import parse_args
  3. @parse_args('v', 'v', 'f', 'i')
  4. def symbolic_foo_forward(g, input1, input2, attr1, attr2):
  5.     return g.op("Foo", input1, input2, attr1_f=attr1, attr2_i=attr2)
  7. # Register custom symbolic function
  8. from torch.onnx import register_custom_op_symbolic
  9. register_custom_op_symbolic('custom_ops::foo_forward', symbolic_foo_forward, 9)
  11. class FooModel(torch.nn.Module):
  12.     def __init__(self, attr1, attr2):
  13.         super(FooModule, self).__init__()
  14.         self.attr1 = attr1
  15.         self.attr2 = attr2
  17.     def forward(self, input1, input2):
  18.         # Calling custom op
  19.         return torch.ops.custom_ops.foo_forward(input1, input2, self.attr1, self.attr2)
  21. model = FooModel(attr1, attr2)
  22. torch.onnx.export(model, (dummy_input1, dummy_input2), 'model.onnx')

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

常见问题解答

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

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

  1. layer_count = 4
  3. model = nn.LSTM(10, 20, num_layers=layer_count, bidirectional=True)
  4. model.eval()
  6. with torch.no_grad():
  7.     input = torch.randn(5, 3, 10)
  8.     h0 = torch.randn(layer_count * 2, 3, 20)
  9.     c0 = torch.randn(layer_count * 2, 3, 20)
  10.     output, (hn, cn) = model(input, (h0, c0))
  12.     # default export
  13.     torch.onnx.export(model, (input, (h0, c0)), 'lstm.onnx')
  14.     onnx_model = onnx.load('lstm.onnx')
  15.     # input shape [5, 3, 10]
  16.     print(onnx_model.graph.input[0])
  18.     # export with `dynamic_axes`
  19.     torch.onnx.export(model, (input, (h0, c0)), 'lstm.onnx',
  20.                     input_names=['input', 'h0', 'c0'],
  21.                     output_names=['output', 'hn', 'cn'],
  22.                     dynamic_axes={'input': {0: 'sequence'}, 'output': {0: 'sequence'}})
  23.     onnx_model = onnx.load('lstm.onnx')
  24.     # input shape ['sequence', 3, 10]
  25.     print(onnx_model.graph.input[0])

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

请签出跟踪与脚本编写。

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

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

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

功能

  1. 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）。 一个内部字典，该字典指定从对应的输入/输出中的动态轴的索引到在导出过程中希望在此输入/输出的该轴上应用的名称的映射。

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

    1. shape(input_1) = ('b', 3, 'w', 'h')
    2. and shape(input_2) = ('b', 4)
    3. and shape(output)  = ('b', 'd', 5)
    1. Then dynamic axes can be defined either as:
    1. (a). ONLY INDICES:

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

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

    1. (b). INDICES WITH CORRESPONDING NAMES:

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

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

    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 或更低，则该参数将被忽略。

  1. torch.onnx.register_custom_op_symbolic(symbolic_name, symbolic_fn, opset_version)¶
  1. torch.onnx.operators.shape_as_tensor(x)¶
  1. torch.onnx.set_training(model, mode)¶

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

  1. torch.onnx.is_in_onnx_export()¶

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