def add_initializers_and_inputs_to_graph(graph: Graph, graph_pb, data_nodes_map: dict):
"""
The function adds nodes specified in the 'initializer' attribute of the pb and input nodes.
:param graph: the Graph to add nodes to
:param graph_pb: the graph protobuf message
:param data_nodes_map: the dictionary with mapping of tensor names to node id and port
:return: the list of Parameter nodes
"""
initializers = Graph()
fill_graph_with_nodes(initializers, graph_pb.initializer, get_id=lambda pb: pb.name, get_attrs=protobuf_attrs)
parameters = []
# first go through all inputs and separate constant from placeholders
for inp in graph_pb.input:
name = str(inp.name)
if graph.has_node(name):
raise Error('Name {} of input node already exists, input names are duplicated.', name)
elif initializers.has_node(name):
graph.add_node(name, kind='op', op='Const', pb=inp, pb_init=initializers.node[name]['pb'])
else:
graph.add_node(name, kind='op', op='Parameter', pb=inp)
parameters.append(Node(graph, name))
assert name not in data_nodes_map, 'Inconsistency between data_nodes_map and graph.nodes'
data_nodes_map[name] = (name, 0)
# go over all initializers and make sure that all of them are added to the graph
for initializer in initializers.nodes():
initializer_id = initializer
if not graph.has_node(initializer_id):
graph.add_node(initializer_id, kind='op', op='Const', pb=initializers.node[initializer]['pb'],
pb_init=initializers.node[initializer]['pb'])
data_nodes_map[initializer] = (initializer_id, 0)
return parameters
def update_body_graph(body_graph: Graph, subgraph_proto: dict,
body_parameter_names: list, body_results: list):
"""
Updates the loop body graph with a sub-graph (for body or condition functions)
:param body_graph: a loop body graph to be updated
:param subgraph_proto: a sub-graph in a protobuf format to be added into the loop body graph
:param body_parameter_names: a (unchanged) list of parameters in the loop body graph
:param body_results: a list of Result nodes that is extended with a list from a sub-graph
"""
# create a map from a node name in original model to a name in a loop body graph assuming
# that names in the original model are unique
# initially, the map contains names for parameters that are common for the body and condition graphs
map_original_name = {}
for idx, pb_node in enumerate(subgraph_proto['input_arg']):
map_original_name[pb_node.name] = body_parameter_names[idx]
# walk through all nodes (non-parameter and non-result nodes) and add into the loop body graph
for pb_node in subgraph_proto['node_def']:
# create an NX node
id = body_graph.unique_id(pb_node.name)
map_original_name[pb_node.name] = id
body_graph.add_node(id, pb=pb_node, kind='op')
if hasattr(body_graph, 'op_names_statistic') and hasattr(
pb_node, 'op'):
body_graph.op_names_statistic[pb_node.op] += 1
# add incoming edges based on data_nodes_map
for dst_port, inp in enumerate(pb_node.input):
orig_src_id = inp.split(":")[0]
# TODO: avoid this temporal workaround for TF 2.4 or higher RNN layers:
# skip control flow dependency
if orig_src_id[0] == '^':
continue
src_id = map_original_name[orig_src_id]
src_port = 0 if len(inp.split(":")) == 1 else int(
inp.split(":")[-1])
assert (body_graph.has_node(src_id))
body_graph.add_edges_from(
[create_tf_edge(src_id + ":" + str(src_port), id, dst_port)])
# create Result nodes in the loop body graph
for output in subgraph_proto['output_arg']:
output_name = subgraph_proto['ret'][output.name]
orig_src_id = output_name.split(":")[0]
src_id = map_original_name[orig_src_id]
src_port = 0 if len(output_name.split(":")) == 1 \
else int(output_name.split(":")[-1])
assert body_graph.has_node(
src_id
), 'The body graph does not contain output with name "{}"'.format(
src_id)
body_results.append(
Node(body_graph, add_opoutput(body_graph, src_id, src_port,
False)))
return True
def apply_pattern(graph: Graph, nodes: list, edges: list, action: callable, node_attrs: list = None,
edge_attrs: list = None):
"""
Search for all matches of a given subgraph defined by [nodes, edges] in graph,
then apply action for each such match.
"""
if not all_edges_in_nodes([node[0] for node in nodes], edges):
log.warning("Incorrect pattern attributes: not all nodes from edges are in nodes. "
"Please, mention all nodes you need in pattern in nodes attribute. ")
matches = []
for match in find_pattern_matches(graph, nodes, edges, node_attrs, edge_attrs):
matches.append(match)
for match in matches:
match = inverse_dict(match)
still_valid = True
for k in match:
if not graph.has_node(match[k]):
# Graph changed significantly
still_valid = False
log.warning("The graph has changed significantly during applying pattern:\n"
"nodes: {}\n"
"edges: {}\n"
"node_attrs: {}\n"
"edge_attrs: {}".format(nodes, edges, node_attrs, edge_attrs))
break
match[k] = Node(graph, match[k])
if still_valid:
action(graph, match)
def find_and_replace_pattern(self, graph: Graph):
reshape_nodes = graph.get_op_nodes(type='Reshape')
for node in reshape_nodes:
if not graph.has_node(node.id):
# the Reshape node has been removed in the previous iteration
continue
if len(node.out_port(0).get_destinations()) == 1:
log.debug('First phase for Reshape: {}'.format(node.soft_get('name')))
next_op = get_next_operation(node)[0]
log.debug('second node: id={}, type={}'.format(next_op.soft_get('id'), next_op.soft_get('type')))
if next_op.has_valid('type') and next_op.type == 'Reshape':
dim_value = next_op.in_port(1).data.get_value()
if dim_value is None or 0 in dim_value or -1 in dim_value:
# we do not fuse reshape sequences with special symbols: 0, -1
continue
# Detected Reshape1 --> data --> Reshape2 pattern without side edges. Remove Reshape1
log.debug('Second phase for Reshape: {}'.format(node.soft_get('name')))
remove_op_node_with_data_node(graph, node)
def replace_sub_graph(self, graph: Graph, match: dict):
box_nms = match['box_nms']
top_k = box_nms.topk
nms_threshold = box_nms.overlap_thresh
ssd_concats = {}
concat_names = ['ssd_concat1', 'ssd_concat0', 'ssd_concat2']
for i, concat_match in enumerate(self.concats_pattern):
for matches in find_pattern_matches(graph, concat_match['nodes'],
concat_match['edges'], None,
None):
for match in matches:
if graph.has_node(match):
n = Node(graph, match)
if n.op == 'Concat':
ssd_concats.update({concat_names[i]: n})
break
assert concat_names[0] in ssd_concats
assert concat_names[1] in ssd_concats
assert concat_names[2] in ssd_concats
graph.remove_nodes_from(graph.get_nodes_with_attributes(op='Result'))
detection_output_node = DetectionOutput(
graph,
dict(name=graph.unique_id() + '/DetectionOutput_',
top_k=top_k,
keep_top_k=top_k,
nms_threshold=nms_threshold,
background_label_id=0,
clip=0,
decrease_label_id=1,
code_type="caffe.PriorBoxParameter.CENTER_SIZE",
confidence_threshold=0.01,
share_location=1,
variance_encoded_in_target=0,
normalized=1)).create_node()
reshape_node = create_op_node_with_second_input(
graph, Reshape, int64_array([0, -1]),
dict(name=graph.unique_id() + '/DetectionOutput_'))
ssd_softmax_node = ssd_concats['ssd_concat0'].out_node().out_node()
ssd_softmax_node.out_port(0).disconnect()
ssd_softmax_node.out_port(0).connect(reshape_node.in_port(0))
reshape_node.out_port(0).connect(detection_output_node.in_port(1))
ssd_concats['ssd_concat2'].axis = 2
self.reshape_priorboxes(ssd_concats['ssd_concat2'])
ssd_concats['ssd_concat1'].out_port(
0).get_connection().set_destination(
detection_output_node.in_port(0))
ssd_concats['ssd_concat2'].out_port(
0).get_connection().set_destination(
detection_output_node.in_port(2))
Result(graph, {
'name': detection_output_node.id + '/Result'
}).create_node([detection_output_node])
def merge_nodes(graph: Graph,
nodes_to_merge_names: list,
inputs_desc: list = None,
outputs_desc: list = None):
"""
Merges nodes specified in the set 'nodes_to_merge_names' into one mega-node, creating new edges between mega-node
and inputs/outputs nodes of the mega-node. The added edges contain name of input/output nodes which will be used for
generation of placeholders and will be saved to the IR xml so IE plug-in know how to map input/output data for the
layer. Also the function adds protobufs of the nodes of the sub-graph and 'Const' ops consumed by nodes in the
sub-graph to the node's attribute 'pbs'.
:param graph: the graph object to operate on.
:param nodes_to_merge_names: list of nodes names that should be merged into a single node.
:param inputs_desc: optional list describing input nodes order.
:param outputs_desc: optional list describing output nodes order.
"""
if not is_connected_component(graph, nodes_to_merge_names):
log.warning(
"The following nodes do not form connected sub-graph: {}".format(
nodes_to_merge_names))
# graph.dump_graph_for_graphviz(nodes_to_dump=nodes_to_merge_names)
new_node_name = graph.unique_id("TFSubgraphCall_")
log.info("Create new node with name '{}' for nodes '{}'".format(
new_node_name, ', '.join(nodes_to_merge_names)))
graph.add_node(new_node_name)
new_node_attrs = graph.node[new_node_name]
new_node_attrs['name'] = new_node_name
set_tf_custom_call_node_attrs(new_node_attrs)
new_node = Node(graph, new_node_name)
added_input_tensors_names = set(
) # set of tensors that are were added as input to the sub-graph
added_new_node_output_tensors = dict(
) # key - tensor name, value - out port
for node_name in nodes_to_merge_names:
node = Node(graph, node_name)
add_node_pb_if_not_yet_added(node, new_node)
# TODO: any improvements?
for in_node_name, edge_attrs in Node(graph, node_name).get_inputs():
in_node = Node(graph, in_node_name)
# internal edges between nodes of the sub-graph
if in_node_name in nodes_to_merge_names:
add_node_pb_if_not_yet_added(in_node, new_node)
continue
# edge outside of sub-graph into sub-graph
if in_node_name not in nodes_to_merge_names:
# we cannot use the 'in_node_name' as a protobuf operation name here
# because the 'in_node_name' could be a sub-graph matched before.
input_tensor_name = node.pb.input[edge_attrs['in']]
if input_tensor_name not in added_input_tensors_names:
if not new_node.has_port('in', edge_attrs['in']):
new_node.add_input_port(edge_attrs['in'])
graph.add_edge(
in_node_name, new_node_name,
**merge_edge_props(
{
'in':
find_input_port(new_node, inputs_desc,
node_name, edge_attrs['in']),
'out':
edge_attrs['out'],
'internal_input_node_name':
input_tensor_name,
'original_dst_node_name':
node_name,
'original_dst_port':
edge_attrs['in'],
'in_attrs': [
'in', 'internal_input_node_name',
'original_dst_node_name',
'original_dst_port', 'placeholder_name'
],
'out_attrs': ['out']
}, edge_attrs))
log.debug(
"Creating edge from outside of sub-graph to inside sub-graph: {} -> {}"
.format(in_node_name, new_node_name))
added_input_tensors_names.add(input_tensor_name)
# edge from inside sub-graph to outside sub-graph
for out_node_name, edge_attrs in Node(graph, node_name).get_outputs():
if out_node_name not in nodes_to_merge_names:
log.debug(
"Creating edge from inside of sub-graph to outside sub-graph: {} -> {}"
.format(new_node_name, out_node_name))
out_name = internal_output_name_for_node(
node_name, edge_attrs['out'])
if out_name not in added_new_node_output_tensors.keys():
added_new_node_output_tensors[out_name] = find_output_port(
new_node, outputs_desc, node_name, edge_attrs['out'])
if not new_node.has_port(
'out', added_new_node_output_tensors[out_name]):
new_node.add_output_port(
added_new_node_output_tensors[out_name])
graph.add_edge(
new_node_name, out_node_name,
**merge_edge_props(
{
'in': edge_attrs['in'],
'out': added_new_node_output_tensors[out_name],
'internal_output_node_name': out_name,
'in_attrs': ['in', 'internal_input_node_name'],
'out_attrs': ['out', 'internal_output_node_name']
}, edge_attrs))
new_node['output_tensors_names'] = [
val for val in
{v: k
for k, v in added_new_node_output_tensors.items()}.values()
]
# add nodes using the same order as in initial GraphDef so we can dump them to IR in "correct" order
new_node['nodes_order'] = [
node for node in graph.graph['initial_nodes_order']
if node in new_node['pbs'].keys()
]
for n in nodes_to_merge_names:
if graph.has_node(
n): # check if not deleted by another (similar) pattern
graph.remove_node(n)
return Node(graph, new_node_name)
def protobuf2nx(graph: Graph, pb):
"""
Convert proto message with ONNX model to equivalent NX representation. All nodes and edges are restored here as
ONNX model has op/data representation, that means that nodes are connected via tensor names. Name of tensors are
defined on demand in nodes, so we have a code similar to Caffe here.
:param graph: the Graph object to load the graph into
:param pb: the ONNX file protobuf message
:return: None
"""
# maps a tensor name to a node produced it and the node port: str -> (node_id, node_port)
data_nodes_map = {}
graph_pb = pb.graph
add_initializers_and_inputs_to_graph(graph, graph_pb, data_nodes_map)
output_ids = []
for outp in graph_pb.output:
name = str(outp.name)
if graph.has_node(name):
log.error(
'Name {} of output node already exists in graph. Ignoring this output. If the output is required,'
' please rename it.'.format(name),
extra={'is_warning': True})
continue
else:
# add fake node on output
graph.add_node(name, kind='op', op='FakeOutput', pb=outp)
output_ids.append(name)
# Go through all nodes in the original model order (because data nodes are defined on-the-fly and order is
# important)
for node in graph_pb.node:
# create an NX node
fw_name = node_id(node)
id = graph.unique_id(fw_name)
graph.add_node(id, pb=node, kind='op')
if hasattr(graph, 'op_names_statistic') and hasattr(node, 'op_type'):
graph.op_names_statistic[node.op_type] += 1
# add incoming edges based on data_nodes_map
for dst_port, inp in enumerate(node.input):
# should add edge inp --> id
if inp not in data_nodes_map:
if inp == '':
# input is omitted; most likely it corresponds to an optional input for an operator
continue
else:
raise Error(
'Reference to {} is not satisfied. A node refer not existing data tensor. ONNX model is not '
'consistent. Protobuf fragment: {}', inp, node)
src_id, src_port = data_nodes_map[inp]
assert (graph.has_node(src_id))
edge_attrs = {
'out': src_port,
'in': dst_port,
'name': inp,
'fw_tensor_debug_info': [(src_id, inp)],
'in_attrs': ['in', 'name'],
'out_attrs': ['out', 'name'],
'data_attrs': ['fw_tensor_debug_info']
}
graph.add_edge(src_id, id, **edge_attrs)
# add outgoing edges to data_nodes_map
for src_port, out in enumerate(node.output):
if out in output_ids:
edge_attrs = {
'out': src_port,
'in': 0,
'name': out,
'fw_tensor_debug_info': [(fw_name, out)],
'in_attrs': ['in', 'name'],
'out_attrs': ['out', 'name'],
'data_attrs': ['fw_tensor_debug_info']
}
graph.add_edge(id, out, **edge_attrs)
if out in data_nodes_map:
log.debug("Detected reuse of blob {}.".format(out))
data_nodes_map[out] = (id, src_port)
graph.graph[
'tensor_mapping'] = data_nodes_map # save main graph tensor names mapping for Loop op parsing
def extract(cls, loop_node):
Loop.update_node_stat(loop_node, {})
body_graph_proto = onnx_attr(loop_node, 'body', 'g', None)
main_graph = loop_node.graph
# create a Graph object for the body and take graph attributes from the main graph
body_graph = Graph()
main_graph_attrs_copy = {}
for attr_key, attr_value in main_graph.graph.items():
if attr_key not in ['tensor_mapping', 'parent_node']:
main_graph_attrs_copy[attr_key] = copy.deepcopy(attr_value)
body_graph.graph.update(main_graph_attrs_copy)
loop_node['body'] = body_graph
# save parent node for nested loops to know which node contains body (and which graph is on upper level)
body_graph.graph['parent_node'] = loop_node
# maps a tensor name to a node produced it and the node port: str -> (node_id, node_port)
data_nodes_map = {}
body_graph.graph['tensor_mapping'] = data_nodes_map # save mapping for possible Loop inside the Loop
body_parameters = add_initializers_and_inputs_to_graph(body_graph, body_graph_proto, data_nodes_map)
external_edges = [] # (src_node, src_out_port), dest_body_parameter_node
# save additional edges information for graph on each level, the first one is the deepest
additional_params = [] # (src_node, src_out_port) -> parameter_node (for manually added Parameters)
# Go through all nodes in the original model order because data nodes are defined on-the-fly and order matters
for pb_node in body_graph_proto.node:
# create an NX node
id = body_graph.unique_id(node_id(pb_node))
body_graph.add_node(id, pb=pb_node, kind='op')
if hasattr(body_graph, 'op_names_statistic') and hasattr(pb_node, 'op_type'):
body_graph.op_names_statistic[pb_node.op_type] += 1
# add incoming edges based on data_nodes_map
for dst_port, inp in enumerate(pb_node.input):
# should add edge src_internal_id --> dst_id
if inp not in data_nodes_map:
if inp == '':
# input is omitted; most likely it corresponds to an optional input for an operator
continue
else:
is_finished = create_cross_body_edge(body_graph, external_edges, additional_params,
inp, id, dst_port)
if not is_finished:
raise Error(
'Reference to "{}" is not satisfied. A node refer not existing data tensor. ONNX '
'model is not consistent. Protobuf fragment: {}', inp, pb_node)
else:
src_id, src_port = data_nodes_map[inp]
create_edge_with_attrs(body_graph, inp, src_id, src_port, id, dst_port)
# add outgoing edges to data_nodes_map
for src_port, out in enumerate(pb_node.output):
if out in data_nodes_map:
log.debug("Detected reuse of blob {}.".format(out))
data_nodes_map[out] = (id, src_port)
body_results = []
for output in body_graph_proto.output:
tensor_name = str(output.name)
node_name, output_port = data_nodes_map[tensor_name]
assert body_graph.has_node(node_name), 'The body graph does not contain output with name "{}"'.format(
node_name)
body_results.append(Node(body_graph, add_opoutput(body_graph, node_name, output_port, False)))
# add 'internal_layer_id' attribute which is a must have attribute for the loop body node
for idx, body_node in enumerate(body_graph.get_op_nodes()):
body_node['internal_layer_id'] = idx
loop_carried_dependencies_count = len(body_graph_proto.input) - 2
scan_outputs_count = len(body_graph_proto.output) - 1 - loop_carried_dependencies_count
# Loop inputs:
# 0 - trip count
# 1 - execution condition
# 2 .. - loop carried dependencies
# Loop outputs:
# 0 .. loop_carried_dependencies_count - 1 - loop carried dependencies
# loop_carried_dependencies_count .. - scan outputs
# Body inputs:
# 0 - iteration number
# 1 - execution condition
# 2 .. - loop carried dependencies
# Body outputs:
# 0 - execution condition
# 1 .. loop_carried_dependencies_count - loop carried dependencies
# loop_carried_dependencies_count + 1 .. - scan outputs
# some of the inputs/outputs may not be connected but the normalization transformation will take care of it
# connection Loop body nodes with external input edges
next_loop_input_port_idx = sorted(loop_node.in_edges().keys())[-1] + 1
cur_graph = body_graph
for external_edges_subg in external_edges:
if 'parent_node' not in cur_graph.graph:
continue
cur_loop_node = cur_graph.graph['parent_node']
parent_graph = cur_loop_node.graph
for (src_node, src_port), body_node, tensor_name in external_edges_subg:
create_edge_with_attrs(parent_graph, tensor_name, src_node, src_port,
cur_loop_node.id, next_loop_input_port_idx)
Loop.connect_body_input(cur_loop_node, next_loop_input_port_idx, body_node)
next_loop_input_port_idx += 1
cur_graph = parent_graph
# mark current iteration input Parameter node
Loop.mark_current_iteration_parameter_node(loop_node, body_parameters[0])
# connect initial value for "execution condition" input of the loop
Loop.connect_body_input(loop_node, 1, body_parameters[1])
# add back edge with "execution condition"
Loop.add_back_edge(loop_node, body_parameters[1], body_results[0])
# mark "execution condition" Result node
Loop.mark_execution_condition_result_node(loop_node, body_results[0])
# connect initial value for "loop carried" dependencies variables
for idx in range(loop_carried_dependencies_count):
Loop.connect_body_input(loop_node, idx + 2, body_parameters[idx + 2])
# add back edge for "loop carried" dependencies variables
for idx in range(loop_carried_dependencies_count):
Loop.add_back_edge(loop_node, body_parameters[idx + 2], body_results[idx + 1])
# connect final value for "loop carried" dependencies variables
for idx in range(loop_carried_dependencies_count):
Loop.connect_body_output(loop_node, idx, body_results[idx + 1])
# connect "scan outputs" and mark axis for concatenation
for idx in range(loop_carried_dependencies_count, loop_carried_dependencies_count + scan_outputs_count):
Loop.connect_body_output(loop_node, idx, body_results[idx + 1], axis=0)
# run function to parse body nodes attributes similar to the main graph
extract_node_attrs(body_graph, lambda node: onnx_op_extractor(node, check_for_duplicates(onnx_op_extractors)))
return cls.enabled
