【精华】ONNX模型API构造

• • • • ONNX模型API构造与代码检查
      • • 1 构造描述张量信息的对象`ValueInfoProto`
        • 2 构造算子节点信息`NodeProto`
        • 3 构造计算图`GraphProto`
        • 4 封装计算图
        • 5 检查代码
      • ONNX Python API 构造模型完整代码

参考博客:https://zhuanlan.zhihu.com/p/516920606

import onnx 
from onnx import helper 
from onnx import TensorProto 
 
a = helper.make_tensor_value_info('a', TensorProto.FLOAT, [10, 10]) 
x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10, 10]) 
b = helper.make_tensor_value_info('b', TensorProto.FLOAT, [10, 10]) 
output = helper.make_tensor_value_info('output', TensorProto.FLOAT, [10, 10]) 

mul = helper.make_node('Mul', ['a', 'x'], ['c']) 
add = helper.make_node('Add', ['c', 'b'], ['output']) 

graph = helper.make_graph([mul, add], 'linear_func', [a, x, b], [output]) 

用 helper.make_model 把计算图 GraphProto 封装进模型 ModelProto

model = helper.make_model(graph) 

onnx.checker.check_model(model) 
print(model) 
onnx.save(model, 'linear_func.onnx') 

import onnx 
from onnx import helper 
from onnx import TensorProto 
 
# input and output 
a = helper.make_tensor_value_info('a', TensorProto.FLOAT, [10, 10]) 
x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10, 10]) 
b = helper.make_tensor_value_info('b', TensorProto.FLOAT, [10, 10]) 
output = helper.make_tensor_value_info('output', TensorProto.FLOAT, [10, 10]) 
 
# Mul 
mul = helper.make_node('Mul', ['a', 'x'], ['c']) 
 
# Add 
add = helper.make_node('Add', ['c', 'b'], ['output']) 
 
# graph and model 
graph = helper.make_graph([mul, add], 'linear_func', [a, x, b], [output]) 
model = helper.make_model(graph) 
 
# save model 
onnx.checker.check_model(model) 
print(model) 
onnx.save(model, 'linear_func.onnx')