从PyTorch到ONNX的端到端AlexNet

从PyTorch到ONNX的端到端AlexNet

这是一个简单的脚本,可将Torchvision中定义的经过预训练的AlexNet导出到ONNX中。运行一轮推理Inference,然后将生成的跟踪模型保存到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运行导出的脚本,将需要安装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 Runtime运行导出的模型,将需要安装ONNX Runtime

按照以下说明进行操作

一旦安装了这些,就可以将后端用于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输出可以同时trace-based and script-based

  • 基于跟踪的trace-based意思是,通过执行一次模型,并导出在此运行期间实际运行的算子进行操作。这意味着,如果模型是动态的,例如根据输入数据更改行为,则导出将不准确。同样,跟踪可能仅对特定的输入大小才有效(这是在跟踪时需要显式输入的原因之一。)建议检查模型跟踪,并确保跟踪的算子看起来合理。如果模型包含诸如for循环和if条件之类的控制流,则基于跟踪的输出,将展开循环和if条件,并输出与此运行完全相同的静态图。如果要使用动态控制流输出模型,则需要使用基于脚本的 输出。
  • 基于脚本的意思是,要导出的模型是ScriptModule。 ScriptModuleTorchScript的核心数据结构,而TorchScript是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)]

为避免将可变标量张量作为固定值常量,导出为ONNX模型的一部分,避免使用torch.Tensor.item()。torch支持将single-element张量隐式转换为数字。例如:

class LoopModel(torch.nn.Module):

    def forward(self, x, y):

        res = []

        arr = x.split(2, 0)

        for i in range(int(y)):

            res += [arr[i].sum(0, False)]

        return torch.stack(res)

 

model = torch.jit.script(LoopModel())

inputs = (torch.randn(16), torch.tensor(8))

 

out = model(*inputs)

torch.onnx.export(model, inputs, 'loop_and_list.onnx', opset_version=11, example_outputs=out)

TorchVision支持

除量化外,所有TorchVision模型均可导出到ONNX。可以在TorchVision中找到更多详细信息。

局限性

  • 仅将元组,列表和变量作为JIT输入/输出支持。也接受字典和字符串,但不建议使用。用户需要仔细验证自己的字典输入,并记住动态查询不可用。
  • PyTorch和ONNX后端(Caffe2,ONNX Runtime等)通常具有某些数字差异的算子实现。根据模型结构,这些差异可能可以忽略不计,但是也可能导致性能的重大差异(尤其是在未经训练的模型上。)允许Caffe2直接调用算子的Torch实现,在精度很重要时,帮助消除这些差异,并记录这些差异。

 

上一篇:百度开源:PaddleOCR与PaddlePaddle / paddle2onnx 实践一


下一篇:onnx c++