ONNX:最流行的模型 IR

概述

ONNX全称为 Open Neural Network Exchange,是一种与框架无关的中间表达(IR)。ONNX的规范及代码主要由微软,亚马逊 ,Facebook 和 IBM 等公司共同开发,以开放源代码的方式托管在Github上。目前官方支持加载ONNX模型并进行推理的深度学习框架有: Caffe2, PyTorch, MXNet,ML.NET,TensorRT 和 Microsoft CNTK,并且 TensorFlow 也非官方的支持 ONNX。随着项目的推进,该 IR 已经被主流的推理框架所支持,逐渐发展成为最为通用的 IR。

image-20221216225359323

ONNX 的数据格式

简介

ONNX本质上一种文件格式,通过Protobuf数据结构存储了神经网络结构权重。其组织格式核心定义在 Open Neural Network Exchange Intermediate Representation (ONNX IR) Specification,其中定义了Model/Graph/Node/ValueInfo/Tensor/Attribute 层面的数据结构。整图通过各节点(Node)的input/output指向关系构建模型图的拓扑结构。

从结果上看

image-20221216225056734

1、左图

例如下面代码中应用 ONNX 表示的 mnist 模型。该模型包含一个输入 Tensor (Input3) 和一个输出 Tensor (Plus214_Output_0),而节点 Node 0 的输入正是 Tensor Input3, 此外,该节点是 Conv 算子,该算子的属性值可以在 Attribute 中找到,而训练好的权重值(weight)Parameter5 可以在 Initializer 中找到,而每个 Node 中 -> 表示拓扑关系,通过指向关系可以构建出计算图。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
[W] --mode is deprecated and will be removed in Polygraphy 0.40.0. Use --show instead.
[I] Loading model: mnist-12.onnx
[I] ==== ONNX Model ====
Name: CNTKGraph | ONNX Opset: 12

---- 1 Graph Input(s) ----
{Input3 [dtype=float32, shape=(1, 1, 28, 28)]}

---- 1 Graph Output(s) ----
{Plus214_Output_0 [dtype=float32, shape=(1, 10)]}

