onnx_helper_api_linear_regression_model_creation.py

1import onnx
2from onnx import helper
3from onnx import TensorProto
4
5# Create one input (ValueInfoProto)
6X = helper.make_tensor_value_info('X', TensorProto.FLOAT, [None, 2])
7
8# Create one output (ValueInfoProto)
9Y = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [None, 1])
10
11# Create weights
12w = helper.make_tensor('W', TensorProto.FLOAT, [2, 1], [1.0, 2.0])
13b = helper.make_tensor('B', TensorProto.FLOAT, [1], [0.5])
14
15# Create a node (NodeProto)
16node_def = helper.make_node(
17    'Gemm', # node name
18    ['X', 'W', 'B'], # inputs
19    ['Y'], # outputs
20)
21
22# Create the graph (GraphProto)
23graph_def = helper.make_graph(
24    [node_def],
25    'test-model',
26    [X],
27    [Y],
28    [w, b],
29)
30
31# Create the model (ModelProto)
32model_def = helper.make_model(graph_def, producer_name='onnx-example')
33
34# Check the model and print its string representation
35onnx.checker.check_model(model_def)
36print(model_def)