forked from onnx/onnx-mlir
-
Notifications
You must be signed in to change notification settings - Fork 0
/
gen_add_onnx.py
33 lines (27 loc) · 903 Bytes
/
gen_add_onnx.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import onnx
from onnx import helper
from onnx import AttributeProto, TensorProto, GraphProto
# Create one input (ValueInfoProto)
X1 = helper.make_tensor_value_info("X1", TensorProto.FLOAT, [3, 2])
X2 = helper.make_tensor_value_info("X2", TensorProto.FLOAT, [3, 2])
# Create one output (ValueInfoProto)
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [3, 2])
# Create a node (NodeProto) - This is based on Pad-11
node_def = helper.make_node(
"Add", # node name
["X1", "X2"], # inputs
["Y"], # outputs
)
# Create the graph (GraphProto)
graph_def = helper.make_graph(
[node_def],
"test-model",
[X1, X2],
[Y],
)
# Create the model (ModelProto)
model_def = helper.make_model(graph_def, producer_name="onnx-example")
print("The model is:\n{}".format(model_def))
onnx.checker.check_model(model_def)
onnx.save(model_def, "add.onnx")
print("The model is checked!")