这是一个简单的脚本，可将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模型均可导出到ONNX。可以在TorchVision中找到更多详细信息。
- 仅将元组，列表和变量作为JIT输入/输出支持。也接受字典和字符串，但不建议使用。用户需要仔细验证自己的字典输入，并记住动态查询不可用。
- PyTorch和ONNX后端（Caffe2，ONNX Runtime等）通常具有某些数字差异的算子实现。根据模型结构，这些差异可能可以忽略不计，但是也可能导致性能的重大差异（尤其是在未经训练的模型上。）允许Caffe2直接调用算子的Torch实现，在精度很重要时，帮助消除这些差异，并记录这些差异。