| 92 | |
| 93 | |
| 94 | class SimpleGraph: |
| 95 | def __init__(self, graph_def: tf.GraphDef): |
| 96 | self._nodes = [ |
| 97 | SimpleNode(name=n.name, op=n.op, inputs=list(n.input)) |
| 98 | for n in graph_def.node |
| 99 | ] |
| 100 | self._name2index = {n.name: i for i, n in enumerate(graph_def.node)} |
| 101 | self._graph_def = graph_def |
| 102 | |
| 103 | for node in graph_def.node: |
| 104 | for inp in node.input: |
| 105 | inp_node_name = tensor_name_to_node_name(inp) |
| 106 | inp_tensor_name = get_canonical_tensor_name(inp) |
| 107 | if inp_node_name not in self._name2index: |
| 108 | raise Exception(f"SimpleNode {node.name}: Unknown input node {inp}") |
| 109 | input_node = self._nodes[self._name2index[inp_node_name]] |
| 110 | # update input node"s [output_node, ..] list |
| 111 | input_node.output_nodes.append(node.name) |
| 112 | # update input node"s {tensor: output_node, ..} dictionary |
| 113 | # TODO: we are missing Graph final output node"s tensors, |
| 114 | # but it is not possible to inspect how many tensors inside |
| 115 | # it, therefore we currently ignore it. |
| 116 | if inp_tensor_name not in input_node.tensors: |
| 117 | input_node.tensors[inp_tensor_name] = [node.name] |
| 118 | else: |
| 119 | input_node.tensors[inp_tensor_name].append(node.name) |
| 120 | |
| 121 | @property |
| 122 | def num_nodes(self) -> int: |
| 123 | """Get total number of nodes in graph.""" |
| 124 | return len(self._nodes) |
| 125 | |
| 126 | @property |
| 127 | def nodes(self) -> List[SimpleNode]: |
| 128 | """Get all nodes in graph.""" |
| 129 | return self._nodes |
| 130 | |
| 131 | def name2index(self, name: str) -> int: |
| 132 | """Get index of node.""" |
| 133 | if name not in self._name2index: |
| 134 | error_msg = "Node {} not exists".format(name) |
| 135 | logging.error(error_msg) |
| 136 | raise Exception(error_msg) |
| 137 | return self._name2index[name] |
| 138 | |
| 139 | def node(self, idx: int) -> SimpleNode: |
| 140 | """Get node with given index.""" |
| 141 | if idx >= len(self._nodes): |
| 142 | error_msg = "Node index {} out of range".format(idx) |
| 143 | logging.error(error_msg) |
| 144 | raise Exception(error_msg) |
| 145 | return self._nodes[idx] |
| 146 | |
| 147 | def name2node(self, name: str) -> SimpleNode: |
| 148 | """Get node by name.""" |
| 149 | return self.node(self._name2index[name]) |
| 150 | |
| 151 | def input_nodes(self, blacklist: List = ["Const"]) -> List[str]: |
no outgoing calls
no test coverage detected