wehub-resource-sync
313 lines
10 KiB
Python
313 lines
10 KiB
Python
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
from onnx_graphsurgeon.logger import G_LOGGER
|
|
from onnx_graphsurgeon.ir.tensor import Tensor, Constant, Variable
|
|
from onnx_graphsurgeon.ir.graph import Graph
|
|
from onnx_graphsurgeon.ir.node import Node
|
|
from onnx_graphsurgeon.importers.onnx_importer import OnnxImporter
|
|
|
|
G_LOGGER.severity = G_LOGGER.ULTRA_VERBOSE
|
|
|
|
from collections import OrderedDict
|
|
import onnx.numpy_helper
|
|
from typing import List
|
|
import numpy as np
|
|
import onnx
|
|
import os
|
|
|
|
TEST_ROOT = os.path.realpath(os.path.dirname(__file__))
|
|
|
|
|
|
class Model(object):
|
|
def __init__(
|
|
self,
|
|
path: str,
|
|
inputs: List[Tensor],
|
|
outputs: List[Tensor],
|
|
nodes: List[Node],
|
|
opset: int = None,
|
|
):
|
|
self.path = path
|
|
self.inputs = inputs
|
|
self.outputs = outputs
|
|
self.nodes = nodes
|
|
self.opset = opset
|
|
|
|
def load(self):
|
|
return onnx.load(self.path)
|
|
|
|
def assert_equal(self, graph: Graph):
|
|
assert graph.inputs == self.inputs
|
|
G_LOGGER.debug("Graph inputs matched")
|
|
|
|
# Break down fields to make debugging failures easier.
|
|
for actual, expected in zip(graph.nodes, self.nodes):
|
|
|
|
def check_tensor_io(actensor, extensor):
|
|
def check_list(aclist, exlist):
|
|
G_LOGGER.debug(
|
|
"Actual node list: {:}\n\nExpected node list: {:}".format(
|
|
aclist, exlist
|
|
)
|
|
)
|
|
assert len(aclist) == len(exlist)
|
|
for acnode, exnode in zip(aclist, exlist):
|
|
assert acnode == exnode
|
|
|
|
G_LOGGER.debug("Checking tensor: {:} inputs".format(actensor.name))
|
|
check_list(actensor.inputs, extensor.inputs)
|
|
G_LOGGER.debug("Checking tensor: {:} outputs".format(actensor.name))
|
|
check_list(actensor.outputs, extensor.outputs)
|
|
|
|
G_LOGGER.debug(
|
|
"Actual Node: {:}\n\nExpected Node: {:}".format(actual, expected)
|
|
)
|
|
assert actual.op == expected.op
|
|
assert actual.inputs == expected.inputs
|
|
# Check I/O of input tensors
|
|
for acinp, exinp in zip(actual.inputs, expected.inputs):
|
|
check_tensor_io(acinp, exinp)
|
|
|
|
assert actual.outputs == expected.outputs
|
|
# Check I/O of output tensors
|
|
for acout, exout in zip(actual.outputs, expected.outputs):
|
|
check_tensor_io(acout, exout)
|
|
|
|
assert actual.name == expected.name
|
|
assert len(actual.attrs) == len(expected.attrs)
|
|
for (ackey, acval), (exkey, exval) in zip(
|
|
actual.attrs.items(), expected.attrs.items()
|
|
):
|
|
assert ackey == exkey
|
|
assert acval == exval
|
|
assert actual == expected
|
|
G_LOGGER.debug("Graph nodes matched")
|
|
|
|
assert graph.outputs == self.outputs
|
|
G_LOGGER.debug("Graph outputs matched")
|
|
|
|
def __str__(self):
|
|
return os.path.basename(self.path)
|
|
|
|
|
|
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 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),
|
|
)
|
|
|
|
|
|
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 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)
|
|
)
|
|
|
|
|
|
# Node includes a subgraph whose I/O names are the same as that of the node.
|
|
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 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 const_foldable():
|
|
path = os.path.join(TEST_ROOT, "models", "const_foldable.onnx")
|
|
return Model(
|
|
path, inputs=None, outputs=None, nodes=None, opset=None
|
|
) # Only used for path.
|
|
|
|
|
|
def shape_cast_elision():
|
|
path = os.path.join(TEST_ROOT, "models", "shape_cast_elision.onnx")
|
|
return Model(
|
|
path, inputs=None, outputs=None, nodes=None, opset=None
|
|
) # Only used for path.
|
|
|
|
|
|
def sparse_nnz_model():
|
|
path = os.path.join(TEST_ROOT, "models", "sparse_nnz.onnx")
|
|
return Model(
|
|
path, inputs=None, outputs=None, nodes=None, opset=None
|
|
) # Only used for path.
|
|
|
|
|
|
def sparse_nnz_rank_model():
|
|
path = os.path.join(TEST_ROOT, "models", "sparse_nnz_rank.onnx")
|
|
return Model(
|
|
path, inputs=None, outputs=None, nodes=None, opset=None
|
|
) # Only used for path.
|