def protobuf2nx(graph, pb: tf_v1.GraphDef):
fill_graph_with_nodes(graph, pb.node, get_id=lambda pb: pb.name, get_attrs=protobuf_attrs)
if hasattr(graph, 'op_names_statistic'):
for node_name in graph.nodes:
node = Node(graph, node_name)
node_pb = node.soft_get('pb', None)
if node_pb is not None:
if hasattr(node_pb, 'op'):
graph.op_names_statistic[node_pb.op] += 1
# Create a library with auxiliary functions used in TensorFlow 2 operations
if hasattr(pb, 'library') and hasattr(pb.library, 'function'):
graph.graph['library'] = {}
for library_function in pb.library.function:
function_name = library_function.signature.name
graph.graph['library'][function_name] = {}
graph.graph['library'][function_name]['input_arg'] = library_function.signature.input_arg
graph.graph['library'][function_name]['output_arg'] = library_function.signature.output_arg
graph.graph['library'][function_name]['node_def'] = library_function.node_def
graph.graph['library'][function_name]['ret'] = library_function.ret
# initial order of nodes in the GraphDef. It is used to specify order in
# which merged nodes are added to the generated sub-graph GraphDef for the TensorFlow offload feature.
graph.graph['initial_nodes_order'] = [node.name for node in pb.node]
# Remove data dependency edges. This is needed for the TF offload case
for _, attrs in list(graph.nodes(data=True)):
pb = attrs['pb']
if '_class' in pb.attr:
index = 0
while index < len(pb.attr['_class'].list.s):
if re.match('^loc:@.*', pb.attr['_class'].list.s[index].decode('utf-8')):
del pb.attr['_class'].list.s[index]
else:
index = index + 1
def protobuf2nx(graph, pb: tf_v1.GraphDef):
fill_graph_with_nodes(graph, pb.node, get_id=lambda pb: pb.name, get_attrs=protobuf_attrs)
# initial order of nodes in the GraphDef. It is used to specify order in
# which merged nodes are added to the generated sub-graph GraphDef for the TensorFlow offload feature.
graph.graph['initial_nodes_order'] = [node.name for node in pb.node]
# Remove data dependency edges. This is needed for the TF offload case
for _, attrs in list(graph.nodes(data=True)):
pb = attrs['pb']
if '_class' in pb.attr:
index = 0
while index < len(pb.attr['_class'].list.s):
if re.match('^loc:@.*', pb.attr['_class'].list.s[index].decode('utf-8')):
del pb.attr['_class'].list.s[index]
else:
index = index + 1
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 protobuf2nx(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. '''
# graph = fill_graph_with_nodes(graph, pb.graph.node, get_id=node_id, get_attrs=protobuf_attrs)
# convert initializers to a NX graph for easier control of model consistency and to use it as a dictionary later
initializers = Graph()
fill_graph_with_nodes(initializers, pb.graph.initializer, get_id=lambda pb: pb.name, get_attrs=protobuf_attrs)
# maps a tensor name to a node produced it and the node port: str -> (node_id, node_port)
data_nodes_map = {}
# first go through all inputs and separate constant from placeholders
for inp in pb.graph.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):
# this is a constant
graph.add_node(name, kind='op', op='Const', pb=inp, pb_init=initializers.node[name]['pb'])
else:
# this is a placeholder
graph.add_node(name, kind='op', op='Parameter', pb=inp)
# add to a tensors map
assert not name in data_nodes_map, 'Inconsistency between data_nodes_map and graph.nodes'
data_nodes_map[name] = (name, 0)
output_ids = []
for outp in pb.graph.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 over all initializer and make sure that all of them are added to the graph
for initializer in initializers.nodes():
initializer_id = 'onnx_initializer_node_' + 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)
# Go through all nodes in the original model order (because data nodes are defined on-the-fly and order is
# important)
for node in pb.graph.node:
# create an NX node
id = graph.unique_id(node_id(node))
graph.add_node(id, pb=node, kind='op')
# 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': [(inp, 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': [(out, 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)