Pytorch1.0.1导出ONNX模型

最新推荐文章于 2024-09-23 08:53:28 发布

丶Shining

最新推荐文章于 2024-09-23 08:53:28 发布

阅读量4.4k

点赞数

分类专栏： Linux部署 pytorch技巧文章标签： pytorch

本文链接：https://blog.csdn.net/xxradon/article/details/93767462

版权

Linux部署同时被 2 个专栏收录

45 篇文章 0 订阅

订阅专栏

pytorch技巧

15 篇文章 2 订阅

订阅专栏

TORCH.ONNX

Example: End-to-end AlexNet from PyTorch to Caffe2

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

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

结果是二进制protobuf文件，其中包含您导出的模型的网络结构和参数（在本例中为AlexNet）。关键字参数使导出器打印出人类可读的网络表示：alexnet.onnxverbose=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运行导出的脚本，您需要安装caffe2：如果您还没有caffe2，请按照安装说明进行操作。

安装完成后，您可以将后端用于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导出是一个基于Trace的出口国，这意味着它通过执行模型一次，导出此运行期间实际运行操作人员进行操作。这意味着如果您的模型是动态的，例如，根据输入数据更改行为，则导出将不准确。类似地，跟踪可能仅对特定输入大小有效（这是我们在跟踪时需要显式输入的一个原因。）我们建议检查模型跟踪并确保跟踪的运算符看起来合理。
PyTorch和Caffe2经常有运算符的实现，但有一些数字差异。根据模型结构，这些差异可以忽略不计，但它们也会导致行为的主要差异（特别是在未经训练的模型上）。在未来的版本中，我们计划允许Caffe2直接调用Torch实现的运算符，以帮助您平滑当精确度很重要时，这些差异，并记录这些差异。

Supported operators

The following operators are supported:

add (nonzero alpha not supported)

sub (nonzero alpha not supported)

mul

div

cat

mm

addmm

neg

sqrt

tanh

sigmoid

mean

sum

prod

t

expand (only when used before a broadcasting ONNX operator; e.g., add)

transpose

view

split

squeeze

prelu (single weight shared among input channels not supported)

threshold (non-zero threshold/non-zero value not supported)

leaky_relu

glu

softmax (only dim=-1 supported)

avg_pool2d (ceil_mode not supported)

log_softmax

unfold (experimental support with ATen-Caffe2 integration)

elu

concat

abs

index_select

pow

clamp

max

min

eq

gt

lt

ge

le

exp

sin

cos

tan

asin

acos

atan

permute

Conv

BatchNorm

MaxPool1d (ceil_mode not supported)

MaxPool2d (ceil_mode not supported)

MaxPool3d (ceil_mode not supported)

Embedding (no optional arguments supported)

RNN

ConstantPadNd

Dropout

FeatureDropout (training mode not supported)

Index (constant integer and tuple indices supported)

The operator set above is sufficient to export the following models:

AlexNet

DCGAN

DenseNet

Inception (warning: this model is highly sensitive to changes in operator implementation)

ResNet

SuperResolution

VGG

word_language_model

增加导出额外的operators是一种高级应用。为此，开发人员需要接触PyTorch的源代码。请按照从源安装PyTorch 的说明进行操作。如果想要的运算符在ONNX中是标准化的存在，则应该很容易添加对导出此类运算符的支持（为运算符添加符号函数）。要确认operators是否标准化，请检查 ONNX 目前支持的operators列表。

如果运算符是ATen运算符，这意味着您可以找到函数的声明torch/csrc/autograd/generated/VariableType.h （在PyTorch install目录中的生成代码中），您应该添加符号函数并按照下面列出的说明操作：torch/onnx/symbolic.py

在torch / onnx /symbolic.py中定义符号函数。确保该函数与在中定义的ATen运算符/函数同名VariableType.h。
第一个参数始终是导出的ONNX图。参数名称必须与名称完全匹配VariableType.h，因为调度是用关键字参数完成的。
参数排序不一定匹配VariableType.h，tensors（输入）始终是第一个，然后是non tensors参数。
在符号函数中，如果运算符已经在ONNX中标准化，我们只需要创建一个节点来表示图中的ONNX运算符。
如果输入参数是张量，但是ONNX要求标量，我们必须明确地进行转换。辅助函数_scalar可以将标量张量转换为python标量，并_if_scalar_type_as可以将Python标量转换为PyTorch张量。

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

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

符号函数应该在Python中实现。所有这些函数都与Python方法交互，这些方法是通过C ++实现的 - 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.elu does not exist
RuntimeError: ONNX export failed: Couldn't export operator elu

导出失败，因为PyTorch不支持导出elu运算符。我们找到了

virtual Tensor elu(const Tensor & input, Scalar alpha, bool inplace) const override;

VariableType.h。这意味着是一个ATen op。我们检查ONNX op 列表，并确认在ONNX中是标准化的。我们将以下行添加到symbolic.py：

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

现在PyTorch能够导出elu op。

symbolic.py， tensor.py， padding.py中有更多示例。

用于指定运算符定义的接口是实验性的; 喜欢冒险的用户应该注意，API可能会在未来的界面中发生变化。

功能

torch.onnx.export（* args，** kwargs ）