onnx_create_simple_mul_model_from_scratch_and_save.py

1import onnx
2from onnx import helper
3from onnx import AttributeProto, TensorProto, GraphProto
4
5# The NodeProto for the 'Mul' operation
6node_def = helper.make_node(
7    'Mul',                  # node name
8    ['X', 'W'],             # inputs
9    ['Y'],                  # outputs
10)
11
12# The ValueInfoProto for inputs
13X = helper.make_tensor_value_info('X', TensorProto.FLOAT, [3, 2])
14W = helper.make_tensor_value_info('W', TensorProto.FLOAT, [3, 2])
15
16# The ValueInfoProto for outputs
17Y = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [3, 2])
18
19# The GraphProto
20graph_def = helper.make_graph(
21    [node_def],
22    'test-model',
23    [X, W],
24    [Y],
25)
26
27# The ModelProto
28model_def = helper.make_model(graph_def, producer_name='onnx-example')
29
30print('The model is:\n{}'.format(model_def))
31onnx.checker.check_model(model_def)
32print('The model is checked!')
33
34# Save the model
35onnx.save(model_def, 'example.onnx')