delete nodes without dependency in each rewriter. check graph integrity after rewriting in verbose mode.

Author: lucienwang1009

tf2onnx/constants.py

@@ -37,3 +37,6 @@
 
 # Environment variables
 ENV_TF2ONNX_DEBUG_MODE = "TF2ONNX_DEBUG_MODE"
+
+# Logging level
+VERBOSE = 15

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.
@@ -591,7 +603,7 @@ def update_node_shape_dtype(self, node, override=False):
         initializers = []
         for i, inp in enumerate(node.inputs):
             if inp is None:
-                if logger.isEnabledFor(logging.INFO):
+                if logger.isEnabledFor(constants.VERBOSE):
                     logger.warning(
                         "[%s] infer a inexistent node: [%s], please check the code",
                         node.name, node.input[i]
@@ -1169,6 +1181,20 @@ def delete_unused_nodes(self, outputs_name):
                     body_graph.delete_unused_nodes(body_graph.outputs)
         self.reset_nodes(related_nodes)
 
+    def delete_nodes_without_dependency(self, to_delete):
+        """Delete nodes in `to_delete` without third-party dependency."""
+        for n in to_delete:
+            can_delete = True
+            for out in n.output:
+                if not can_delete:
+                    break
+                for consumer in self.find_output_consumers(out):
+                    if consumer not in to_delete:
+                        can_delete = False
+                        break
+            if can_delete:
+                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.delete_nodes_without_dependency(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.delete_nodes_without_dependency(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.delete_nodes_without_dependency(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 set(match.get_nodes()) if n != output]
+        g.delete_nodes_without_dependency(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.delete_nodes_without_dependency(set(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.delete_nodes_without_dependency(set(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 set(match.get_nodes()) if n != input_node]
+            g.delete_nodes_without_dependency(to_delete)
 
     return ops
 
@@ -654,6 +649,14 @@ def run_rewriters(g, funcs, continue_on_error):
             else:
                 raise ex
 
+        if logger.isEnabledFor(constants.VERBOSE):
+            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():

tf2onnx/verbose_logging.py

@@ -15,14 +15,12 @@
 
 from . import constants
 
-VERBOSE = 15
-
-_logging.addLevelName(VERBOSE, "VERBOSE")
+_logging.addLevelName(constants.VERBOSE, "VERBOSE")
 
 
 def _verbose(self, message, *args, **kwargs):
-    if self.isEnabledFor(VERBOSE):
-        self._log(VERBOSE, message, args, **kwargs)  # pylint: disable=protected-access
+    if self.isEnabledFor(constants.VERBOSE):
+        self._log(constants.VERBOSE, message, args, **kwargs)  # pylint: disable=protected-access
 
 
 def getLogger(name=None):  # pylint: disable=invalid-name, function-redefined
@@ -47,7 +45,7 @@ def basicConfig(**kwargs):  # pylint: disable=invalid-name, function-redefined
     set_tf_verbosity(_logging.getLogger().getEffectiveLevel())
 
 
-_LOG_LEVELS = [FATAL, ERROR, WARNING, INFO, VERBOSE, DEBUG]
+_LOG_LEVELS = [FATAL, ERROR, WARNING, INFO, constants.VERBOSE, DEBUG]
 
 
 def get_verbosity_level(verbosity, base_level=INFO):