从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。 ScriptModule是TorchScript中的核心数据结构，而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实现，在精度很重要时，帮助消除这些差异，并记录这些差异。