ONNX概念
onnx不仅仅一种模型参数存储的格式,还是一套完整的用于描述计算函数的编程语言,它的作用就是定义计算图,他本身无法进行。
这个概念和Verilog有点像,一个是硬件描述语言,一个是模型描述语言。
onnx模型或者说计算图,是由这几部分组成:
- input,output:输入输出
- node:节点,即算子,算子的固定参数保存在attribute中
- initializer:初始化器,用于加载函数
此外onnx还允许在模型中添加一些元数据,用于记录作者,模型版本等信息,类似于注释一样。onnx模型中的元数据有:
- doc_string:人类可读的文档,可以用markdown
- domain:不知道干啥用的,反正存个模型名字
- metadata_props:是个字典类型的,不知道干啥使的
- model_author:模型的作者
- model_license:模型的版权协议
- model_version:模型版本
- producer_name:训练模型的框架
- producer_version:训练框架的版本
- training_info:训练信息
Onnx存储方式是使用protobuf来存储,protobuf是Protocol Buffers的简称,它是Google公司开发的一种数据描述语言,很适合做数据存储或 RPC 数据交换格式。可用于通讯协议、数据存储等领域的语言无关、平台无关、可扩展的序列化结构数据格式。它的定位类似于xml和json这种。
Onnx支持的operator
Onnx将支持的operator存储到两个domain中,一个domain是ai.onnx,存储了大量深度学习常用operator,另一个是ai.onnx.ml,存储了树模型以及一些机器学习中的operator。
Onnx支持的数据类型
基本上主流的数据类型都支持,例如fp32,fp16,int4,int8,int16等.Onnx中支持的数据类型:
0: onnx.TensorProto.UNDEFINED
1: onnx.TensorProto.FLOAT
2: onnx.TensorProto.UINT8
3: onnx.TensorProto.INT8
4: onnx.TensorProto.UINT16
5: onnx.TensorProto.INT16
6: onnx.TensorProto.INT32
7: onnx.TensorProto.INT64
8: onnx.TensorProto.STRING
9: onnx.TensorProto.BOOL
10: onnx.TensorProto.FLOAT16
11: onnx.TensorProto.DOUBLE
12: onnx.TensorProto.UINT32
13: onnx.TensorProto.UINT64
14: onnx.TensorProto.COMPLEX64
15: onnx.TensorProto.COMPLEX128
16: onnx.TensorProto.BFLOAT16
17: onnx.TensorProto.FLOAT8E4M3FN
18: onnx.TensorProto.FLOAT8E4M3FNUZ
19: onnx.TensorProto.FLOAT8E5M2
20: onnx.TensorProto.FLOAT8E5M2FNUZ
21: onnx.TensorProto.UINT4
22: onnx.TensorProto.INT4
ONNX是强类型的语言,不支持数据类型之间的隐式转换,想要转换必须添加显式转换。
ONNX支持2维的稀疏张量,类型为SparseTensorProto
稀疏向量是指0比较多的向量,可以用特殊的方式来存储以减少空间占用。
opset version
ONNX的算子库是不断更新的,opset version就是算子库版本,onnx版本和opset的版本是对应的,每次算子库版本更新都会引入新的算子。模型本身也要指定一个算子版本来表示模型所依赖的算子的版本。例如6,7,13,14版本的算子库对Add算子进行了更新,如果模型指定的算子库版本是15,那么Add算子将使用14版本的实现。
Onnx控制流
ONNX支持控制流算子例如if,但是这样会降低性能,最好避免控制流算子。
if算子
根据条件决定执行哪个子图,但是子图的输出shape和num必须是一样的,子图的输出将作为if算子的输出。
下面是一个简单的模型的搭建过程,手搓模型就是这么搓的。
import onnx
import numpy as np
# Given a bool scalar input cond.
# return constant tensor x if cond is True, otherwise return constant tensor y.
cond = onnx.helper.make_tensor_value_info( # 创建输入
"cond", onnx.TensorProto.BOOL, []
)
then_out = onnx.helper.make_tensor_value_info( # 创建then输出
"then_out", onnx.TensorProto.FLOAT, [5]
)
else_out = onnx.helper.make_tensor_value_info( # 创建else输出
"else_out", onnx.TensorProto.FLOAT, [5]
)
x = np.array([1, 2, 3, 4, 5]).astype(np.float32) # 创建then输出的值
y = np.array([5, 4, 3, 2, 1]).astype(np.float32) # 创建else输出的值
then_const_node = onnx.helper.make_node( # 创建then输出的节点
"Constant",
inputs=[],
outputs=["then_out"],
value=onnx.numpy_helper.from_array(x),
)
else_const_node = onnx.helper.make_node( # 创建else输出的节点
"Constant",
inputs=[],
outputs=["else_out"],
value=onnx.numpy_helper.from_array(y),
)
then_body = onnx.helper.make_graph( # 创建then的子图
[then_const_node], "then_body", [], [then_out]
)
else_body = onnx.helper.make_graph( # 创建else的子图
[else_const_node], "else_body", [], [else_out]
)
if_node = onnx.helper.make_node( # 创建if节点
"If",
inputs=["cond"],
outputs=["res"],
then_branch=then_body,
else_branch=else_body,
)
res = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, [5]) # 创建输出,这个输出是if节点的输出
graph = onnx.helper.make_graph( # 创建主图
[if_node], "test_if", [cond], [res]
)
onnx.save_model( # 保存模型
onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 11)]),
"if.onnx",
)
模型结构图
Scan循环算子
scan算子有一个全局变量,只需要开始时输入进行初始化,还有一个输入,每次循环都输入,这两个是并在一个输入里面的,scan的输出也是,一个输出是全局变量最后的值,另一个输出是每次循环的输出结合到一起,比如:全局变量是N维度,每次输入是M维,每次输出也是M维,循环k次,那么scan的输入就是:[N,kM],输出是:[N,kM]。
# Given an input sequence [x1, ..., xN], sum up its elements using a scan
# returning the final state (x1+x2+...+xN) as well the scan_output
# [x1, x1+x2, ..., x1+x2+...+xN]
#
# create graph to represent scan body
import numpy as np
import onnx
sum_in = onnx.helper.make_tensor_value_info(
"sum_in", onnx.TensorProto.FLOAT, [2]
)
next = onnx.helper.make_tensor_value_info("next", onnx.TensorProto.FLOAT, [2])
sum_out = onnx.helper.make_tensor_value_info(
"sum_out", onnx.TensorProto.FLOAT, [2]
)
scan_out = onnx.helper.make_tensor_value_info(
"scan_out", onnx.TensorProto.FLOAT, [2]
)
add_node = onnx.helper.make_node(
"Add", inputs=["sum_in", "next"], outputs=["sum_out"]
)
id_node = onnx.helper.make_node(
"Identity", inputs=["sum_out"], outputs=["scan_out"]
)
scan_body = onnx.helper.make_graph(
[add_node, id_node], "scan_body", [sum_in, next], [sum_out, scan_out]
)
# create scan op node
node = onnx.helper.make_node(
"Scan",
inputs=["initial", "x"],
outputs=["y", "z"],
num_scan_inputs=1,
body=scan_body,
)
# create inputs for sequence-length 3, inner dimension 2
initial = np.array([0, 0]).astype(np.float32).reshape((2,))
x = np.array([1, 2, 3, 4, 5, 6]).astype(np.float32).reshape((3, 2))
# final state computed = [1 + 3 + 5, 2 + 4 + 6]
y = np.array([9, 12]).astype(np.float32).reshape((2,))
# scan-output computed
z = np.array([1, 2, 4, 6, 9, 12]).astype(np.float32).reshape((3, 2))
# create graph
initial_info = onnx.helper.make_tensor_value_info(
"initial", onnx.TensorProto.FLOAT, initial.shape
)
x_info = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, x.shape)
y_info = onnx.helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, y.shape)
z_info = onnx.helper.make_tensor_value_info("z", onnx.TensorProto.FLOAT, z.shape)
graph = onnx.helper.make_graph(
[node], "test_scan", [initial_info, x_info], [y_info, z_info]
)
# create model
model = onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 11)])
# save model
onnx.save_model(model, "scan.onnx")
#inference
import onnxruntime as rt
sess = rt.InferenceSession("scan.onnx")
res = sess.run(None, {"initial": initial, "x": x})
print(res)
LOOP算子
就常见的循环,跟scan差不多,不过没有全局变量,输入也是每次循环输入一个,输出是每个循环的输出结合到一起,结合方式有两种:一种是结合成一个大的tensor,另一种是多个tensor结合成一个sequence,前者要求每个循环的输出的shape必须可以相互兼容才能结合。
扩展算子
ONNX允许自定义算子,这部分内容比较多,后面再写blog专门讲
Function
我的理解是:有的模型中的层是多个算子结合到一起形成的,比如yolo中的C3,这种不需要再自定义算子,可以把几个需要的算子结合到一起形成一个Function。官方说这样做的好处是:可以减小代码量,可以给推理引擎额外信息,推理引擎可以用这些信息做优化,比如为一些Function进行底层实现。
官方文档原文:Functions are one way to extend ONNX specifications. Some model requires the same combination of operators. This can be avoided by creating a function itself defined with existing ONNX operators. Once defined, a function behaves like any other operators. It has inputs, outputs and attributes.
There are two advantages of using functions. The first one is to have a shorter code and easier to read. The second one is that any onnxruntime can leverage that information to run predictions faster. The runtime could have a specific implementation for a function not relying on the implementation of the existing operators
类型推理
Onnx可以推理出模型的输出的数据类型以及大小,但是对于有自定义算子的模型不可以。
工具
- netron:可视化工具,自行百度下载
- onnx2py.py:根据onnx模型反生成一个py文件,这个py脚本可以生成这个模型,用于让用户修改模型,例如想要修改一个模型,可以先生成“可以生成这个模型的脚本”,然后再修改这个脚本,再用这个修改过的脚本生成模型。
- onnx-graphsurgeon:TensorRT做的一个工具,可以用于修改onnx模型,名字翻译过来就是图手术刀
如果觉得有帮助,点赞收藏+关注!thanks!