def _get_node_id(self, node):
try:
return node.id
except AttributeError:
G_LOGGER.critical(
"Encountered a node not in the graph:\n{:}.\n\nTo fix this, please append the node to this graph's `nodes` attribute."
.format(node))
def __init__(self,
name: str,
values: Union[np.ndarray, LazyValues],
data_location: int = None):
"""
Represents a Tensor whose value is known.
Args:
name (str): The name of the tensor.
values (numpy.ndarray): The values in this tensor, in the form of a NumPy array.
data_location (int):
An enum value indicating the location where the tensor data is stored.
Generally, this will come from onnx.TensorProto.DataLocation.
"""
self.name = name
self.inputs = misc.SynchronizedList(self,
field_name="outputs",
initial=[])
self.outputs = misc.SynchronizedList(self,
field_name="inputs",
initial=[])
if not isinstance(values, np.ndarray) and not isinstance(
values, LazyValues):
G_LOGGER.critical(
"Provided `values` argument is not a NumPy array or a LazyValues instance. "
"Please provide a NumPy array or LazyValues instance to construct a Constant. "
"Note: Provided `values` parameter was: {:}".format(values))
self._values = values
self.data_location = data_location
def cleanup(self, remove_unused_node_outputs=False):
"""
Removes unused nodes and tensors from the graph.
A node or tensor is considered unused if it does not contribute to any of the graph outputs.
Additionally, any producer nodes of graph input tensors are removed from the graph.
*Note: This function will never modify graph output tensors.*
Args:
remove_unused_node_outputs (bool): Whether to remove unused output tensors of nodes. This will never remove
empty-tensor (i.e. optional, but omitted) outputs. Defaults to False.
Returns:
self
"""
with self.node_ids():
# Graph inputs cannot have producers
for inp in self.inputs:
inp.inputs.clear()
used_node_ids, used_tensors = self._get_used_node_ids()
inputs = []
for inp in self.inputs:
if inp in used_tensors:
inputs.append(inp)
else:
G_LOGGER.debug("Removing unused input: {:}".format(inp))
self.inputs = inputs
nodes = []
for node in self.nodes:
if self._get_node_id(node) in used_node_ids:
nodes.append(node)
else:
node.inputs.clear()
node.outputs.clear()
G_LOGGER.verbose("Removing unused node: {:}".format(node))
# Last pass to remove any hanging tensors - tensors without outputs
if remove_unused_node_outputs:
graph_output_names = set(
[tensor.name for tensor in self.outputs])
for node in nodes:
def is_hanging_tensor(tensor):
return not tensor.is_empty() and len(
tensor.outputs
) == 0 and tensor.name not in graph_output_names
to_remove = [
out for out in node.outputs if is_hanging_tensor(out)
]
for out in to_remove:
if out in node.outputs:
node.outputs.remove(out)
self.nodes = nodes
return self
def register_func(func):
if hasattr(Graph, func.__name__):
G_LOGGER.warning("Registered function: {:} is hidden by a Graph attribute or function with the same name. This function will never be called!".format(func.__name__))
for opset in opsets:
Graph.OPSET_FUNC_MAP[opset][func.__name__] = func
return func
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
def attrs_to_dict(attrs):
attr_dict = OrderedDict()
for attr in attrs:
def process_attr(attr_str: str):
processed = getattr(attr,
ONNX_PYTHON_ATTR_MAPPING[attr_str])
if attr_str == "STRING":
processed = processed.decode()
elif attr_str == "TENSOR":
processed = OnnxImporter.import_tensor(processed)
elif attr_str == "GRAPH":
processed = OnnxImporter.import_graph(
processed,
misc.combine_dicts(tensor_map,
subgraph_tensor_map))
elif attr_str == "FLOATS" or attr_str == "INTS":
processed = list(processed)
elif attr_str == "STRINGS":
processed = [p.decode() for p in processed]
return processed
if attr.type in ATTR_TYPE_MAPPING:
attr_str = ATTR_TYPE_MAPPING[attr.type]
if attr_str in ONNX_PYTHON_ATTR_MAPPING:
attr_dict[attr.name] = process_attr(attr_str)
else:
G_LOGGER.warning(
"Attribute of type {:} is currently unsupported. Skipping attribute."
.format(attr_str))
else:
G_LOGGER.warning(
"Attribute type: {:} was not recognized. Was the graph generated with a newer IR version than the installed `onnx` package? Skipping attribute."
.format(attr.type))
return attr_dict
def add_to_tensor_map(tensor):
if not tensor.is_empty():
if check_duplicates and tensor.name in tensor_map and not (tensor_map[tensor.name] is tensor):
G_LOGGER.critical("Found distinct tensors that share the same name:\n[id: {:}] {:}\n[id: {:}] {:}"
.format(id(tensor_map[tensor.name]), tensor_map[tensor.name], id(tensor), tensor))
tensor_map[tensor.name] = tensor
def __init__(self,
nodes: Sequence[Node] = None,
inputs: Sequence[Tensor] = None,
outputs: Sequence[Tensor] = None,
name=None,
doc_string=None,
opset=None):
"""
Args:
nodes (Sequence[Node]): A list of the nodes in this graph.
inputs (Sequence[Tensor]): A list of graph input Tensors.
outputs (Sequence[Tensor]): A list of graph output Tensors.
name (str): The name of the graph. Defaults to "onnx_graphsurgeon_graph".
doc_string (str): A doc_string for the graph. Defaults to "".
"""
self.nodes = misc.default_value(nodes, [])
self.inputs = list(misc.default_value(inputs, []))
self.outputs = list(misc.default_value(outputs, []))
self.name = misc.default_value(name, "onnx_graphsurgeon_graph")
self.__name__ = self.name
self.doc_string = misc.default_value(doc_string, "")
self.opset = misc.default_value(opset, Graph.DEFAULT_OPSET)
# Printing graphs can be very expensive
G_LOGGER.ultra_verbose(lambda: "Created Graph: {:}".format(self))
# For layer() function
self.name_idx = 0
def __getattr__(self, name):
try:
return super().__getattribute__(name)
except AttributeError as err:
if self.opset not in Graph.OPSET_FUNC_MAP or name not in Graph.OPSET_FUNC_MAP[self.opset]:
G_LOGGER.error("No function: {:} registered for opset: {:}".format(name, self.opset))
raise err
return lambda *args, **kwargs: Graph.OPSET_FUNC_MAP[self.opset][name](self, *args, **kwargs)
def get_opset(model: onnx.ModelProto):
try:
return model.opset_import[0].version
except:
G_LOGGER.warning(
"Model does not contain opset information! Using default opset."
)
return None
def fold_constants(self):
"""
Folds constants in-place in the graph. The graph must be topologically sorted prior to
calling this function (see `toposort()`).
This function will not remove constants after folding them. In order to get rid of
these hanging nodes, you can run the `cleanup()` function.
*Note: Due to how this function is implemented, the graph must be exportable to ONNX,
and evaluable in ONNX-Runtime. Additionally, ONNX-Runtime must be installed.*
Returns:
self
"""
import onnxruntime
from onnx_graphsurgeon.exporters.onnx_exporter import export_onnx
temp_graph = copy.deepcopy(self)
# Since the graph is topologically sorted, this should find all constant nodes in the graph.
graph_constants = {
tensor.name: tensor
for tensor in temp_graph.tensors().values()
if isinstance(tensor, Constant)
}
for node in temp_graph.nodes:
if all([inp.name in graph_constants for inp in node.inputs]):
graph_constants.update({out.name: out for out in node.outputs})
# Next build a graph with just the constants, and evaluate - no need to evaluate constants
outputs_to_evaluate = [
tensor for tensor in graph_constants.values()
if isinstance(tensor, Variable)
]
if not outputs_to_evaluate:
G_LOGGER.warning(
"Could not find any operations in this graph that can be folded. This could mean that constant folding has already been run on this graph. Skipping."
)
return self
output_names = [out.name for out in outputs_to_evaluate]
temp_graph.outputs = outputs_to_evaluate
temp_graph.cleanup()
# Determining types is not trivial, and ONNX-RT does its own type inference.
sess = onnxruntime.InferenceSession(
export_onnx(temp_graph, do_type_check=False).SerializeToString())
constant_values = sess.run(output_names, {})
# Finally, replace the Variables in the original graph with constants.
graph_tensors = self.tensors()
for name, values in zip(output_names, constant_values):
graph_tensors[name].to_constant(values)
graph_tensors[name].inputs.clear() # Constants do not need inputs
return self
def export_value_info_proto(tensor: Variable, do_type_check: bool) -> onnx.ValueInfoProto:
if do_type_check and tensor.dtype is None:
G_LOGGER.critical("Graph input and output tensors must include dtype information. Please set the dtype attribute for: {:}".format(tensor))
if tensor.dtype is not None:
onnx_tensor = onnx.helper.make_tensor_value_info(tensor.name, dtype_to_onnx(tensor.dtype), tensor.shape)
else:
onnx_tensor = onnx.helper.make_empty_tensor_value_info(tensor.name)
return onnx_tensor
def register_func(func):
if hasattr(Graph, func.__name__):
G_LOGGER.warning("Registered function: {:} is hidden by a Graph attribute or function with the same name. This function will never be called!".format(func.__name__))
# Default behavior is to register functions for all opsets.
if opsets is None:
Graph.GLOBAL_FUNC_MAP[func.__name__] = func
else:
for opset in opsets:
Graph.OPSET_FUNC_MAP[opset][func.__name__] = func
return func
def check_tensor(name: str):
if name not in tensor_map:
if name:
G_LOGGER.debug(
"Tensor: {:} was not generated during shape inference, or shape inference was not run on this model. Creating a new Tensor."
.format(name))
tensor_map[name] = Variable(name)
else:
# Empty tensors are not tracked by the graph, as these represent optional inputs/outputs that have been omitted.
G_LOGGER.verbose("Generating empty tensor")
return Variable.empty()
return tensor_map[name]
def process_io(io):
new_io = []
for elem in io:
if isinstance(elem, Tensor):
new_io.append(elem)
elif isinstance(elem, str):
tensor = Variable(name=self._generate_name(elem))
new_io.append(tensor)
elif isinstance(elem, np.ndarray):
new_io.append(Constant(name=self._generate_name("onnx_graphsurgeon_constant"), values=elem))
else:
G_LOGGER.critical("Unrecognized type passed to Graph.layer: {:}.\n\tHint: Did you forget to unpack a list with `*`?\n\tPlease use Tensors, strings, or NumPy arrays.".format(elem))
return new_io
def __getattr__(self, name):
try:
return super().__getattribute__(name)
except AttributeError as err:
# Opset specific ops always take priority over global ops.
if self.opset in Graph.OPSET_FUNC_MAP and name in Graph.OPSET_FUNC_MAP[self.opset]:
return lambda *args, **kwargs: Graph.OPSET_FUNC_MAP[self.opset][name](self, *args, **kwargs)
if name in Graph.GLOBAL_FUNC_MAP:
return lambda *args, **kwargs: Graph.GLOBAL_FUNC_MAP[name](self, *args, **kwargs)
G_LOGGER.error("No function: {:} registered for opset: {:}".format(name, self.opset))
raise err
def get_opset(model: onnx.ModelProto):
try:
for importer in OnnxImporter.get_import_domains(model):
if importer.domain == "" or importer.domain == "ai.onnx":
return importer.version
G_LOGGER.warning(
"Model does not contain ONNX domain opset information! Using default opset."
)
return None
except:
G_LOGGER.warning(
"Model does not contain opset information! Using default opset."
)
return None
def infer_model(path):
model = onnx.load(path)
graph = gs.import_onnx(model)
feed_dict = {}
for tensor in graph.inputs:
feed_dict[tensor.name] = np.random.random_sample(size=tensor.shape).astype(tensor.dtype)
output_names = [out.name for out in graph.outputs]
sess = onnxruntime.InferenceSession(model.SerializeToString())
outputs = sess.run(output_names, feed_dict)
G_LOGGER.info("Inference outputs: {:}".format(outputs))
return outputs
def test_examples(example_dir, artifacts):
example_dir = os.path.join(EXAMPLES_ROOT, example_dir)
readme = os.path.join(example_dir, "README.md")
commands = load_commands_from_readme(readme)
for command in commands:
G_LOGGER.info(command)
assert sp.run(["bash", "-c", command], cwd=example_dir, env={"PYTHONPATH": ROOT_DIR}).returncode == 0
for artifact in artifacts:
artifact_path = os.path.join(example_dir, artifact.name)
assert os.path.exists(artifact_path)
if artifact.infer:
assert infer_model(artifact_path)
os.remove(artifact_path)
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)
def __init__(self, name: str, values: np.ndarray):
"""
Represents a Tensor whose value is known.
Args:
name (str): The name of the tensor.
values (numpy.ndarray): The values in this tensor, in the form of a NumPy array.
dtype (numpy.dtype): The data type of the tensor.
shape (Sequence[Union[int, str]]): The shape of the tensor.
"""
self.name = name
self.inputs = misc.SynchronizedList(self, field_name="outputs", initial=[])
self.outputs = misc.SynchronizedList(self, field_name="inputs", initial=[])
if not isinstance(values, np.ndarray):
G_LOGGER.critical("Provided `values` argument is not a NumPy array (please provide a NumPy array to construct a Constant): {:}".format(values))
self.values = np.array(values)
def __eq__(self, other):
"""
Check whether two nodes are equal by comparing name, attributes, op, inputs, and outputs.
"""
G_LOGGER.verbose("Comparing node: {:} with {:}".format(
self.name, other.name))
attrs_match = self.name == other.name and self.op == other.op and self.attrs == other.attrs
inputs_match = len(self.inputs) == len(other.inputs) and all([
inp == other_inp
for inp, other_inp in zip(self.inputs, other.inputs)
])
outputs_match = len(self.outputs) == len(other.outputs) and all([
out == other_out
for out, other_out in zip(self.outputs, other.outputs)
])
return attrs_match and inputs_match and outputs_match
def get_tensor(name: str, check_outer_graph=True):
# Prioritize the subgraph even if check_outer_graph is set
if name in subgraph_tensor_map:
return subgraph_tensor_map[name]
if check_outer_graph and name in tensor_map:
return tensor_map[name]
if not name:
# Empty tensors are not tracked by the graph, as these represent optional inputs/outputs that have been omitted.
G_LOGGER.verbose("Generating empty tensor")
return Variable.empty()
G_LOGGER.verbose("Tensor: {:} was not generated during shape inference, or shape inference was not run on this model. Creating a new Tensor.".format(name))
subgraph_tensor_map[name] = Variable(name)
return subgraph_tensor_map[name]
def cleanup(self, remove_unused_node_outputs=True):
"""
Removes unused nodes and tensors from the graph.
A node or tensor is considered unused if it does not contribute to any of the graph outputs.
Note: This function will never modify graph output tensors.
Optional Args:
remove_unused_node_outputs (bool): Whether to remove unused output tensors of nodes. This will never remove
empty tensor outputs. If this is set to False, outputs of nodes kept in the graph will not be modified.
Returns:
self
"""
with self.node_ids():
used_node_ids, used_tensors = self._get_used_node_ids()
inputs = []
for inp in self.inputs:
if inp in used_tensors:
inputs.append(inp)
else:
G_LOGGER.debug("Removing unused input: {:}".format(inp))
self.inputs = inputs
nodes = []
for node in self.nodes:
if self._get_node_id(node) in used_node_ids:
nodes.append(node)
else:
node.inputs.clear()
node.outputs.clear()
G_LOGGER.verbose("Removing unused node: {:}".format(node))
# Last pass to remove any hanging tensors - tensors without outputs
if remove_unused_node_outputs:
graph_output_names = set([tensor.name for tensor in self.outputs])
for node in nodes:
def is_hanging_tensor(tensor):
return not tensor.is_empty() and len(tensor.outputs) == 0 and tensor.name not in graph_output_names
[node.outputs.remove(out) for out in node.outputs if is_hanging_tensor(out)]
self.nodes = nodes
return self
def get_hierarchy_level(node):
# Return all nodes that contribute to this node.
def get_input_nodes(node):
inputs = {}
for tensor in node.inputs:
for node in tensor.inputs:
inputs[self._get_node_id(node)] = node
return inputs.values()
if self._get_node_id(node) in hierarchy_levels:
return hierarchy_levels[self._get_node_id(node)].level
# The level of a node is the level of it's highest input + 1.
try:
max_input_level = max([get_hierarchy_level(input_node) for input_node in get_input_nodes(node)] + [-1])
except RecursionError:
G_LOGGER.critical("Cycle detected in graph! Are there tensors with duplicate names in the graph?")
return max_input_level + 1
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 assert_equal(self, graph: Graph):
assert graph.inputs == self.inputs
G_LOGGER.debug("Graph inputs matched")
for actual, expected in zip(graph.nodes, self.nodes):
G_LOGGER.debug("Actual Node: {:}.\n\nExpected Node: {:}".format(
actual, expected))
# Break down fields to make debugging failures easier.
assert actual.op == expected.op
assert actual.inputs == expected.inputs
assert actual.outputs == expected.outputs
assert actual.name == expected.name
for (akey, aval), (ekey, eval) in zip(actual.attrs.items(),
expected.attrs.items()):
assert akey == ekey
assert aval == eval
assert actual == expected
G_LOGGER.debug("Graph nodes matched")
assert graph.outputs == self.outputs
G_LOGGER.debug("Graph outputs matched")
def import_graph(onnx_graph: onnx.GraphProto,
tensor_map: "OrderedDict[str, Tensor]" = None,
opset=None,
import_domains: onnx.OperatorSetIdProto = None) -> Graph:
"""
Imports a Graph from an ONNX Graph.
Args:
onnx_graph (onnx.GraphProto): The ONNX graph to import.
tensor_map (OrderedDict[str, Tensor]): A mapping of tensor names to Tensors. This is generally only useful for subgraph import.
opset (int): The ONNX opset to use for this graph.
"""
tensor_map = copy.copy(misc.default_value(
tensor_map, OrderedDict())) # Outer graph tensors, read-only
subgraph_tensor_map = OrderedDict() # Tensors in this subgraph
# Retrieves a Tensor from subgraph_tensor_map or the outer graph (tensor_map) if present, otherwise imports the tensor
# If overwrite=True, this function will overwrite previously imported tensors
# if the new tensor has more information available.
def get_tensor(onnx_tensor: Union[onnx.ValueInfoProto,
onnx.TensorProto],
overwrite=False,
check_outer_graph=True) -> Tensor:
# Prioritize the subgraph even if check_outer_graph is set
if onnx_tensor.name in subgraph_tensor_map:
if overwrite:
tensor = OnnxImporter.import_tensor(onnx_tensor)
if isinstance(subgraph_tensor_map[onnx_tensor.name],
Variable):
subgraph_tensor_map[
onnx_tensor.name].dtype = subgraph_tensor_map[
onnx_tensor.name].dtype or tensor.dtype
subgraph_tensor_map[
onnx_tensor.name].shape = subgraph_tensor_map[
onnx_tensor.name].shape or tensor.shape
return subgraph_tensor_map[onnx_tensor.name]
if check_outer_graph and onnx_tensor.name in tensor_map:
return tensor_map[onnx_tensor.name]
subgraph_tensor_map[onnx_tensor.name] = OnnxImporter.import_tensor(
onnx_tensor)
return subgraph_tensor_map[onnx_tensor.name]
# Import initializers contents into Constants.
G_LOGGER.verbose("Importing initializers")
for initializer in onnx_graph.initializer:
get_tensor(initializer)
# Import all tensors whose shapes are known. Tensors may be repeated, and some of these
# duplicates may not include shape/dtype information, so overwrite is set to True
# so that we can capture all the information available about the tensor
G_LOGGER.verbose("Importing tensors with known shapes")
for tensor in onnx_graph.value_info:
get_tensor(tensor, overwrite=True)
# Import graph inputs and outputs. Initializers are not considered to be inputs.
# Graph inputs and outputs can never come from the outer graph!
initializer_names = set(
[tensor.name for tensor in onnx_graph.initializer])
G_LOGGER.verbose("Importing graph inputs")
graph_inputs = [] # List[Tensor]
for inp in onnx_graph.input:
if inp.name not in initializer_names:
tensor = get_tensor(inp, check_outer_graph=False)
graph_inputs.append(tensor)
G_LOGGER.verbose("Importing graph outputs")
graph_outputs = [] # List[Tensor]
for out in onnx_graph.output:
tensor = get_tensor(out, check_outer_graph=False)
graph_outputs.append(tensor)
G_LOGGER.verbose("Importing nodes")
nodes = [] # List[Node]
for onnx_node in onnx_graph.node:
node = OnnxImporter.import_node(onnx_node, tensor_map,
subgraph_tensor_map)
nodes.append(node)
return Graph(nodes=nodes,
inputs=graph_inputs,
outputs=graph_outputs,
name=onnx_graph.name,
doc_string=onnx_graph.doc_string,
opset=opset,
import_domains=import_domains)
