Merge pull request #625 from lucienwang1009/fix_rewriter

fix rewriter bugs

Author: zhijxu-MS

tests/test_backend.py

@@ -925,6 +925,22 @@ def test_leaky_relu_int(self):
                 self._run_test_case([_OUTPUT], {_INPUT: x_val})
                 tf.reset_default_graph()
 
+    @skip_caffe2_backend("fails on caffe2 with dim issue")
+    @check_onnxruntime_incompatibility("Mul")
+    def test_leaky_relu_with_dependency(self):
+        x_val = 1000 * np.random.random_sample([1000, 100]).astype(np.float32)
+        x = tf.placeholder(x_val.dtype, [None] * x_val.ndim, name=_TFINPUT)
+        # simulate leaky_relu
+        alpha = tf.constant(0.5)
+        y = alpha * x
+        x_ = tf.maximum(y, x)
+        dependency = y - 1
+
+        _ = tf.identity(x_, name=_TFOUTPUT)
+        _ = tf.identity(dependency, name=_TFOUTPUT1)
+        self._run_test_case([_OUTPUT, _OUTPUT1], {_INPUT: x_val})
+        tf.reset_default_graph()
+
     @skip_caffe2_backend("fails on caffe2 with dim issue")
     @check_onnxruntime_incompatibility("Mul")
     def test_leaky_relu_float(self):

tests/test_loops.py

@@ -12,7 +12,7 @@
 import tensorflow as tf
 
 from backend_test_base import Tf2OnnxBackendTestBase
