def ext_weights():
path = os.path.join(TEST_ROOT, "models", "ext_weights.onnx")
model = onnx.load(path)
inputs = [Variable("input", shape=(1, 3), dtype=np.float32)]
outputs = [Variable("output", shape=(1, 3), dtype=np.float32)]
a = Constant("a", values=np.ones((1, 3), dtype=np.float32))
b = Constant("b", values=np.ones((1, 3), dtype=np.float32))
d = Constant("d", values=np.ones((1, 3), dtype=np.float32))
c = Variable("c")
e = Variable("e")
nodes = [
Node(op="Add", inputs=[a, b], outputs=[c]),
Node(op="Add", inputs=[c, d], outputs=[e]),
Node(op="Add", inputs=[inputs[0], e], outputs=outputs),
]
return Model(
path,
inputs=inputs,
outputs=outputs,
nodes=nodes,
opset=OnnxImporter.get_opset(model),
)
def nested_dup_names():
path = os.path.join(TEST_ROOT, "models", "nested_dup_names.onnx")
model = onnx.load(path)
# Inner
subgraph_inputs = [Variable("X", shape=(2, 2), dtype=np.float32)]
subgraph_outputs = [Variable("Y", shape=(2, 2), dtype=np.float32)]
subgraph_node = Node(op="Identity",
inputs=subgraph_inputs,
outputs=subgraph_outputs)
subgraph = Graph(nodes=[subgraph_node],
inputs=subgraph_inputs,
outputs=subgraph_outputs)
# Outer - problem happens if outer node has same I/O names as subgraph
inputs = [Variable("X", shape=(2, 2), dtype=np.float32)]
outputs = [Variable("Y", shape=(2, 2), dtype=np.float32)]
node = Node(op="Nested",
inputs=inputs,
outputs=outputs,
attrs={"body": subgraph})
return Model(
path,
inputs=inputs,
outputs=outputs,
nodes=[node],
opset=OnnxImporter.get_opset(model),
)
def initializer_is_output_model():
path = os.path.join(TEST_ROOT, "models", "initializer_is_output.onnx")
model = onnx.load(path)
X = Constant(name="X", values=np.ones((64, 64), dtype=np.float32))
return Model(path,
inputs=[],
outputs=[X],
nodes=[],
opset=OnnxImporter.get_opset(model))
def identity_model():
path = os.path.join(TEST_ROOT, "models", "identity.onnx")
model = onnx.load(path)
x = Variable(name="x", dtype=np.float32, shape=(1, 1, 2, 2))
y = Variable(name="y", dtype=np.float32, shape=(1, 1, 2, 2))
node = Node(op="Identity", inputs=[x], outputs=[y])
return Model(path,
inputs=[x],
outputs=[y],
nodes=[node],
opset=OnnxImporter.get_opset(model))
def dim_param_model():
path = os.path.join(TEST_ROOT, "models", "dim_param.onnx")
model = onnx.load(path)
x = Variable(name="Input:0", dtype=np.float32, shape=("dim0", 16, 128))
y = Variable(name="Output:0", dtype=np.float32, shape=("dim0", 16, 128))
node = Node(op="Identity", inputs=[x], outputs=[y])
return Model(path,
inputs=[x],
outputs=[y],
nodes=[node],
opset=OnnxImporter.get_opset(model))
def import_onnx(onnx_model: "onnx.ModelProto") -> Graph:
"""
Import an onnx-graphsurgeon Graph from the provided ONNX model.
Args:
onnx_model (onnx.ModelProto): The ONNX model.
Returns:
Graph: A corresponding onnx-graphsurgeon Graph.
"""
from onnx_graphsurgeon.importers.onnx_importer import OnnxImporter
return OnnxImporter.import_graph(onnx_model.graph,
opset=OnnxImporter.get_opset(onnx_model))
def scan_model():
path = os.path.join(TEST_ROOT, "models", "scan.onnx")
model = onnx.load(path)
# Body graph
sum_in = Variable(name="sum_in", dtype=np.float32, shape=(2, ))
next = Variable(name="next", dtype=np.float32, shape=(2, ))
sum_out = Variable(name="sum_out", dtype=np.float32, shape=(2, ))
scan_out = Variable(name="scan_out", dtype=np.float32, shape=(2, ))
body_nodes = [
Node(op="Add", inputs=[sum_in, next], outputs=[sum_out]),
Node(op="Identity", inputs=[sum_out], outputs=[scan_out]),
]
body_graph = Graph(nodes=body_nodes,
inputs=[sum_in, next],
outputs=[sum_out, scan_out],
name="scan_body")
# Outer graph
inputs = [
Variable(name="initial", dtype=np.float32, shape=(2, )),
Variable(name="x", dtype=np.float32, shape=(3, 2)),
]
outputs = [
Variable(name="y", dtype=np.float32, shape=(2, )),
Variable(name="z", dtype=np.float32, shape=(3, 2)),
]
attrs = OrderedDict()
attrs["body"] = body_graph
attrs["num_scan_inputs"] = 1
scan_node = Node(op="Scan", inputs=inputs, outputs=outputs, attrs=attrs)
return Model(
path,
inputs=inputs,
outputs=outputs,
nodes=[scan_node],
opset=OnnxImporter.get_opset(model),
)
def lstm_model():
path = os.path.join(TEST_ROOT, "models", "lstm.onnx")
model = onnx.load(path)
onnx_graph = model.graph
def load_initializer(index: int) -> np.ndarray:
return onnx.numpy_helper.to_array(onnx_graph.initializer[index])
# Optional inputs are represented by empty tensors
X = Variable(name="X", dtype=np.float32, shape=(4, 3, 6))
W = Constant(name="W", values=load_initializer(0))
R = Constant(name="R", values=load_initializer(1))
B = Constant(name="B", values=load_initializer(2))
initial_c = Constant(name="initial_c", values=load_initializer(3))
Y = Variable(name="Y", dtype=np.float32, shape=(4, 1, 3, 5))
Y_h = Variable(name="Y_h", dtype=np.float32, shape=(1, 3, 5))
Y_c = Variable(name="Y_c", dtype=np.float32, shape=(1, 3, 5))
attrs = OrderedDict()
attrs["direction"] = "forward"
attrs["hidden_size"] = 5
node = Node(
op="LSTM",
attrs=attrs,
inputs=[X, W, R, B,
Variable.empty(),
Variable.empty(), initial_c],
outputs=[Y, Y_h, Y_c],
)
# Initializers will not be included in the graph inputs.
return Model(
path,
inputs=[X],
outputs=[Y, Y_h, Y_c],
nodes=[node],
opset=OnnxImporter.get_opset(model),
)