---- 8 Initializer(s) ----
Initializer | Parameter5 [dtype=float32, shape=[8, 1, 5, 5]] | Values:
[[[[-0.00890567 -0.23690744 -0.50882167 -0.06456178 0.14181185]
[-0.59197617 -0.47528538 -0.04934815 0.7682156 0.2634652 ]
[-0.49176344 0.05561765 1.0189645 0.5547042 -0.4416644 ]
...

Initializer | Parameter194 [dtype=float32, shape=[1, 10]] | Values:
[[-0.04485603 0.00779166 0.06810082 0.02999374 -0.12640963 0.14021875
-0.0552849 -0.04938382 0.08432205 -0.05454041]]

---- 12 Node(s) ----
Node 0 | Convolution28 [Op: Conv]
{Input3 [dtype=float32, shape=(1, 1, 28, 28)],
Initializer | Parameter5 [dtype=float32, shape=(8, 1, 5, 5)]}
-> {Convolution28_Output_0 [dtype=float32, shape=(1, 8, 28, 28)]}
---- Attributes ----
Convolution28.kernel_shape = [5, 5]
Convolution28.strides = [1, 1]
Convolution28.auto_pad = SAME_UPPER
Convolution28.group = 1
Convolution28.dilations = [1, 1]

...

Node 11 | Plus214 [Op: Add]
{Times212_Output_0 [dtype=float32, shape=(1, 10)],
Initializer | Parameter194 [dtype=float32, shape=(1, 10)]}
-> {Plus214_Output_0 [dtype=float32, shape=(1, 10)]}

此外,每个节点对应的 Op 可能包含一些先验参数的设置可以在 Attributes 中找到,例如

1
2
3
4
5
6
7
8
9
10
Node 0    | Convolution28 [Op: Conv]
{Input3 [dtype=float32, shape=(1, 1, 28, 28)],
Initializer | Parameter5 [dtype=float32, shape=(8, 1, 5, 5)]}
-> {Convolution28_Output_0 [dtype=float32, shape=(1, 8, 28, 28)]}
---- Attributes ----
Convolution28.kernel_shape = [5, 5]
Convolution28.strides = [1, 1]
Convolution28.auto_pad = SAME_UPPER
Convolution28.group = 1
Convolution28.dilations = [1, 1]
Initializer 中就可以找到已经训练好的模型参数 Values, 例如 Parameter 5。

2、右图

ONNX 不仅能表示正常训练的模型,它也包括对量化 INT8 模型表示的,例如 mnit-12-int8模型,相比较 mnit-12 模型最突出的特点是 Initializer 个数明显增多,不仅要有权重(从浮点 -> int8), 还需要额外的 scale + zero_point 参数,不仅原来的权重需要 Initializer , 相关的 Op 也需要额外的量化参数。所以 Initializer 数量增加几倍。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
$ polygraphy inspect model mnist-12-int8.onnx --model-type onnx --mode basic
[W] --mode is deprecated and will be removed in Polygraphy 0.40.0. Use --show instead.
[I] Loading model: mnist-12-int8.onnx
[I] ==== ONNX Model ====
Name: CNTKGraph | ONNX Opset: 12 | Other Opsets: {'ai.onnx.preview.training': 1, 'ai.onnx.training': 1, 'com.ms.internal.nhwc': 1, 'org.pytorch.aten': 1, 'com.microsoft': 1, 'ai.onnx.contrib': 1000, 'com.microsoft.nchwc': 1, 'com.microsoft.experimental': 1, 'ai.onnx.ml': 3}

---- 1 Graph Input(s) ----
{Input3 [dtype=float32, shape=(1, 1, 28, 28)]}

---- 1 Graph Output(s) ----
{Plus214_Output_0 [dtype=float32, shape=(1, 10)]}

---- 29 Initializer(s) ----
{Initializer | Input3_zero_point [dtype=uint8, shape=[]] | Values:
0
Initializer | Input3_scale [dtype=float32, shape=[]] | Values:
1.0
Initializer | Parameter5_quantized [dtype=int8, shape=[8, 1, 5, 5]] | Values:
[[[[ -1 -30 -63 -8 18]
[ -74 -59 -6 96 33]
[ -61 7 127 69 -55]
[ -20 69 74 -37 -76]
[ 5 28 -27 -59 -36]]]
...
[ -8 -43 -96 -64 -43]]]]
Initializer | Parameter5_scale [dtype=float32, shape=[8]] | Values:
[0.00802334 0.00446989 0.00765891 0.00375404 0.00538003 0.00577387
0.00440539 0.00446567]
Initializer | Parameter5_zero_point [dtype=int8, shape=[8]] | Values:
[0 0 0 0 0 0 0 0]
Initializer | Convolution28_zero_point [dtype=uint8, shape=[]] | Values:
0
Initializer | Convolution28_scale [dtype=float32, shape=[]] | Values:
3.6796057
Initializer | ConvAddFusion_Add_B_Parameter6_quantized [dtype=int32, shape=[8]] | Values:
[-20 -97 12 -4 -12 -23 5 -27]
...
---- 11 Node(s) ----
Node 0 | Input3_Convolution28_QuantizeLinear [Op: QuantizeLinear]
{Input3 [dtype=float32, shape=(1, 1, 28, 28)],
Initializer | Input3_scale [dtype=float32, shape=()],
Initializer | Input3_zero_point [dtype=uint8, shape=()]}
-> {Input3_Convolution28_QuantizeLinear}

Node 1 | Convolution28_quant [Op: QLinearConv]
{Input3_Convolution28_QuantizeLinear,
Initializer | Input3_scale [dtype=float32, shape=()],
Initializer | Input3_zero_point [dtype=uint8, shape=()],
Initializer | Parameter5_quantized [dtype=int8, shape=(8, 1, 5, 5)],
Initializer | Parameter5_scale [dtype=float32, shape=(8,)],
Initializer | Parameter5_zero_point [dtype=int8, shape=(8,)],
Initializer | Convolution28_scale [dtype=float32, shape=()],
Initializer | Convolution28_zero_point [dtype=uint8, shape=()],
Initializer | ConvAddFusion_Add_B_Parameter6_quantized [dtype=int32, shape=(8,)]}
-> {Convolution28_QuantizeLinear_0}
---- Attributes ----
Convolution28_quant.auto_pad = SAME_UPPER
Convolution28_quant.dilations = [1, 1]
Convolution28_quant.group = 1
Convolution28_quant.kernel_shape = [5, 5]
Convolution28_quant.strides = [1, 1]
...

Node 8 | gemm_MatMul_quant [Op: QLinearMatMul]
{Pooling160_Output_0_reshape0_gemm_MatMul_QuantizeLinear,
Initializer | Pooling160_Output_0_reshape0_scale [dtype=float32, shape=()],
Initializer | Pooling160_Output_0_reshape0_zero_point [dtype=uint8, shape=()],
Initializer | Parameter193_reshape1_quantized [dtype=int8, shape=(256, 10)],
Initializer | Parameter193_reshape1_scale [dtype=float32, shape=(10,)],
Initializer | Parameter193_reshape1_zero_point [dtype=int8, shape=(10,)],
Initializer | Plus214_Output_0_MatMul_scale [dtype=float32, shape=()],
Initializer | Plus214_Output_0_MatMul_zero_point [dtype=uint8, shape=()]}
-> {Plus214_Output_0_MatMul_QuantizeLinear_0}

Node 9 | gemm_Add_quant [Op: QLinearAdd]
{Plus214_Output_0_MatMul_QuantizeLinear_0,
Initializer | Plus214_Output_0_MatMul_scale [dtype=float32, shape=()],
Initializer | Plus214_Output_0_MatMul_zero_point [dtype=uint8, shape=()],
Initializer | Parameter194_quantized [dtype=uint8, shape=(1, 10)],
Initializer | Parameter194_scale [dtype=float32, shape=()],
Initializer | Parameter194_zero_point [dtype=uint8, shape=()],
Initializer | Plus214_Output_0_scale [dtype=float32, shape=()],
Initializer | Plus214_Output_0_zero_point [dtype=uint8, shape=()]}
-> {Plus214_Output_0_QuantizeLinear_0}

Node 10 | Plus214_Output_0_DequantizeLinear_0 [Op: DequantizeLinear]
{Plus214_Output_0_QuantizeLinear_0,
Initializer | Plus214_Output_0_scale [dtype=float32, shape=()],
Initializer | Plus214_Output_0_zero_point [dtype=uint8, shape=()]}
-> {Plus214_Output_0 [dtype=float32, shape=(1, 10)]}

Proto 定义中看

Graphs

从模型 Model 定义中可以看到,一个模型中有一个 Graph 的属性。而这个 Graph的定义如下:

Name Type Description
name string 计算图的名字
node Node[] 节点列表,根据输入/输出数据依赖关系形成部分有序的计算图。它是按拓扑顺序的。
initializer Tensor[] 张量的列表。当初始化式与图形输入具有相同的名称时,它将为该输入指定一个默认值。当初始化式的名称与所有图形输入不同时,它指定一个常量值。列表的顺序未指定。
doc_string string 该模型的文档。
input ValueInfo[] 计算图的输入参数,可能由' initializer '中的默认值初始化。
output ValueInfo[] 计算图的输出参数。一旦执行写入了所有输出参数,计算图的执行就完成了。
value_info ValueInfo[] 用于存储不是输入或输出的值的类型和形状信息。

原来 ValueInfo 和 Tensor 是需要从其内容信息上区分的, ValueInfo 是躯壳的话, Tensor 是具体的模型参数值。

Node

Name Type Description
name string 节点的可选名称,仅用于调试目的。
input string[] 节点用来将输入值传播到节点运算符的值的名称。 它必须引用图形输入、Initializer 或其他节点输出。
output string[] 节点用于从节点调用的运算符捕获数据的输出的名称。它要么在图中引入一个值,要么引用一个图输出。
op_type string 要调用的算子的符号标识符。
domain string 算子集的命名空间,其中包含以op_type命名的操作符。
attribute Attribute[] 命名属性,算子参数化的一种形式,用于常量值而不是传播值。
doc_string string 该Node的文档。
  • 计算图中的边(Edge)由一个节点的输出在后续节点的输入中按名称引用来建立。
  • 给定节点的输出将新名称引入图中。节点输出的值由节点算子计算所得,节点输入可以指其他节点输出、图输入和图初始化器(Initializer)。当节点输出的名称与图输出的名称重合时,图输出的值就是该节点计算出的对应输出值。嵌套子图中的节点输入可以引用外部图中的名称(作为节点输出、图输入或图初始值设定项)。
  • 该图必须对所有节点输出使用单一静态分配,这意味着所有节点输出名称在一个图中必须是唯一的。 在嵌套子图的情况下,节点输出名称必须不同于嵌套子图中可见的外部范围的名称。
  • 节点间依赖性使得不能在计算图中创建循环。
  • 节点中输入和输出的数量、它们的类型、节点中指定的属性集及其类型必须满足节点操作员签名施加的约束。
  • 定义顶层计算图的节点列表必须按拓扑排序; 也就是说,如果节点 K 在图中跟随节点 N,则 N 的任何数据输入都不能引用 K 的输出。
  • 节点属性用于将静态值传递给运算符。

更多关于数据结构的信息可以参考 standard-data-types

ONNX支持的功能

基于ONNX模型,官方提供了一系列相关工具:模型转化/模型优化(simplifier等)/模型部署(Runtime)/模型可视化(Netron等)等

更多工具可以参考ONNX 支持的工具

1、Netron

上述图即为Netron打开

2、模型导出

torch.onnx.export 导出模型

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import torch
import onnx

model_torch = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
model_torch.eval()
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy_input, output_fn="resnet50.onnx",
do_constant_folding=True,
opset_version=14,
input_names=['input'],
output_names=['output'],
keep_initializers_as_inputs=True,
export_params=True,
dynamic_axes = {
'input':{0: 'batch_size'},
'output':{0: 'batch_size'}
},
verbose=True)
# Load the ONNX model
model_onnx = onnx.load("resnet50_torch.onnx")
# Check that the model is well formed
onnx.checker.check_model(model_onnx)
# Print a human readable representation of the graph
# print(onnx.helper.printable_graph(model_onnx.graph))

Scikit-learn 模型导出 onnx 模型

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
# Convert into ONNX format
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType


iris = load_iris()
X, y = iris.data, iris.target
X_train, X_test, y_train, y_test = train_test_split(X, y)
clr = RandomForestClassifier()
clr.fit(X_train, y_train)

initial_type = [('float_input', FloatTensorType([None, 4]))]
onx = convert_sklearn(clr, initial_types=initial_type)
with open("rf_iris.onnx", "wb") as f:
f.write(onx.SerializeToString())

3、onnx simplifier

1
2
3
4
5
6
7
8
9
10
11
import onnx
import onnxsim

model_onnxdo = onnx.load("model_name.onnx")
# useless output, we can cutoff when simplify
unused_output = ["onnx::Sigmoid_467", "onnx::Sigmoid_753", "onnx::Sigmoid_1039"]
model_onnx, check = onnxsim.simplify(model_onnx,
dynamic_input_shape=True,
input_shapes={'input': [1, 3, image_size, image_size]},
unused_output=unused_output)
assert check, 'assert simplification check failed'

simplify的基本流程如下:

step 1. 利用onnxruntime推理计算图,得到各个节点的输入输出的 infer shape step 2. 基于ONNX支持的优化方法进行ONNX模型的优化(如fuse_bn_into_conv) step 3. 对ONNX模型的常量OP进行折叠: 1.基于get_constant_nodes获取常量 OP 2.基于add_features_to_output将所有静态节点的输出扩展到 ONNX 图的输出节点列表中(主要为了后续步骤方便获取常量节点输出) 3.将1.中得到的常量OP从图中移除(断开连线),同时将其节点参数构建为其他节点的输入参数 4.清理图中的孤立节点(3.中断开连线的节点)

其中 v0.3.10 onnxsim/onnx_simplifier.py simplify 函数实现如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
def simplify(model: Union[str, onnx.ModelProto],    # onnx ModelProto object or file path
check_n: int = 0, # The simplified model will be checked for `check_n` times by random inputs
perform_optimization: bool = True, # Whether to run onnx optimizer on the model
skip_fuse_bn: bool = False, # Skip fuse_bn_into_conv onnx optimizer
input_shapes: Optional[TensorShapesWithOptionalKey] = None, #
skipped_optimizers: Optional[Sequence[str]] = None, # Skip some specific onnx optimizers
skip_shape_inference=False, # Skip shape inference (sometimes shape inference will crash)
input_data: Optional[Tensors] = None, # Feed custom input data for checking if needed
dynamic_input_shape: bool = False, # Indicates whether the input shape should be dynamic.
custom_lib: Optional[str] = None, # onnxruntime custom ops's shared library
include_subgraph: bool = False, # Simplify subgraph (e.g. true graph and false graph of "If" operator) instead of only the main graph
unused_output: Optional[Sequence[str]] = None) -> Tuple[onnx.ModelProto, bool]: # name of unused outputs that will be eliminated from the model
"""
:param input_shapes: If the model has dynamic input shape, user must pass a fixed input shape for generating random inputs and checking equality.
(Also see "dynamic_input_shape" param)
:param dynamic_input_shape: input_shapes is also needed even if dynamic_input_shape is True,
the value of input_shapes will be used when generating random inputs for checking equality.
If 'dynamic_input_shape' is False, the input shape in simplified model will be overwritten
by the value of 'input_shapes' param.
:return: A tuple (simplified model, success(True) or failed(False))
"""
if input_shapes is None:
input_shapes = {}
if input_data is None:
input_data = {}

if isinstance(model, str):
model = onnx.load(model)
assert(isinstance(model, onnx.ModelProto))
onnx.checker.check_model(model)
model_ori = model
model = copy.deepcopy(model)

input_names = get_input_names(model)
for input_name, data in input_data.items():
if input_name not in input_names:
raise RuntimeError(
'The model doesn\'t have input named "{}"'.format(input_name))

shape = list(input_data[input_name].shape)

# special case for single constant variables (with shape [])
if len(shape) == 0:
shape = [input_data[input_name].size]
if input_name in input_shapes and shape != input_shapes[input_name]:
raise RuntimeError('The shape of input_data[{}] is not the same with input_shape[{}]'.format(
input_name, input_name))
elif input_name not in input_shapes:
input_shapes[input_name] = shape

if unused_output is not None:
model = remove_unused_output(model, unused_output)

updated_input_shapes = check_and_update_input_shapes(
model, input_shapes, dynamic_input_shape)

def infer_shapes_and_optimize(model: onnx.ModelProto) -> onnx.ModelProto:
def infer_shapes_if_applicable(model: onnx.ModelProto) -> onnx.ModelProto:
"""将Inference Shape信息添加到模型的ValeInfo中"""
if not skip_shape_inference:
model = infer_shapes(model)
return model

def optimize_if_applicable(model: onnx.ModelProto) -> onnx.ModelProto:
if perform_optimization:
model = optimize(model, skip_fuse_bn, skipped_optimizers)
return model

return fixed_point(model, infer_shapes_if_applicable, optimize_if_applicable)

def constant_folding(model: onnx.ModelProto) -> onnx.ModelProto:
const_nodes = get_constant_nodes(
model, dynamic_input_shape=dynamic_input_shape)
res = forward_for_node_outputs(model,
const_nodes,
input_shapes=updated_input_shapes,
input_data=input_data,
custom_lib=custom_lib)
const_nodes = clean_constant_nodes(const_nodes, res)
model = eliminate_const_nodes(model, const_nodes, res)
onnx.checker.check_model(model)
return model

model = fixed_point(model, constant_folding, infer_shapes_and_optimize)

check_ok = check(model_ori, model, check_n,
input_shapes=updated_input_shapes, custom_lib=custom_lib)

return model, check_ok

def optimize(model: onnx.ModelProto, skip_fuse_bn: bool, skipped_optimizers: Optional[Sequence[str]]) -> onnx.ModelProto:
"""
:param model: The onnx model.
:return: The optimized onnx model.
Before simplifying, use this method to generate value_info, which is used in `forward_all`
After simplifying, use this method to fold constants generated in previous step into initializer,
and eliminate unused constants.
"""

onnx.checker.check_model(model)
onnx.helper.strip_doc_string(model)
optimizers_list = onnxoptimizer.get_fuse_and_elimination_passes()
if skip_fuse_bn:
optimizers_list.remove('fuse_bn_into_conv')
if skipped_optimizers is not None:
for opt in skipped_optimizers:
try:
optimizers_list.remove(opt)
except ValueError:
pass

model = onnxoptimizer.optimize(model, optimizers_list,
fixed_point=True)
onnx.checker.check_model(model)
return model

4、onnx 模型编辑 onnx-graphsurgeon

TensorRT 官方提供了关于 onnx-graphsurgeon 操作 ONNX 模型的 example ,例如剥离子图移除Node子图替换等等。

onnx-graphsurgeon-sample

例如我们将原来的 step 这个节点进行修改,修改成直接传入已经构造好的 postion_ids 的方式:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
import onnx
import onnx_graphsurgeon as gs
import numpy as np

def edit_onnx_sample(onnx_path):
graph = gs.import_onnx(onnx.load(onnx_path))
tensors = graph.tensors()
# replace subgraph
shape_node = [node for node in graph.nodes if node.name == "Mul_101"][0]
inputs_mask = tensors["step"]
inputs_mask.name = "position_ids"
step_matrix = inputs_mask.to_variable(dtype=np.float32, shape= [1, 1, 512])
add_node = [node for node in graph.nodes if node.name == "Add_114"][0]
add_node.inputs = [shape_node.outputs[0], step_matrix]
graph.cleanup()
onnx.save(gs.export_onnx(graph), onnx_path.replace(".onnx", "-replace.onnx"))

if __name__ == "__main__":
edit_onnx_sample(path="decoder.onnx")

支持的算子

ONNX 官网 Operators 文档中列出所有 ONNX 算子。 对于每个算子,列出使用指南、参数、示例和逐行版本历史记录。 例如 Abs 是 Opset version 1 开始支持的,其中在版本号为 6 和 13 的时候进行过相关内容的变更, 要查看相信变更情况,可以进入到具体的 diff 页面进行查看;同样,Acos 是从版本 7 开始支持的。

operator versions differences
Abs 13, 6, 1 13/6, 13/1, 6/1
Acos 7
Acosh 9

相关链接