Extracts useful information from the graph and returns them.
(graph_def)
| 123 | |
| 124 | |
| 125 | def _extract_graph_summary(graph_def): |
| 126 | """Extracts useful information from the graph and returns them.""" |
| 127 | name_to_input_name = {} # Keyed by the dest node name. |
| 128 | name_to_node = {} # Keyed by node name. |
| 129 | |
| 130 | # Keeps track of node sequences. It is important to still output the |
| 131 | # operations in the original order. |
| 132 | name_to_seq_num = {} # Keyed by node name. |
| 133 | seq = 0 |
| 134 | for node in graph_def.node: |
| 135 | n = _node_name(node.name) |
| 136 | name_to_node[n] = node |
| 137 | name_to_input_name[n] = [_node_name(x) for x in node.input] |
| 138 | # Prevent colocated nodes from being lost. |
| 139 | if "_class" in node.attr: |
| 140 | for colocated_node_name in node.attr["_class"].list.s: |
| 141 | colocated_node_decoded = colocated_node_name.decode("utf-8") |
| 142 | if colocated_node_decoded.startswith("loc:@"): |
| 143 | name_to_input_name[n].append(colocated_node_decoded[5:]) |
| 144 | name_to_seq_num[n] = seq |
| 145 | seq += 1 |
| 146 | return name_to_input_name, name_to_node, name_to_seq_num |
| 147 | |
| 148 | |
| 149 | def _assert_nodes_are_present(name_to_node, nodes): |