ONNX 概述

参考：ONNX with 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/tmp/cache/conda/envs/tvmz/lib/python3.10/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

Arguments:
    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/tmp/cache/conda/envs/tvmz/lib/python3.10/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: 9
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: 19
}

备注

空的形状（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 中的每个对象（参见 Protos）都可以用 SerializeToString 方法序列化。

with open("linear_regression.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

反序列化，加载序列化的模型：

import onnx

with open("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("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("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 语义。下一个示例修改了前一个示例，将输入 A 和 B 更改为初始化式。（参见 array）。

• onnx.numpy_helper.to_array: 将 ONNX 转换为 NumPy 数组。

• onnx.numpy_helper.from_array: 将N umPy 数组转换为 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>"

节点属性

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: 9
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: 19
}

评估与运行时

完整 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([[0.5315706 ],
       [2.439557  ],
       [0.55338883],
       [0.8112479 ]], dtype=float32)]

评估节点

评估器还可以评估简单的节点，以检查算子在特定输入上的行为。

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 上工作。

逐步评估

转换库接受使用机器学习框架（如 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([[-5.9394155],
       [ 0.3874052],
       [-0.8096788],
       [-1.4719555]], dtype=float32)]

------ verbose=2

MatMul(X, A) -> XA
Add(XA, B) -> Y
[array([[ 3.3348103 ],
       [ 1.0723081 ],
       [-0.3836289 ],
       [ 0.30403274]], dtype=float32)]

------ verbose=3

 +I X: float32:(4, 2) in [-2.0797390937805176, 1.9476134777069092]
 +I A: float32:(2, 1) in [-0.00496311392635107, 0.2595963776111603]
 +I B: float32:(1, 1) in [1.089706540107727, 1.089706540107727]
MatMul(X, A) -> XA
 + XA: float32:(4, 1) in [-0.5365331172943115, -0.1079740896821022]
Add(XA, B) -> Y
 + Y: float32:(4, 1) in [0.5531734228134155, 0.9817324280738831]
[array([[0.5531734 ],
       [0.9525344 ],
       [0.9817324 ],
       [0.72496355]], dtype=float32)]

------ verbose=4

 +I X: float32:(4, 2):1.3000822067260742,-0.7955483198165894,0.6271983981132507,0.798268735408783,0.1631830334663391...
 +I A: float32:(2, 1):[0.5944057106971741, -1.264738917350769]
 +I B: float32:(1, 1):[0.09203124046325684]
MatMul(X, A) -> XA
 + XA: float32:(4, 1):[1.7789373397827148, -0.6367912888526917, -2.1531877517700195, -1.0709211826324463]
Add(XA, B) -> Y
 + Y: float32:(4, 1):[1.8709685802459717, -0.5447600483894348, -2.0611565113067627, -0.9788899421691895]
[array([[ 1.8709686 ],
       [-0.54476005],
       [-2.0611565 ],
       [-0.97888994]], dtype=float32)]

评估自定义节点

以下示例仍然实现了线性回归，但在 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.3906131, -1.8486918],
       [-1.4727641,  0.1704061],
       [-0.6817981, -1.8363199],
       [-1.1007469, -1.8389864]], dtype=float32)]