ICode9

精准搜索请尝试: 精确搜索
首页 > 其他分享> 文章详细

从PyTorch到ONNX的端到端AlexNet

2020-12-21 08:01:54  阅读:238  来源: 互联网

标签:ONNX torch PyTorch loop onnx input model 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实现,在精度很重要时,帮助消除这些差异,并记录这些差异。

 

标签:ONNX,torch,PyTorch,loop,onnx,input,model,AlexNet
来源: https://www.cnblogs.com/wujianming-110117/p/14166245.html

本站声明: 1. iCode9 技术分享网(下文简称本站)提供的所有内容,仅供技术学习、探讨和分享;
2. 关于本站的所有留言、评论、转载及引用,纯属内容发起人的个人观点,与本站观点和立场无关;
3. 关于本站的所有言论和文字,纯属内容发起人的个人观点,与本站观点和立场无关;
4. 本站文章均是网友提供,不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属;如您发现该文章侵犯了您的权益,可联系我们第一时间进行删除;
5. 本站为非盈利性的个人网站,所有内容不会用来进行牟利,也不会利用任何形式的广告来间接获益,纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。

专注分享技术,共同学习,共同进步。侵权联系[81616952@qq.com]

Copyright (C)ICode9.com, All Rights Reserved.

ICode9版权所有