Merge pull request #1056 from xadupre/input

Use the same syntax to replace an node input

Author: xadupre

tf2onnx/custom_opsets/ms.py

@@ -85,8 +85,8 @@ def version_1(cls, ctx, node, **kwargs):
         # set node's attrs, Note: output_padding, group are left default.
         conv_dims_attr(node, "dilations")
         # set node's inputs from (output_shape, filter, input_tensor) to (input_tensor, filter, pads, Bias)
-        node.input[0] = node.input[2]
-        node.input[2] = pads.output[0]
+        ctx.replace_input(node, node.input[0], node.input[2], 0)
+        ctx.replace_input(node, node.input[2], pads.output[0], 2)
         conv_convert_inputs(ctx, node, with_kernel=True)
         node.attr.pop("data_format")
         node.attr.pop("padding")

tf2onnx/graph.py

@@ -56,6 +56,11 @@ def input(self):
 
     @input.setter
     def input(self, val):
+        # The setter can catch that all inputs are change
+        # but it cannot catch that one input is changed.
+        # That's method replace_input and replace_inputs must
+        # be used to change inputs to let the graph instance
+        # update its internal indices.
         self._input = copy.deepcopy(val)
 
     @property
@@ -1134,13 +1139,21 @@ def dump_node_statistics(self):
         return op_cnt
 
     @staticmethod
-    def remove_input(node, to_be_removed):
+    def remove_input(node, to_be_removed, input_index=None):
         """Remove input from Node.
         Args:
             node: the node we expect the input on
             to_be_removed: the node name we want to remove
+            input_index: if not None, index of the input to be removed,
+                the method is more efficient if *input_index* is specified,
+                otherwise, it has to look for every input named *old_input*.
         """
         assert isinstance(node, Node) and isinstance(to_be_removed, six.text_type)
+        if input_index is not None:
+            assert node.input[input_index] == to_be_removed
+            del node.input[input_index]
+            return True
+
         for i, name in enumerate(node.input):
             if name == to_be_removed:
                 del node.input[i]
