Use the same syntax to replace an node input (2) + optimisation

Author: sdpython

tests/test_internals.py

@@ -139,7 +139,7 @@ def test_rewrite_subgraph(self):
             op_name = utils.make_name("ReplacedOp")
             out_name = utils.port_name(op_name)
             new_node = g.make_node("Sub", inputs=input_node.input, outputs=[out_name], name=op_name)
-            g.replace_all_inputs(ops, output_node.output[0], new_node.output[0])
+            g.replace_all_inputs(None, output_node.output[0], new_node.output[0])  # ops
             for n in set(match.get_nodes()):
                 g.remove_node(n.name)
         g.topological_sort(ops)

tf2onnx/graph.py

@@ -425,6 +425,8 @@ def __init__(self, nodes, output_shapes=None, dtypes=None, target=None, opset=No
         self._nodes = []
         self._nodes_by_name = {}
         self._output_to_node_name = {}
+        self._input_to_node_name = {}
+        self._input_to_graph = {}
         self.shapes = {}
         self.graph_name = graph_name or "tf2onnx"
         self._is_subgraph = is_subgraph
@@ -576,6 +578,9 @@ def make_node(self, op_type, inputs, attr=None, output_count=1, outputs=None, sk
 
         onnx_node = helper.make_node(op_type, inputs, outputs, name=name, domain=domain, **raw_attr)
 
+        for name2 in onnx_node.input:
+            self._register_input_name(name2, onnx_node)
+
         if op_type in ["If", "Loop", "Scan"]:
             # we force the op containing inner graphs not skipped during conversion.
             skip_conversion = False
@@ -613,6 +618,8 @@ def append_node(self, node):
             self._output_to_node_name[name] = node.name
             self.set_dtype(name, output_dtypes[i])
             self.set_shape(name, output_shapes[i])
+        for name in node.input:
+            self._register_input_name(name, node)
 
     def remove_node(self, node_name):
         """Remove node in current graph."""
@@ -633,6 +640,12 @@ def remove_node(self, node_name):
             if op_output in self._dtypes:
                 del self._dtypes[op_output]
 
+        for op_input in node.input:
+            if op_input not in self._input_to_node_name:
+                raise RuntimeError(
+                    "Input %r of node %r not found." % (op_input, node_name))
+            self._unregister_input_name(op_input, node)
+
         self._nodes.remove(node)
         node.graph = None
 
@@ -656,16 +669,32 @@ def reset_nodes(self, ops):
         self.contained_graphs = remained_sub_graphs
         self._nodes_by_name = {op.name: op for op in ops}
         self._output_to_node_name = {}
+        self._input_to_node_name = {}
         for op in ops:
             for op_output in op.output:
                 self._output_to_node_name[op_output] = op.name
+            if op.type == 'Placeholder':
+                inps = [op.name]
+            elif op.type == 'Const':
+                inps = [op.name]
+            else:
+                inps = op.input
+            for op_input in inps:
+                self._register_input_name(op_input, op)
 
         for n in self._order_sensitive_inputs:
             if n not in ops:
                 self._order_sensitive_inputs.remove(n)
         for o in self.outputs:
             if o not in self._output_to_node_name:
                 raise ValueError("graph output " + o + " not exist")
+        for i in self.inputs:
+            if i.name.startswith('Placeholder'):
+                continue
+            if i.name.startswith('keras_learning_phase'):
+                continue
+            if i.name not in self._input_to_node_name:
+                raise ValueError("graph input %r not exist in graph." % i.name)
 
         self._dtypes = remained_dtypes
         self._output_shapes = remained_shapes
@@ -782,6 +811,14 @@ def get_node_by_output_in_current_graph(self, output):
             ret = self._nodes_by_name.get(name)
         return ret
 
+    def get_node_by_input_in_current_graph(self, input_name):
+        """Get nodes by node input id."""
+        names = self._output_to_node_name.get(input_name)
+        ret = None
+        if name:
+            ret = [self._nodes_by_name.get(name) for name in names]
+        return ret
+
     def get_node_by_name(self, name):
         """Get node by name."""
         ret = self._nodes_by_name.get(name)
@@ -792,6 +829,8 @@ def set_node_by_name(self, node):
         self._nodes_by_name[node.name] = node
         for op_output in node.output:
             self._output_to_node_name[op_output] = node.name
+        for name in node.input:
+            self._register_input_name(name, node)
 
     def change_node_name(self, node, new_name):
         """Remove node in current graph."""
@@ -1138,8 +1177,7 @@ def dump_node_statistics(self):
 
         return op_cnt
 
-    @staticmethod
-    def remove_input(node, to_be_removed, i=None):
+    def remove_input(self, node, to_be_removed, i=None):
         """Remove input from Node.
         Args:
             node: the node we expect the input on
@@ -1149,11 +1187,16 @@ def remove_input(node, to_be_removed, i=None):
         assert isinstance(node, Node) and isinstance(to_be_removed, six.text_type)
         if i is not None:
             assert node.input[i] == to_be_removed
+            if node.input[i] in self._input_to_node_name:
+                to_ops = self._input_to_node_name[node.input[i]]
+                if node.name in to_ops:
+                    to_ops.remove(node.name)
             del node.input[i]
             return True
 
         for i2, name in enumerate(node.input):
             if name == to_be_removed:
+                self._unregister_input_name(node.input[i2], node)
                 del node.input[i2]
                 break
         # don't remove output from parent since others might depend on it
@@ -1205,43 +1248,90 @@ def insert_new_node_on_output(self, op_type, output_name, name, domain=None, **k
         new_output = port_name(name)
         new_node = self.make_node(op_type, [output_name], attr=kwargs, outputs=[new_output], name=name, domain=domain)
 
-        to_replace = [n for n in self.get_nodes() if n != new_node]
+        to_replace = [self.get_node_by_name(n) for n in self._input_to_node_name[output_name]]
+        to_replace = [n for n in to_replace if n != new_node]
         self.replace_all_inputs(to_replace, output_name, new_output)
         return new_node
 
     def find_output_consumers(self, output_name):
         """Find all nodes consuming a given output."""
+        if output_name in self._input_to_node_name:
+            ops = self._input_to_node_name[output_name]
+            ops = [self.get_node_by_name(n) for n in ops]
+        else:
+            ops = self.get_nodes()
         nodes = []
-        for node in self.get_nodes():
+        for node in ops:
+            if node is None:
+                continue
             if output_name in node.input:
                 nodes.append(node)
 
-            # find consumers in sub graphs
-            body_graphs = node.get_body_graphs()
-            if body_graphs:
-                for g in body_graphs.values():
-                    nodes.extend(g.find_output_consumers(output_name))
+        # find consumers in sub graphs
+        if output_name in self._input_to_graph:
+            for _, g in self._input_to_graph[output_name].items():
+                nodes.extend(g.find_output_consumers(output_name))
         return nodes
 
-    @staticmethod
-    def replace_all_inputs(ops, old_input, new_input):
-        """Replace all inputs pointing to old_input with new_input."""
+    def _register_input_name(self, input_name, node, only_graph=False):
+        "Register node taking a specific input."
+        if not only_graph:
+            if input_name not in self._input_to_node_name:
+                self._input_to_node_name[input_name] = set()
+            self._input_to_node_name[input_name].add(node.name)
+        if self.parent_graph is not None:
+            if input_name not in self.parent_graph._input_to_graph:
+                self.parent_graph._input_to_graph[input_name] = {}
+            self.parent_graph._input_to_graph[input_name][id(self)] = self
+            self.parent_graph._register_input_name(input_name, node, only_graph=True)
+
+    def _unregister_input_name(self, input_name, node, only_graph=False):
+        "Unregister node taking a specific input."
+        node_name = node.name
+        if not only_graph:
+            if node_name in self._input_to_node_name[input_name]:
+                if node_name in self._input_to_node_name[input_name]:
+                    self._input_to_node_name[input_name].remove(node_name)
+        if (self.parent_graph is not None and
+                input_name in self.parent_graph._input_to_graph and
+                id(self) in self.parent_graph._input_to_graph[input_name]):
+            del self.parent_graph._input_to_graph[input_name][id(self)]
+            self.parent_graph._unregister_input_name(input_name, node, only_graph=True)
+
+    def replace_all_inputs(self, ops, old_input, new_input):
+        """
+        Replace all inputs pointing to old_input with new_input.
+        *ops* is used if defined, otherwise _input_to_node_name
+        is used to determine the impacted nodes.
+        """
         if old_input == new_input:
             return
+        if new_input not in self._input_to_node_name:
+            self._input_to_node_name[new_input] = set()
+
+        if ops is not None:
+            keep_ops = True
+        elif old_input in self._input_to_node_name:
+            ops = [self.get_node_by_name(n) for n in self._input_to_node_name[old_input]]
+            keep_ops = False
+        else:
+            ops = []
+            keep_ops = False
 
         for node in ops:
+            assert node is not None
             if old_input in node.input and new_input in node.output:
                 raise RuntimeError("creating a circle in the graph is not allowed: " + node.name)
+            self._register_input_name(new_input, node)
 
             for i, input_name in enumerate(node.input):
                 if input_name == old_input:
-                    node.input[i] = new_input
+                    self.replace_input(node, node.input[i], new_input, i)
 
-            # modify references in sub graphs
-            body_graphs = node.get_body_graphs()
-            if body_graphs:
-                for g in body_graphs.values():
-                    g.replace_all_inputs(g.get_nodes(), old_input, new_input)
+        # modify references in sub graphs
+        if old_input in self._input_to_graph:
+            for _, g in self._input_to_graph[old_input].items():
+                g.replace_all_inputs(g.get_nodes() if keep_ops else None, old_input, new_input)
 
     def replace_input(self, node, old_input, new_input, i=None):
         """Replace one input in a node."""
@@ -1257,11 +1347,31 @@ def replace_input(self, node, old_input, new_input, i=None):
             is_replaced = True
         else:
             raise RuntimeError("Unable to replace input %r into %r for node %r." % (old_input, new_input, node.name))
+
+        to_ops = self._input_to_node_name.get(old_input, None)
+        if to_ops is not None:
+            if node.name in to_ops:
+                # A node may take twice the same entry.
+                to_ops.remove(node.name)
+
+        self._register_input_name(new_input, node)
         return is_replaced
 
     def replace_inputs(self, node, new_inputs):
         """Replace node inputs."""
         assert isinstance(node, Node) and isinstance(new_inputs, list)
+
+        for old_input in node.input:
+            to_ops = self._input_to_node_name.get(old_input, None)
+            if to_ops is not None and old_input in to_ops:
+                # To avoid issues when a node
+                # takes twice the same entry.
+                to_ops.remove(old_input)
+
+        for input_name in new_inputs:
+            assert isinstance(input_name, six.text_type)
+            self._register_input_name(input_name, node)
+
         node.input = new_inputs
         return True
 

tf2onnx/onnx_opset/controlflow.py

@@ -493,7 +493,7 @@ class TensorListStack:
     def version_7(cls, ctx, node, **kwargs):
         if node.inputs[0].is_while():
             ctx.remove_node(node.name)
-            ctx.replace_all_inputs(ctx.get_nodes(), node.output[0], node.input[0])
+            ctx.replace_all_inputs(None, node.output[0], node.input[0])  # ctx.get_nodes()
 
 
 @tf_op(["While", "StatelessWhile"])
@@ -583,7 +583,7 @@ def version_7(cls, ctx, node, **kwargs):
         for idx, n in reversed(to_remove):
             ctx.remove_node(n.name)
             # make the node output bad
-            ctx.replace_all_inputs(ctx.get_nodes(), n.output[0], "@@ALLOC")
+            ctx.replace_all_inputs(None, n.output[0], "@@ALLOC")  # ctx.get_nodes()
             del body.func_inputs[idx]
             del cond_graph.func_inputs[idx]
             del tf_while_inputs[idx]
@@ -619,7 +619,7 @@ def version_7(cls, ctx, node, **kwargs):
 
         # shift output consumers
         for k, v in output_map.items():
-            ctx.replace_all_inputs(ctx.get_nodes(), k, v)
+            ctx.replace_all_inputs(None, k, v)  # ctx.get_nodes()
 
         wire_while_body(ctx, body, loop_node.inputs, body_input_to_state_var, cond_input_to_state_var, output_shapes,
                         output_dtypes, body_name, node.name, cond_graph, tf_while_inputs, removed_scan_outputs)

tf2onnx/onnx_opset/math.py

@@ -695,4 +695,4 @@ def atan2(y, x):
             "Add", inputs=[atan_node.output[0], pi_part.output[0]],
             op_name_scope=node.name + 'all',
             shapes=[shape], dtypes=[onnx_dtype])
-        ctx.replace_all_inputs(ctx.get_nodes(), node.output[0], last_node.output[0])
+        ctx.replace_all_inputs(None, node.output[0], last_node.output[0])  # ctx.get_nodes()

tf2onnx/onnx_opset/misc.py

@@ -30,7 +30,7 @@ def version_1(cls, ctx, node, **kwargs):
             # if identity has a const as input, remove it
             input_name = node.input[0]
             output_name = node.output[0]
-            ctx.replace_all_inputs(ctx.get_nodes(), output_name, input_name)
+            ctx.replace_all_inputs(None, output_name, input_name)  # ctx.get_nodes()
             ctx.remove_node(node.name)
 
 

tf2onnx/onnx_opset/quantize.py

@@ -78,4 +78,4 @@ def version_10(cls, ctx, node, **kwargs):
             "DequantizeLinear", [new_node.output[0], pb_scale.name, zero_point.name],
             op_name_scope=node.name, attr={"axis": axis},
             shapes=[shape], dtypes=[dtype])
-        ctx.replace_all_inputs(ctx.get_nodes(), node.output[0], last_node.output[0])
+        ctx.replace_all_inputs(None, node.output[0], last_node.output[0])  # ctx.get_nodes()

tf2onnx/onnx_opset/rnn.py

@@ -153,7 +153,7 @@ def make_sigmoid(i, w, b):
         h_node = ctx.make_node("Mul", [co_node.output[0], o])
 
         def replace_output(old_output, new_output):
-            ctx.replace_all_inputs(ctx.get_nodes(), old_output, new_output)
+            ctx.replace_all_inputs(None, old_output, new_output)  # ctx.get_nodes()
             ctx.copy_dtype(old_output, new_output)
             ctx.copy_shape(old_output, new_output)
 

tf2onnx/onnx_opset/tensor.py

@@ -116,7 +116,7 @@ def version_1(cls, ctx, node, **kwargs):
             # if identity has a const as input, remove it
             input_name = node.input[0]
             output_name = node.output[0]
-            ctx.replace_all_inputs(ctx.get_nodes(), output_name, input_name)
+            ctx.replace_all_inputs(None, output_name, input_name)  # ctx.get_nodes()
             ctx.remove_node(node.name)
 
 
@@ -126,7 +126,7 @@ class IdentityN:
     def version_1(cls, ctx, node, **kwargs):
         ctx.remove_node(node.name)
         for input_name, output_name in zip(node.input, node.output):
-            ctx.replace_all_inputs(ctx.get_nodes(), output_name, input_name)
+            ctx.replace_all_inputs(None, output_name, input_name)  # ctx.get_nodes()
 
 
 @tf_op("Reshape")
@@ -1051,7 +1051,7 @@ def version_1(cls, ctx, node, **kwargs):
         # concat all unqueezes
         concat = ctx.make_node("Concat", inputs, op_name_scope=node.name, attr={"axis": axis},
                                shapes=shapes, dtypes=dtypes)
-        ctx.replace_all_inputs(ctx.get_nodes(), node.output[0], concat.output[0])
+        ctx.replace_all_inputs(None, node.output[0], concat.output[0])  # ctx.get_nodes()
 
 
 @tf_op("Unpack")