這是一個簡單的腳本,可以將 Torchvision 中定義的經(jīng)過預(yù)訓練的 AlexNet 導出到 ONNX 中。 它運行一輪推斷,然后將生成的跟蹤模型保存到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 文件,其中包含您導出的模型的網(wǎng)絡(luò)結(jié)構(gòu)和參數(shù)(在本例中為 AlexNet)。 關(guān)鍵字參數(shù)verbose=True
使導出程序打印出人類可讀的網(wǎng)絡(luò)表示形式:
## 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 運行導出的腳本,您將需要安裝 <cite>caffe2</cite> :如果尚未安裝,請按照安裝說明進行操作。
一旦安裝了這些,就可以將后端用于 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 運行導出的模型,您將需要安裝 <cite>ONNX Runtime</cite> :請按照這些說明進行操作。
一旦安裝了這些,就可以將后端用于 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 導出器可以是基于跟蹤的和基于腳本的導出器。
我們允許混合跟蹤和腳本編寫。 您可以組合跟蹤和腳本以適合模型部分的特定要求。 看看這個例子:
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)
使用基于跟蹤的導出器,我們得到結(jié)果 ONNX 圖,該圖展開了 for 循環(huán):
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)
為了利用基于腳本的導出器捕獲動態(tài)循環(huán),我們可以在腳本中編寫循環(huán),然后從常規(guī) nn.Module 中調(diào)用它:
## 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'])
現(xiàn)在,導出的 ONNX 圖變?yōu)椋?/p>
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)
動態(tài)控制流已正確捕獲。 我們可以在具有不同循環(huán)范圍的后端進行驗證。
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)]
data = torch.zeros(3, 4)
index = torch.tensor(1)
new_data = torch.arange(4).to(torch.float32)
# Assigning to left hand side indexing is not supported in exporting.
# class InPlaceIndexedAssignment(torch.nn.Module):
# def forward(self, data, index, new_data):
# data[index] = new_data
# return data
class InPlaceIndexedAssignmentONNX(torch.nn.Module):
def forward(self, data, index, new_data):
new_data = new_data.unsqueeze(0)
index = index.expand(1, new_data.size(1))
data.scatter_(0, index, new_data)
return data
out = InPlaceIndexedAssignmentONNX()(data, index, new_data)
torch.onnx.export(InPlaceIndexedAssignmentONNX(), (data, index, new_data), 'inplace_assign.onnx')
# caffe2
import caffe2.python.onnx.backend as backend
import onnx
onnx_model = onnx.load('inplace_assign.onnx')
rep = backend.prepare(onnx_model)
out_caffe2 = rep.run((torch.zeros(3, 4).numpy(), index.numpy(), new_data.numpy()))
assert torch.all(torch.eq(out, torch.tensor(out_caffe2)))
# onnxruntime
import onnxruntime
sess = onnxruntime.InferenceSession('inplace_assign.onnx')
out_ort = sess.run(None, {
sess.get_inputs()[0].name: torch.zeros(3, 4).numpy(),
sess.get_inputs()[1].name: index.numpy(),
sess.get_inputs()[2].name: new_data.numpy(),
})
assert torch.all(torch.eq(out, torch.tensor(out_ort)))
x = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])
# This is not exportable
class Model(torch.nn.Module):
def forward(self, x):
return x.unbind(0)
# This is exportable.
# Note that in this example we know the split operator will always produce exactly three outputs,
# Thus we can export to ONNX without using tensor list.
class AnotherModel(torch.nn.Module):
def forward(self, x):
return [torch.squeeze(out, 0) for out in torch.split(x, [1,1,1], dim=0)]
支持以下運算符:
上面設(shè)置的運算符足以導出以下模型:
為操作員添加導出支持是的高級用法。 為此,開發(fā)人員需要觸摸 PyTorch 的源代碼。 請按照說明從源代碼安裝 PyTorch。 如果所需的運算符在 ONNX 中已標準化,則應(yīng)該容易添加對導出此類運算符的支持(為該運算符添加符號功能)。 要確認操作員是否標準化,請檢查 ONNX 操作員列表
。
如果該運算符是 ATen 運算符,則意味著您可以在torch/csrc/autograd/generated/VariableType.h
中找到該函數(shù)的聲明(可在 PyTorch 安裝目錄的生成代碼中找到),您應(yīng)在torch/onnx/symbolic_opset<version>.py
中添加符號函數(shù),并按照以下說明進行操作 :
torch/onnx/symbolic_opset<version>.py
中定義符號功能,例如 torch / onnx / symbolic_opset9.py 。 確保函數(shù)具有與VariableType.h
中定義的 ATen 運算符/函數(shù)相同的名稱。VariableType.h
中的名稱完全匹配,因為分配是通過關(guān)鍵字參數(shù)完成的。VariableType.h
中的匹配,張量(輸入)始終是第一個,然后是非張量參數(shù)。_scalar
可以將標量張量轉(zhuǎn)換為 python 標量,_if_scalar_type_as
可以將 Python 標量轉(zhuǎn)換為 PyTorch 張量。如果該運算符是非 ATen 運算符,則必須在相應(yīng)的 PyTorch Function 類中添加符號函數(shù)。 請閱讀以下說明:
symbolic
的符號函數(shù)。forward
中的名稱完全匹配。forward
的輸出匹配。符號函數(shù)應(yīng)在 Python 中實現(xiàn)。 所有這些功能都與通過 C ++-Python 綁定實現(xiàn)的 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
運算符缺失的符號函數(shù)的示例。 我們嘗試導出模型,并看到如下錯誤消息:
UserWarning: ONNX export failed on elu because torch.onnx.symbolic_opset9.elu does not exist
RuntimeError: ONNX export failed: Couldn't export operator elu
導出失敗,因為 PyTorch 不支持導出elu
運算符。 我們在VariableType.h
中找到virtual Tensor elu(const Tensor & input, Scalar alpha, bool inplace) const override;
。 這意味著elu
是 ATen 運算符。 我們檢查 ONNX 操作員列表,并確認Elu
在 ONNX 中已標準化。 我們在symbolic_opset9.py
中添加以下行:
def elu(g, input, alpha, inplace=False):
return g.op("Elu", input, alpha_f=_scalar(alpha))
現(xiàn)在,PyTorch 能夠?qū)С?code>elu運算符。
symbolic_opset9.py 和 symbolic_opset10.py 中還有更多示例。
用于指定操作員定義的界面是實驗性的; 冒險的用戶應(yīng)注意,API 可能會在將來的界面中更改。
按照本教程[使用自定義 C ++運算符擴展 TorchScript ] 之后,您可以在 PyTorch 中創(chuàng)建并注冊自己的自定義 ops 實現(xiàn)。 將這種模型導出到 ONNX 的方法如下:
## Create custom symbolic function
from torch.onnx.symbolic_helper import parse_args
@parse_args('v', 'v', 'f', 'i')
def symbolic_foo_forward(g, input1, input2, attr1, attr2):
return g.op("Foo", input1, input2, attr1_f=attr1, attr2_i=attr2)
## Register custom symbolic function
from torch.onnx import register_custom_op_symbolic
register_custom_op_symbolic('custom_ops::foo_forward', symbolic_foo_forward, 9)
class FooModel(torch.nn.Module):
def __init__(self, attr1, attr2):
super(FooModule, self).__init__()
self.attr1 = attr1
self.attr2 = attr2
def forward(self, input1, input2):
# Calling custom op
return torch.ops.custom_ops.foo_forward(input1, input2, self.attr1, self.attr2)
model = FooModel(attr1, attr2)
torch.onnx.export(model, (dummy_input1, dummy_input2), 'model.onnx')
根據(jù)自定義運算符的不同,您可以將其導出為現(xiàn)有 ONNX 操作之一或組合。 您也可以將其導出為 ONNX 中的自定義操作。 在這種情況下,您將需要通過匹配的自定義操作實現(xiàn)來擴展選擇的后端,例如 Caffe2 定制操作, ONNX Runtime 定制操作。
問:我已經(jīng)導出了我的 lstm 模型,但是它的輸入大小似乎是固定的?
跟蹤器將示例輸入形狀記錄在圖中。 如果模型應(yīng)接受動態(tài)形狀的輸入,則可以在導出 api 中使用參數(shù) <cite&dynamic_axes</cite& 。
layer_count = 4
model = nn.LSTM(10, 20, num_layers=layer_count, bidirectional=True)
model.eval()
with torch.no_grad():
input = torch.randn(5, 3, 10)
h0 = torch.randn(layer_count * 2, 3, 20)
c0 = torch.randn(layer_count * 2, 3, 20)
output, (hn, cn) = model(input, (h0, c0))
# default export
torch.onnx.export(model, (input, (h0, c0)), 'lstm.onnx')
onnx_model = onnx.load('lstm.onnx')
# input shape [5, 3, 10]
print(onnx_model.graph.input[0])
# export with `dynamic_axes`
torch.onnx.export(model, (input, (h0, c0)), 'lstm.onnx',
input_names=['input', 'h0', 'c0'],
output_names=['output', 'hn', 'cn'],
dynamic_axes={'input': {0: 'sequence'}, 'output': {0: 'sequence'}})
onnx_model = onnx.load('lstm.onnx')
# input shape ['sequence', 3, 10]
print(onnx_model.graph.input[0])
問:如何導出帶有循環(huán)的模型?
請簽出跟蹤與腳本編寫。
問:ONNX 是否支持隱式標量數(shù)據(jù)類型轉(zhuǎn)換?
不,但是出口商將嘗試處理該部分。 標量在 ONNX 中轉(zhuǎn)換為恒定張量。 導出器將嘗試找出標量的正確數(shù)據(jù)類型。 但是,對于無法執(zhí)行此操作的情況,您將需要手動提供數(shù)據(jù)類型信息。 這通常發(fā)生在腳本模型中,其中未記錄數(shù)據(jù)類型。 我們正在嘗試改進數(shù)據(jù)類型在導出器中的傳播,以便將來不再需要手動更改。
class ImplicitCastType(torch.jit.ScriptModule):
@torch.jit.script_method
def forward(self, x):
# Exporter knows x is float32, will export '2' as float32 as well.
y = x + 2
# Without type propagation, exporter doesn't know the datatype of y.
# Thus '3' is exported as int64 by default.
return y + 3
# The following will export correctly.
# return y + torch.tensor([3], dtype=torch.float32)
x = torch.tensor([1.0], dtype=torch.float32)
torch.onnx.export(ImplicitCastType(), x, 'models/implicit_cast.onnx',
example_outputs=ImplicitCastType()(x))
torch.onnx.export(model, args, f, export_params=True, verbose=False, training=False, input_names=None, output_names=None, aten=False, export_raw_ir=False, operator_export_type=None, opset_version=None, _retain_param_name=True, do_constant_folding=False, example_outputs=None, strip_doc_string=True, dynamic_axes=None, keep_initializers_as_inputs=None)?
將模型導出為 ONNX 格式。 這個導出器運行一次您的模型,以便跟蹤要導出的模型執(zhí)行情況。 目前,它支持一組有限的動態(tài)模型(例如 RNN)。
參數(shù)
model(*args)
是模型的有效調(diào)用。 任何非 Tensor 參數(shù)將被硬編碼到導出的模型中; 任何 Tensor 參數(shù)將按照在 args 中出現(xiàn)的順序成為導出模型的輸入。 如果 args 是一個 Tensor,則相當于用該 Tensor 的 1 元元組調(diào)用了它。 (注意:當前不支持將關(guān)鍵字參數(shù)傳遞給模型。如果需要,請給我們喊叫。)model.state_dict().values()
指定一個字典,用于指定輸入/輸出的動態(tài)軸,例如:-KEY:輸入和/或輸出名稱-VALUE:給定鍵的動態(tài)軸的索引,以及可能用于導出動態(tài)軸的名稱。 通常,該值是根據(jù)以下方式之一或兩者的組合定義的:(1)。 指定提供的輸入的動態(tài)軸的整數(shù)列表。 在這種情況下,將在導出過程中自動生成名稱并將其應(yīng)用于提供的輸入/輸出的動態(tài)軸。 或(2)。 一個內(nèi)部字典,該字典指定從對應(yīng)的輸入/輸出中的動態(tài)軸的索引到在導出過程中希望在此輸入/輸出的該軸上應(yīng)用的名稱的映射。
例。 如果我們的輸入和輸出具有以下形狀:
shape(input_1) = ('b', 3, 'w', 'h')
and shape(input_2) = ('b', 4)
and shape(output) = ('b', 'd', 5)
Then dynamic axes can be defined either as:
(a). ONLY INDICES:
dynamic_axes = {'input_1':[0,2,3],'input_2':[0],'output':[0,1]}
其中將為導出的動態(tài)軸生成自動名稱
(b). INDICES WITH CORRESPONDING NAMES:
dynamic_axes = {'input_1':{0:'batch',1:'width',2:'height'},'input_2':{0:'batch'},'output':{0:'batch', 1:“檢測”}
提供的名稱將應(yīng)用于導出的動態(tài)軸
(c). MIXED MODE OF (a) and (b)
dynamic_axes = {'input_1':[0,2,3],'input_2':{0:'batch'},'output':[0,1]}
torch.onnx.register_custom_op_symbolic(symbolic_name, symbolic_fn, opset_version)?
torch.onnx.operators.shape_as_tensor(x)?
torch.onnx.set_training(model, mode)?
上下文管理器將“模型”的訓練模式臨時設(shè)置為“模式”,當我們退出 with 塊時將其重置。 如果模式為“無”,則為無操作。
torch.onnx.is_in_onnx_export()?
檢查它是否在 ONNX 導出的中間。 此函數(shù)在 torch.onnx.export()的中間返回 True。 torch.onnx.export 應(yīng)該使用單線程執(zhí)行。
更多建議: