ONNX 概述#
源代码：ONNX & ONNX with Python
备注
ONNX 是一种开放格式，旨在表示机器学习模型。ONNX 定义了一组通用的算子——机器学习和深度学习模型的基本构建块——以及一种通用的文件格式，以便 AI 开发人员能够使用多种框架、工具、运行时和编译器来处理模型。
ONNX Python 接口#
参考：onnx python
ONNX 是强类型的。必须为函数的输入和输出定义形状和类型。下面介绍一些常用的 make function：
make_tensor_value_info：声明给定形状和类型的变量（输入或输出）。make_node: 创建由运算（算子类型）、输入和输出定义的节点。make_graph: 用于使用前两个函数创建的 ONNX graph 对象。make_model: 将 graph 和附加元数据合并在一起。
在整个创建过程中，需要为 graph 中每个节点的每个输入和输出命名。graph 的输入和输出由 onnx 对象定义，字符串用于引用中间结果。
from onnx import TensorProto
from onnx.helper import (
make_model, make_node, make_graph,
make_tensor_value_info
)
# from onnx.checker import check_model
创建输入变量：
make_tensor_value_info?
Signature:
make_tensor_value_info(
name: str,
elem_type: int,
shape: Optional[Sequence[Union[str, int, NoneType]]],
doc_string: str = '',
shape_denotation: Optional[List[str]] = None,
) -> onnx.onnx_ml_pb2.ValueInfoProto
Docstring: Makes a ValueInfoProto based on the data type and shape.
File: /media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/onnx/helper.py
Type: function
X = make_tensor_value_info('X', TensorProto.FLOAT, [None, None])
A = make_tensor_value_info('A', TensorProto.FLOAT, [None, None])
B = make_tensor_value_info('B', TensorProto.FLOAT, [None, None])
查看 X：
X
name: "X"
type {
tensor_type {
elem_type: 1
shape {
dim {
}
dim {
}
}
}
}
创建输出变量：
Y = make_tensor_value_info('Y', TensorProto.FLOAT, [None])
将上述变量组织为计算图的节点：
node1 = make_node('MatMul', ['X', 'A'], ['XA'])
node2 = make_node('Add', ['XA', 'B'], ['Y'])
查看 node1：
node1
input: "X"
input: "A"
output: "XA"
op_type: "MatMul"
将节点组织为计算图：
make_graph?
Signature:
make_graph(
nodes: Sequence[onnx.onnx_ml_pb2.NodeProto],
name: str,
inputs: Sequence[onnx.onnx_ml_pb2.ValueInfoProto],
outputs: Sequence[onnx.onnx_ml_pb2.ValueInfoProto],
initializer: Optional[Sequence[onnx.onnx_ml_pb2.TensorProto]] = None,
doc_string: Optional[str] = None,
value_info: Optional[Sequence[onnx.onnx_ml_pb2.ValueInfoProto]] = None,
sparse_initializer: Optional[Sequence[onnx.onnx_ml_pb2.SparseTensorProto]] = None,
) -> onnx.onnx_ml_pb2.GraphProto
Docstring:
Construct a GraphProto
Args:
nodes: list of NodeProto
name (string): graph name
inputs: list of ValueInfoProto
outputs: list of ValueInfoProto
initializer: list of TensorProto
doc_string (string): graph documentation
value_info: list of ValueInfoProto
sparse_initializer: list of SparseTensorProto
Returns:
GraphProto
File: /media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/onnx/helper.py
Type: function
graph = make_graph(
[node1, node2], # nodes
'lr', # a name
[X, A, B], # inputs
[Y] # outputs
)
graph
node {
input: "X"
input: "A"
output: "XA"
op_type: "MatMul"
}
node {
input: "XA"
input: "B"
output: "Y"
op_type: "Add"
}
name: "lr"
input {
name: "X"
type {
tensor_type {
elem_type: 1
shape {
dim {
}
dim {
}
}
}
}
}
input {
name: "A"
type {
tensor_type {
elem_type: 1
shape {
dim {
}
dim {
}
}
}
}
}
input {
name: "B"
type {
tensor_type {
elem_type: 1
shape {
dim {
}
dim {
}
}
}
}
}
output {
name: "Y"
type {
tensor_type {
elem_type: 1
shape {
dim {
}
}
}
}
}
将计算图变换为模型：
onnx_model = make_model(graph)
onnx_model
ir_version: 10
graph {
node {
input: "X"
input: "A"
output: "XA"
op_type: "MatMul"
}
node {
input: "XA"
input: "B"
output: "Y"
op_type: "Add"
}
name: "lr"
input {
name: "X"
type {
tensor_type {
elem_type: 1
shape {
dim {
}
dim {
}
}
}
}
}
input {
name: "A"
type {
tensor_type {
elem_type: 1
shape {
dim {
}
dim {
}
}
}
}
}
input {
name: "B"
type {
tensor_type {
elem_type: 1
shape {
dim {
}
dim {
}
}
}
}
}
output {
name: "Y"
type {
tensor_type {
elem_type: 1
shape {
dim {
}
}
}
}
}
}
opset_import {
version: 21
}
备注
空的形状（None）意味着任何形状。
访问 ONNX graph#
ONNX graph 也可以通过查看计算图中每个对象的字段来检查。
查看输入列表：
print(onnx_model.graph.input)
[name: "X"
type {
tensor_type {
elem_type: 1
shape {
dim {
}
dim {
}
}
}
}
, name: "A"
type {
tensor_type {
elem_type: 1
shape {
dim {
}
dim {
}
}
}
}
, name: "B"
type {
tensor_type {
elem_type: 1
shape {
dim {
}
dim {
}
}
}
}
]
更优雅的打印输入信息：
def shape2tuple(shape):
return tuple(getattr(d, 'dim_value', 0) for d in shape.dim)
for obj in onnx_model.graph.input:
print("name=%r dtype=%r shape=%r" % (
obj.name, obj.type.tensor_type.elem_type,
shape2tuple(obj.type.tensor_type.shape)))
name='X' dtype=1 shape=(0, 0)
name='A' dtype=1 shape=(0, 0)
name='B' dtype=1 shape=(0, 0)
同样可以查看输出信息：
for obj in onnx_model.graph.output:
print("name=%r dtype=%r shape=%r" % (
obj.name, obj.type.tensor_type.elem_type,
shape2tuple(obj.type.tensor_type.shape)))
name='Y' dtype=1 shape=(0,)
查看节点信息：
for node in onnx_model.graph.node:
print("name=%r type=%r input=%r output=%r" % (
node.name, node.op_type, node.input, node.output))
name='' type='MatMul' input=['X', 'A'] output=['XA']
name='' type='Add' input=['XA', 'B'] output=['Y']
ONNX 序列化与反序列化#
onnx 中的每个对象（参见 Protos）都可以用 SerializeToString 方法序列化。
!mkdir -p .temp
with open(".temp/linear_regression.onnx", "wb") as f:
f.write(onnx_model.SerializeToString())
反序列化，加载序列化的模型：
import onnx
with open(".temp/linear_regression.onnx", "rb") as f:
onnx_model = onnx.load(f)
数据也可以序列化：
import numpy as np
from onnx.numpy_helper import from_array
numpy_tensor = np.array([0, 1, 4, 5, 3], dtype=np.float32)
print(type(numpy_tensor))
onnx_tensor = from_array(numpy_tensor)
print(type(onnx_tensor))
serialized_tensor = onnx_tensor.SerializeToString()
print(type(serialized_tensor))
with open(".temp/saved_tensor.pb", "wb") as f:
f.write(serialized_tensor)
<class 'numpy.ndarray'>
<class 'onnx.onnx_ml_pb2.TensorProto'>
<class 'bytes'>
反序列化数据：
from onnx import TensorProto
from onnx.numpy_helper import to_array
with open(".temp/saved_tensor.pb", "rb") as f:
serialized_tensor = f.read()
print(type(serialized_tensor))
onnx_tensor = TensorProto()
onnx_tensor.ParseFromString(serialized_tensor)
print(type(onnx_tensor))
numpy_tensor = to_array(onnx_tensor)
print(numpy_tensor)
<class 'bytes'>
<class 'onnx.onnx_ml_pb2.TensorProto'>
[0. 1. 4. 5. 3.]
也可以使用便捷函数 load_tensor_from_string：
from onnx import load_tensor_from_string
with open("saved_tensor.pb", "rb") as f:
serialized = f.read()
proto = load_tensor_from_string(serialized)
print(type(proto))
<class 'onnx.onnx_ml_pb2.TensorProto'>
ONNX 初始化器与默认值#
之前的模型假设线性回归的系数也是模型的输入。那不太方便。它们应该作为常量或初始化项成为模型本身的一部分，以遵循 onnx 语义。下一个示例修改了前一个示例，将输入 A 和 B 更改为初始化式。（参见 array）。
onnx.numpy_helper.to_array: 将 ONNX 转换为 NumPy 数组。onnx.numpy_helper.from_array: 将 NumPy 数组转换为 ONNX。
import numpy as np
from onnx import numpy_helper, TensorProto
from onnx.helper import (
make_model, make_node, make_graph,
make_tensor_value_info)
from onnx.checker import check_model
# initializers
value = np.array([0.5, -0.6], dtype=np.float32)
A = numpy_helper.from_array(value, name='A')
value = np.array([0.4], dtype=np.float32)
C = numpy_helper.from_array(value, name='C')
X = make_tensor_value_info('X', TensorProto.FLOAT, [None, None])
Y = make_tensor_value_info('Y', TensorProto.FLOAT, [None])
node1 = make_node('MatMul', ['X', 'A'], ['AX'])
node2 = make_node('Add', ['AX', 'C'], ['Y'])
graph = make_graph([node1, node2], 'lr', [X], [Y], [A, C])
onnx_model = make_model(graph)
check_model(onnx_model)
查看初始化值：
for init in onnx_model.graph.initializer:
print(init)
dims: 2
data_type: 1
name: "A"
raw_data: "\000\000\000?\232\231\031\277"
dims: 1
data_type: 1
name: "C"
raw_data: "\315\314\314>"
ONNX 节点属性#
from onnx import TensorProto
from onnx.helper import (
make_model, make_node, make_graph,
make_tensor_value_info)
from onnx.checker import check_model
# unchanged
X = make_tensor_value_info('X', TensorProto.FLOAT, [None, None])
A = make_tensor_value_info('A', TensorProto.FLOAT, [None, None])
B = make_tensor_value_info('B', TensorProto.FLOAT, [None, None])
Y = make_tensor_value_info('Y', TensorProto.FLOAT, [None])
# 添加属性
node_transpose = make_node('Transpose', ['A'], ['tA'], perm=[1, 0])
# unchanged except A is replaced by tA
node1 = make_node('MatMul', ['X', 'tA'], ['XA'])
node2 = make_node('Add', ['XA', 'B'], ['Y'])
# node_transpose is added to the list
graph = make_graph([node_transpose, node1, node2],
'lr', [X, A, B], [Y])
onnx_model = make_model(graph)
check_model(onnx_model)
# the work is done, let's display it...
print(onnx_model)
ir_version: 10
graph {
node {
input: "A"
output: "tA"
op_type: "Transpose"
attribute {
name: "perm"
ints: 1
ints: 0
type: INTS
}
}
node {
input: "X"
input: "tA"
output: "XA"
op_type: "MatMul"
}
node {
input: "XA"
input: "B"
output: "Y"
op_type: "Add"
}
name: "lr"
input {
name: "X"
type {
tensor_type {
elem_type: 1
shape {
dim {
}
dim {
}
}
}
}
}
input {
name: "A"
type {
tensor_type {
elem_type: 1
shape {
dim {
}
dim {
}
}
}
}
}
input {
name: "B"
type {
tensor_type {
elem_type: 1
shape {
dim {
}
dim {
}
}
}
}
}
output {
name: "Y"
type {
tensor_type {
elem_type: 1
shape {
dim {
}
}
}
}
}
}
opset_import {
version: 21
}
ONNX 评估与运行时#
完整 API 的描述见 onnx.reference。它接受一个模型(ModelProto，文件名，…)。方法 run 返回字典中指定的一组给定输入的输出。
import numpy
from onnx import numpy_helper, TensorProto
from onnx.helper import (
make_model, make_node, set_model_props, make_tensor,
make_graph, make_tensor_value_info)
from onnx.checker import check_model
from onnx.reference import ReferenceEvaluator
X = make_tensor_value_info('X', TensorProto.FLOAT, [None, None])
A = make_tensor_value_info('A', TensorProto.FLOAT, [None, None])
B = make_tensor_value_info('B', TensorProto.FLOAT, [None, None])
Y = make_tensor_value_info('Y', TensorProto.FLOAT, [None])
node1 = make_node('MatMul', ['X', 'A'], ['XA'])
node2 = make_node('Add', ['XA', 'B'], ['Y'])
graph = make_graph([node1, node2], 'lr', [X, A, B], [Y])
onnx_model = make_model(graph)
check_model(onnx_model)
sess = ReferenceEvaluator(onnx_model)
x = numpy.random.randn(4, 2).astype(numpy.float32)
a = numpy.random.randn(2, 1).astype(numpy.float32)
b = numpy.random.randn(1, 1).astype(numpy.float32)
feeds = {'X': x, 'A': a, 'B': b}
print(sess.run(None, feeds))
[array([[-1.8057449],
[-2.0268912],
[-1.369731 ],
[-1.8334708]], dtype=float32)]
ONNX 评估节点#
评估器还可以评估简单的节点，以检查算子在特定输入上的行为。
import numpy
from onnx import numpy_helper, TensorProto
from onnx.helper import make_node
from onnx.reference import ReferenceEvaluator
node = make_node('EyeLike', ['X'], ['Y'])
sess = ReferenceEvaluator(node)
x = numpy.random.randn(4, 2).astype(numpy.float32)
feeds = {'X': x}
print(sess.run(None, feeds))
[array([[1., 0.],
[0., 1.],
[0., 0.],
[0., 0.]], dtype=float32)]
类似的代码也可以在 GraphProto 或 FunctionProto 上工作。
ONNX 逐步评估#
转换库接受使用机器学习框架（如 PyTorch、scikit-learn 等）训练的现有模型，并将其转换为 ONNX graph。复杂的模型通常不会在第一次尝试中工作，查看中间结果可能有助于找到未正确转换的部分。参数详细显示有关中间结果的信息。
import numpy
from onnx import numpy_helper, TensorProto
from onnx.helper import (
make_model, make_node, set_model_props, make_tensor,
make_graph, make_tensor_value_info)
from onnx.checker import check_model
from onnx.reference import ReferenceEvaluator
X = make_tensor_value_info('X', TensorProto.FLOAT, [None, None])
A = make_tensor_value_info('A', TensorProto.FLOAT, [None, None])
B = make_tensor_value_info('B', TensorProto.FLOAT, [None, None])
Y = make_tensor_value_info('Y', TensorProto.FLOAT, [None])
node1 = make_node('MatMul', ['X', 'A'], ['XA'])
node2 = make_node('Add', ['XA', 'B'], ['Y'])
graph = make_graph([node1, node2], 'lr', [X, A, B], [Y])
onnx_model = make_model(graph)
check_model(onnx_model)
for verbose in [1, 2, 3, 4]:
print()
print(f"------ verbose={verbose}")
print()
sess = ReferenceEvaluator(onnx_model, verbose=verbose)
x = numpy.random.randn(4, 2).astype(numpy.float32)
a = numpy.random.randn(2, 1).astype(numpy.float32)
b = numpy.random.randn(1, 1).astype(numpy.float32)
feeds = {'X': x, 'A': a, 'B': b}
print(sess.run(None, feeds))
------ verbose=1
[array([[ 0.9199427 ],
[ 0.3668961 ],
[ 0.04145484],
[-0.4361933 ]], dtype=float32)]
------ verbose=2
MatMul(X, A) -> XA
Add(XA, B) -> Y
[array([[3.8724718],
[4.0825634],
[0.4058497],
[2.005159 ]], dtype=float32)]
------ verbose=3
+I X: float32:(4, 2) in [-1.1875110864639282, 1.6281424760818481]
+I A: float32:(2, 1) in [-1.774684190750122, -1.7485564947128296]
+I B: float32:(1, 1) in [0.5859348177909851, 0.5859348177909851]
MatMul(X, A) -> XA
+ XA: float32:(4, 1) in [-2.718858480453491, 3.8624463081359863]
Add(XA, B) -> Y
+ Y: float32:(4, 1) in [-2.1329236030578613, 4.448380947113037]
[array([[ 1.2496312 ],
[ 4.448381 ],
[-2.1329236 ],
[ 0.39690298]], dtype=float32)]
------ verbose=4
+I X: float32:(4, 2):1.1294219493865967,0.30637702345848083,-0.029152901843190193,0.06715787202119827,0.20514409244060516...
+I A: float32:(2, 1):[-0.08160115778446198, 1.6542348861694336]
+I B: float32:(1, 1):[-0.9010018110275269]
MatMul(X, A) -> XA
+ XA: float32:(4, 1):[0.4146574139595032, 0.1134738028049469, 2.200101137161255, -0.3942927122116089]
Add(XA, B) -> Y
+ Y: float32:(4, 1):[-0.4863443970680237, -0.7875280380249023, 1.299099326133728, -1.2952945232391357]
[array([[-0.4863444 ],
[-0.78752804],
[ 1.2990993 ],
[-1.2952945 ]], dtype=float32)]
ONNX 评估自定义节点#
以下示例仍然实现了线性回归，但在 A 中添加了单位矩阵：
\[
Y = X(A + I) + B
\]
import numpy
from onnx import numpy_helper, TensorProto
from onnx.helper import (
make_model, make_node, set_model_props, make_tensor,
make_graph, make_tensor_value_info)
from onnx.checker import check_model
from onnx.reference import ReferenceEvaluator
X = make_tensor_value_info('X', TensorProto.FLOAT, [None, None])
A = make_tensor_value_info('A', TensorProto.FLOAT, [None, None])
B = make_tensor_value_info('B', TensorProto.FLOAT, [None, None])
Y = make_tensor_value_info('Y', TensorProto.FLOAT, [None])
node0 = make_node('EyeLike', ['A'], ['Eye'])
node1 = make_node('Add', ['A', 'Eye'], ['A1'])
node2 = make_node('MatMul', ['X', 'A1'], ['XA1'])
node3 = make_node('Add', ['XA1', 'B'], ['Y'])
graph = make_graph([node0, node1, node2, node3], 'lr', [X, A, B], [Y])
onnx_model = make_model(graph)
check_model(onnx_model)
with open("linear_regression.onnx", "wb") as f:
f.write(onnx_model.SerializeToString())
sess = ReferenceEvaluator(onnx_model, verbose=2)
x = numpy.random.randn(4, 2).astype(numpy.float32)
a = numpy.random.randn(2, 2).astype(numpy.float32) / 10
b = numpy.random.randn(1, 2).astype(numpy.float32)
feeds = {'X': x, 'A': a, 'B': b}
print(sess.run(None, feeds))
EyeLike(A) -> Eye
Add(A, Eye) -> A1
MatMul(X, A1) -> XA1
Add(XA1, B) -> Y
[array([[ 1.5872786 , 2.8542266 ],
[ 0.64192116, -1.225003 ],
[-1.0039632 , 1.1880492 ],
[-0.31344515, 0.10079634]], dtype=float32)]