-from common import unittest_main, check_tf_min_version
+from common import unittest_main, check_tf_min_version, check_onnxruntime_min_version
 
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test
@@ -160,6 +160,10 @@ def b(i, out_ta):
         output_names_with_port = ["i:0", "output_ta:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
 
+    @check_onnxruntime_min_version(
+        "0.5.0",
+        "disable this case due to onnxruntime loop issue: https://github.com/microsoft/onnxruntime/issues/1272"
+    )
     def test_while_loop_with_cond_init_false(self):
         i = tf.placeholder(tf.int32, (), name="input_1")
         inputs = tf.placeholder(tf.float32, (10,), name="input_2")

tf2onnx/graph.py

@@ -575,6 +575,18 @@ def reset_nodes(self, ops):
         self._dtypes = remained_dtypes
         self._output_shapes = remained_shapes
 
+    def check_integrity(self):
+        """
+        Check graph integrity. Every node's input needs to associate with a node.
+        Return broken outputs.
+        """
+        broken_outputs = set()
+        for node in self.get_nodes():
+            for inp in node.input:
+                if self.get_node_by_output(inp) is None:
+                    broken_outputs.add(inp)
+        return list(broken_outputs)
+
     def update_node_shape_dtype(self, node, override=False):
         """Try the best to infer shapes and dtypes for outputs of the node,
         by default, we respect TF shapes and dtypes.
@@ -1169,6 +1181,16 @@ def delete_unused_nodes(self, outputs_name):
                     body_graph.delete_unused_nodes(body_graph.outputs)
         self.reset_nodes(related_nodes)
 
+    def safe_remove_nodes(self, to_delete):
+        """Delete nodes in `to_delete` without third-party node consuming it."""
+        delete_set = set(to_delete)
+        for n in delete_set:
+            out_consumers = set()
+            for out in n.output:
+                out_consumers |= set(self.find_output_consumers(out))
+            if out_consumers.issubset(delete_set):
+                self.remove_node(n.name)
+
 
 class GraphUtil(object):
     """Utilities for Graph manipulation."""

tf2onnx/rewriter/leakyrelu_rewriter.py

@@ -42,6 +42,8 @@ def rewrite_leakyrelu(g, ops):
                                     shapes=[g.get_shape(max_node.output[0])], dtypes=[g.get_dtype(max_node.output[0])])
             ops.append(leakyrelu)
             g.replace_all_inputs(ops, max_node.output[0], leakyrelu.output[0])
+            to_delete = [max_node, mul_node]
+            g.safe_remove_nodes(to_delete)
 
     return ops
 

tf2onnx/rewriter/random_uniform.py

@@ -63,8 +63,7 @@ def rewrite_random_uniform_fold_const(g, ops):
         to_delete = list(set(match.get_nodes()))
         new_node = create_onnx_random_uniform_op(g, tmax, tmin, ru_op, output, to_delete)
         g.replace_all_inputs(ops, output.output[0], new_node.output[0])
-        for n in to_delete:
-            g.remove_node(n.name)
+        g.safe_remove_nodes(to_delete)
 
     return ops
 

tf2onnx/rewriter/thresholded_relu_rewriter.py

@@ -34,6 +34,7 @@ def rewrite_thresholded_relu(g, ops):
         greater_input_node = match.get_op('greater_input')
         mul_node = match.get_op("mul")
         mul_input_node = match.get_op('mul_input')
+        cast_node = match.get_op('cast')
 
         greater_input_edge_name = _find_edge_name_between_nodes(greater_input_node, greater_node)
         mul_input_edge_name = _find_edge_name_between_nodes(mul_input_node, mul_node)
@@ -42,6 +43,7 @@ def rewrite_thresholded_relu(g, ops):
             thresholded_relu = g.make_node("ThresholdedRelu", inputs=[mul_input_edge_name], attr={"alpha": theta},
                                            shapes=[g.get_shape(mul_node.output[0])],
                                            dtypes=[g.get_dtype(mul_node.output[0])])
-            ops.append(thresholded_relu)
             g.replace_all_inputs(ops, mul_node.output[0], thresholded_relu.output[0])
+            to_delete = [cast_node, mul_node]
+            g.safe_remove_nodes(to_delete)
     return ops

tf2onnx/tfonnx.py

@@ -141,9 +141,8 @@ def rewrite_transpose(g, ops):
         dims = [i for i in range(len(shape) - 1, -1, -1)]
         output.set_attr("perm", dims)
         g.remove_input(output, output.input[1])
-        for n in set(match.get_nodes()):
-            if n != output:
-                g.remove_node(n.name)
+        to_delete = [n for n in match.get_nodes() if n != output]
+        g.safe_remove_nodes(to_delete)
     return ops
 
 
@@ -175,8 +174,7 @@ def rewrite_random_normal(g, ops):
                                    attr={"shape": shape, "mean": mean, "scale": 1.0, "dtype": dtype})
 
         g.replace_all_inputs(ops, output.output[0], new_node.output[0])
-        for n in set(match.get_nodes()):
-            g.remove_node(n.name)
+        g.safe_remove_nodes(match.get_nodes())
     return ops
 
 
@@ -208,8 +206,7 @@ def rewrite_dropout(g, ops):
             dtypes=[g.get_dtype(inputs2.input[0])]
         )
         g.replace_all_inputs(ops, outputs.output[0], new_node.output[0])
-        for n in set(match.get_nodes()):
-            g.remove_node(n.name)
+        g.safe_remove_nodes(match.get_nodes())
 
     # remove dropout if its ratio is 1.0
     for node in g.get_nodes():
@@ -294,10 +291,8 @@ def rewrite_flatten(g, ops):
 
             g.set_shape(out_name, input_shape[:-2] + [new_dim])
             g.replace_all_inputs(ops, reshape_node.output[0], out_name)
-
-            for n in set(match.get_nodes()):
-                if n != input_node:
-                    g.remove_node(n.name)
+            to_delete = [n for n in match.get_nodes() if n != input_node]
+            g.safe_remove_nodes(to_delete)
 
     return ops
 
@@ -654,6 +649,14 @@ def run_rewriters(g, funcs, continue_on_error):
             else:
                 raise ex
 
+        if utils.is_debug_mode():
+            broken_outputs = g.check_integrity()
+            if broken_outputs:
+                logging.error(
+                    "After rewriter %s, graph breaks at outputs %s",
+                    func.__name__, broken_outputs
+                )
+
     if g.contained_graphs:
         for dict_val in g.contained_graphs.values():
             for attr_name, b_g in dict_val.items():