@@ -1171,7 +1184,7 @@ def insert_new_node_on_input(self, node, op_type, input_name, name=None, domain=
         new_node = self.make_node(op_type, input_name, attr=kwargs, outputs=[new_output], name=name, domain=domain)
         for i, n in enumerate(node.input):
             if n == input_name[0]:
-                node.input[i] = new_output
+                self.replace_input(node, node.input[i], new_output, i)
                 break
         return new_node
 
@@ -1232,17 +1245,32 @@ def replace_all_inputs(ops, old_input, new_input):
                 for g in body_graphs.values():
                     g.replace_all_inputs(g.get_nodes(), old_input, new_input)
 
-    @staticmethod
-    def replace_input(node, old_input, new_input):
-        """Replace node."""
+    def replace_input(self, node, old_input, new_input, input_index=None):
+        """
+        Replace one input in a node.
+        The method is more efficient if *input_index* is specified.
+        Otherwise, it renames every output named *old_input*.
+        """
         assert isinstance(node, Node) and isinstance(old_input, six.text_type) and isinstance(new_input, six.text_type)
         is_replaced = False
-        for i, input_name in enumerate(node.input):
-            if input_name == old_input:
-                node.input[i] = new_input
-                is_replaced = True
+        if input_index is None:
+            for i, input_name in enumerate(node.input):
+                if input_name == old_input:
+                    node.input[i] = new_input
+                    is_replaced = True
+        elif node.input[input_index] == old_input:
+            node.input[input_index] = new_input
+            is_replaced = True
+        else:
+            raise RuntimeError("Unable to replace input %r into %r for node %r." % (old_input, new_input, node.name))
         return is_replaced
 
+    def replace_inputs(self, node, new_inputs):
+        """Replace node inputs."""
+        assert isinstance(node, Node) and isinstance(new_inputs, list)
+        node.input = new_inputs
+        return True
+
     def _extract_sub_graph_nodes(self, dest_node, input_checker=None):
         """Return nodes of subgraph ending with dest_node.
         Args:

tf2onnx/onnx_opset/common.py

@@ -40,8 +40,8 @@ def version_1(cls, ctx, node, **kwargs):
                         shape1 = node.inputs[1].scalar_to_dim1()
             if shape0 and shape1 and len(shape0) < len(shape1) and node.type in ["Mul", "Add"]:
                 tmp = node.input[0]
-                node.input[0] = node.input[1]
-                node.input[1] = tmp
+                ctx.replace_input(node, node.input[0], node.input[1], 0)
+                ctx.replace_input(node, node.input[1], tmp, 1)
         else:
             node.set_attr("broadcast", 0)
 
@@ -65,5 +65,5 @@ def version_6(cls, ctx, node, **kwargs):
                         shape1 = node.inputs[1].scalar_to_dim1()
             if shape0 and shape1 and len(shape0) < len(shape1) and node.type in ["Mul", "Add"]:
                 tmp = node.input[0]
-                node.input[0] = node.input[1]
-                node.input[1] = tmp
+                ctx.replace_input(node, node.input[0], node.input[1], 0)
+                ctx.replace_input(node, node.input[1], tmp, 1)

tf2onnx/onnx_opset/controlflow.py

@@ -374,7 +374,7 @@ def version_9(cls, ctx, node, **kwargs):
             broadcast_shape = [cond_shape[0]] + [1] * (input_rank - 1)
             shape_const = ctx.make_const(utils.make_name(node.name), np.array(broadcast_shape, dtype=np.int64))
             reshape = ctx.make_node("Reshape", [node.input[0], shape_const.output[0]])
-            ctx.replace_input(node, node.input[0], reshape.output[0])
+            ctx.replace_input(node, node.input[0], reshape.output[0], 0)
 
 
 @tf_op("Where")
@@ -456,7 +456,7 @@ class TensorListGetItem:
     def version_7(cls, ctx, node, **kwargs):
         ctx.ta_reads.append(node.input[0])
         node.type = "Gather"
-        node.input = [node.input[0], node.input[1]]
+        ctx.replace_inputs(node, [node.input[0], node.input[1]])
         ctx.insert_new_node_on_input(node, "Unsqueeze", node.input[1], name=node.child_name(), axes=[0])
         ctx.insert_new_node_on_output("Squeeze", node.output[0], name=node.child_name(), axes=[0])
 
@@ -531,7 +531,7 @@ def version_7(cls, ctx, node, **kwargs):
         else:
             maximum_iterations_name = utils.make_name(node.inputs[1].name)
         ctx.make_const(maximum_iterations_name, np.array(maximum_iterations, dtype=np.int64))
-        node.input[1] = maximum_iterations_name
+        ctx.replace_input(node, node.input[1], maximum_iterations_name, 1)
 
         cond_name = node.get_attr_str("cond")
         cond_graph = find_function(cond_name)
@@ -642,7 +642,7 @@ def wire_while_body(parent_g, g, loop_node_inputs, body_input_to_state_var, cond
     for n in g.inputs:
         if n.output[0] in body_input_to_state_var:
             n.type = "Identity"
-            n.input = [body_input_to_state_var[n.output[0]]]
+            g.replace_inputs(n, [body_input_to_state_var[n.output[0]]])
 
     # onnx will pass in cond as argument
     cond_node = g.make_node("Placeholder", [], name=utils.make_name("cond"),
@@ -673,7 +673,7 @@ def wire_while_body(parent_g, g, loop_node_inputs, body_input_to_state_var, cond
             node.type = "Identity"
             g.set_shape(node.output[0], g.get_shape(node.input[2]))
             g.set_dtype(node.output[0], g.get_dtype(node.input[2]))
-            node.input = [node.input[2]]
+            g.replace_inputs(node, [node.input[2]])
             scan_outputs.append(node.output[0])
 
     if len(scan_outputs) != len(removed_scan_outputs):
@@ -729,7 +729,7 @@ def wire_if_branch(parent_g, g, inputs, output_shapes, output_dtypes, scope, par
     for node in g.inputs:
         parent_name = binding.get(node.output[0])
         if parent_name and parent_name != "@@ALLOC":
-            node.input = [parent_name]
+            g.replace_inputs(node, [parent_name])
             node.type = "Identity"
         else:
             to_remove.append(node)
@@ -753,7 +753,7 @@ def inline_subgraph(parent, g, scope, binding):
     for node in g.inputs:
         parent_name = binding.get(node.output[0])
         if parent_name and parent_name != "@@ALLOC":
-            node.input = [parent_name]
+            g.replace_inputs(node, [parent_name])
             node.type = "Identity"
         else:
             to_remove.append(node)

tf2onnx/onnx_opset/generator.py

@@ -40,7 +40,7 @@ def version_1(cls, ctx, node, **kwargs):
         node.set_attr("seed", float(seed.f))
         if len(node.input) > 0:
             shape = node.inputs[0].get_tensor_value()
-            ctx.remove_input(node, node.input[0])
+            ctx.remove_input(node, node.input[0], 0)
             node.set_attr("shape", shape)
             ctx.set_shape(node.output[0], shape)
 
@@ -53,7 +53,7 @@ def version_9(cls, ctx, node, **kwargs):
             node.set_attr("seed", float(seed.f))
             cast_node = ctx.make_node("Cast", node.input, attr={'to': onnx_pb.TensorProto.INT64})
             const_node = ctx.make_node("ConstantOfShape", cast_node.output)
-            node.input = const_node.output
+            ctx.replace_inputs(node, const_node.output.copy())
             node.type = node.type + 'Like'
 
 
@@ -103,8 +103,8 @@ def version_7(cls, ctx, node, **kwargs):
         ctx.set_shape(tile_shape_int64.output[0], fill_shape)
 
         tmp = node.input[0]
-        node.input[0] = node.input[1]
-        node.input[1] = tmp
+        ctx.replace_input(node, node.input[0], node.input[1], 0)
+        ctx.replace_input(node, node.input[1], tmp, 1)
         node.type = "Tile"
         ctx.set_dtype(node.output[0], new_dtype)
 
@@ -128,13 +128,13 @@ def version_9(cls, ctx, node, **kwargs):
         value = np.array([node.inputs[1].get_tensor_value()]).astype(utils.map_onnx_to_numpy_type(dtype))
         value_proto = numpy_helper.from_array(value)
         node.set_attr("value", value_proto)
-        del node.input[1]
+        ctx.remove_input(node, node.input[1], 1)
 
     @classmethod
     def version_11(cls, ctx, node, **kwargs):
         # cls.version_7(ctx, node, **kwargs)
         node.type = "Expand"
-        node.input = [node.input[1], node.input[0]]
+        ctx.replace_inputs(node, [node.input[1], node.input[0]])
         # cast shape to int64 if needed
         if ctx.get_dtype(node.input[1]) != onnx_pb.TensorProto.INT64:
             ctx.insert_new_node_on_input(node, "Cast", node.input[1], to=onnx_pb.TensorProto.INT64)
@@ -156,7 +156,7 @@ def version_7(cls, ctx, node, **kwargs):
             output_dtype = onnx_pb.TensorProto.INT32
         node.set_attr("dtype", output_dtype)
         node.set_attr("sample_size", sample_size)
-        ctx.remove_input(node, node.input[1])
+        ctx.remove_input(node, node.input[1], 1)
 
 
 @tf_op("ZerosLike")

tf2onnx/onnx_opset/math.py

@@ -123,7 +123,7 @@ def make_min_or_max_op(ctx, op_type, inputs, outputs,
             # use add as 'broadcast' op
             add_node = ctx.make_node("Add", [input_node.output[0], sub_node.output[0]],
                                      op_name_scope=input_node.name)
-            node.input[i] = add_node.output[0]
+            ctx.replace_input(node, node.input[i], add_node.output[0], i)
     return final_node
 
 
@@ -293,7 +293,7 @@ def version_1(cls, ctx, node, **kwargs):
             # workaround a bug in caffe2 pre Feb2018, pow(a, b) becomes np.exp(np.log(a) * b)
             node.type = "Log"
             b = node.input[1]
-            ctx.remove_input(node, node.input[1])
+            ctx.remove_input(node, node.input[1], 1)
             op_name = utils.make_name(node.name)
             mul_op = ctx.insert_new_node_on_output("Mul", node.output[0], name=op_name)
             mul_op.input.append(b)

tf2onnx/onnx_opset/nn.py

@@ -130,7 +130,7 @@ def conv_convert_inputs(ctx, node, with_kernel=False, new_kernel_shape=None,
                 ctx.make_const(shape_name, np.array(new_kernel_shape, dtype=np.int64))
 
                 reshape = ctx.make_node("Reshape", [kernel_name, shape_name])
-                ctx.replace_input(node, kernel_name, reshape.output[0])
+                ctx.replace_input(node, kernel_name, reshape.output[0], 1)
 
                 reshape.skip_conversion = True
 
@@ -453,11 +453,11 @@ def version_1(cls, ctx, node, **kwargs):
         conv_dims_attr(node, "dilations", spatial=spatial)
 
         # remove output_shapes input
-        ctx.remove_input(node, node.input[0])
+        ctx.remove_input(node, node.input[0], 0)
         # swap data and kernel
         t = node.input[0]
-        node.input[0] = node.input[1]
-        node.input[1] = t
+        ctx.replace_input(node, node.input[0], node.input[1], 0)
+        ctx.replace_input(node, node.input[1], t, 1)
 
         conv_convert_inputs(ctx, node, with_kernel=True, spatial=spatial)
 
@@ -539,8 +539,8 @@ def _convert(cls, ctx, node, **kwargs):
         else:
             kernel_shape_tf = node.inputs[1].get_tensor_value()
             strides_tf = node.inputs[2].get_tensor_value()
-            ctx.remove_input(node, node.input[2])
-            ctx.remove_input(node, node.input[1])
+            ctx.remove_input(node, node.input[2], 2)
+            ctx.remove_input(node, node.input[1], 1)
 
         kernel_shape_hw = parse_dims_attr(node, kernel_shape_tf, spatial)
         strides_hw = parse_dims_attr(node, strides_tf, spatial)
@@ -605,7 +605,7 @@ def version_7(cls, ctx, node, **kwargs):
                 ctx.make_const(shape_name, np.array(new_broadcast_shape, dtype=np.int64))
                 op_name = node.input[1]
                 reshape_node = ctx.make_node("Reshape", [op_name, shape_name])
-                ctx.replace_input(node, op_name, reshape_node.output[0])
+                ctx.replace_input(node, op_name, reshape_node.output[0], 1)
                 ctx.set_shape(reshape_node.output[0], new_broadcast_shape)
 
 
@@ -629,9 +629,9 @@ def version_1(cls, ctx, node, **kwargs):
         if mode in [None, "constant"] and len(node.input) == 3:
             const_val = node.inputs[2].get_tensor_value()
             node.set_attr("value", const_val)
-            ctx.remove_input(node, node.input[2])
+            ctx.remove_input(node, node.input[2], 2)
 
-        ctx.remove_input(node, node.input[1])
+        ctx.remove_input(node, node.input[1], 1)
         node.set_attr("pads", paddings)
 
         origin_dtype = ctx.get_dtype(node.output[0])
@@ -707,14 +707,14 @@ def version_6(cls, ctx, node, **kwargs):
                                       dtype=val_type)
             new_mean_node_name = utils.make_name(node.name)
             ctx.make_const(new_mean_node_name, new_mean_value)
-            node.input[3] = new_mean_node_name
+            ctx.replace_input(node, node.input[3], new_mean_node_name, 3)
 
         if var_shape != scale_shape:
             new_var_value = np.array(np.resize(node.inputs[4].get_tensor_value(as_list=False), scale_shape),
                                      dtype=val_type)
             new_val_node_name = utils.make_name(node.name)
             ctx.make_const(new_val_node_name, new_var_value)
-            node.input[4] = new_val_node_name
+            ctx.replace_input(node, node.input[4], new_val_node_name, 4)
 
     @classmethod
     def version_9(cls, ctx, node, **kwargs):
@@ -865,7 +865,7 @@ def version_7(cls, ctx, node, **kwargs):
         scaler = [1., 1., float(nh) / h, float(nw) / w]
         node.set_attr("scales", scaler)
         node.set_attr("mode", mode)
-        ctx.remove_input(node, node.input[1])
+        ctx.remove_input(node, node.input[1], 1)
         node.data_format = "NHWC"
         conv_convert_inputs(ctx, node, with_kernel=False)
 

tf2onnx/onnx_opset/quantize.py

@@ -70,7 +70,7 @@ def version_10(cls, ctx, node, **kwargs):
             op_name_scope=node.name, attr={"axis": axis},
             shapes=[shape], dtypes=[idtype])
         output_name = new_node.output[0]
-        node.input[0] = output_name
+        ctx.replace_input(node, node.input[0], output_name, 0)
 
         ctx.remove_node(node.name)
 

tf2onnx/onnx_opset/reduction.py

@@ -43,7 +43,7 @@ def version_1(cls, ctx, node, **kwargs):
             axes = [val + input_rank if val < 0 else val for val in axes]
 
         node.set_attr("axes", axes)
-        ctx.remove_input(node, node.input[1])
+        ctx.remove_input(node, node.input[1], 1)
         keep_dims = node.get_attr("keep_dims")
         if keep_dims:
             del node.attr['keep_dims']
@@ -82,7 +82,7 @@ def version_1(cls, ctx, node, **kwargs):
 
         node.set_attr("axis", axis)
         node.set_attr("keepdims", 0)
-        ctx.remove_input(node, node.input[1])
+        ctx.remove_input(node, node.input[1], 1)
 
     @classmethod
     def version_11(cls, ctx, node, **kwargs):