Added error checking to dropout_rewriter to avoid incorrect matches

Author: TomWildenhain-Microsoft

tf2onnx/graph.py

@@ -76,6 +76,7 @@ def output(self, val):
             utils.make_sure(o not in self.graph._output_to_node_name, "output %s already in output mapping", o)
             self.graph._output_to_node_name[o] = self.name
 
+    # TODO(tomwildenhain): Rename to "input_nodes"
     @property
     def inputs(self):
         """Input node objects."""
@@ -151,6 +152,16 @@ def is_const(self):
         """Return True if node is a constant."""
         return self.type in ["Const", "ConstV2"]
 
+    def is_scalar(self):
+        """Return True if node is a constant with a scalar value."""
+        if not self.is_const():
+            return False
+        t = self.get_attr("value", default=None)
+        if t is None:
+            return False
+        t = numpy_helper.to_array(helper.get_attribute_value(t))
+        return t.shape == tuple()
+
     def is_graph_input(self):
         return self.type in ["Placeholder", "PlaceholderWithDefault", "PlaceholderV2"]
 
@@ -1318,6 +1329,7 @@ def safe_to_remove_nodes(self, to_delete):
                 safe_to_remove.append(n)
         return safe_to_remove
 
+    # TODO(tomwildenhain): Remove this function
     def safe_remove_nodes(self, to_delete):
         """Delete nodes in `to_delete` without third-party node consuming it."""
         delete_set = set(to_delete)
@@ -1328,6 +1340,20 @@ def safe_remove_nodes(self, to_delete):
             if out_consumers.issubset(delete_set):
                 self.remove_node(n.name)
 
+    def is_safe_to_remove_nodes(self, to_delete, outputs_to_ignore=None):
+        """Returns true if the outputs of all the nodes in to_delete have no third-party nodes consuming them"""
+        delete_set = set(to_delete)
+        outputs_to_ignore_set = set(outputs_to_ignore or [])
+        for n in delete_set:
+            out_consumers = set()
+            for out in n.output:
+                if out in outputs_to_ignore_set:
+                    continue
+                out_consumers |= set(self.find_output_consumers(out))
+            if not out_consumers.issubset(delete_set):
+                return False
+        return True
+
 
 class GraphUtil(object):
     """Utilities for Graph manipulation."""

tf2onnx/rewriter/dropout_rewriter.py

@@ -5,6 +5,7 @@
 tf2onnx.rewriter - rewrite tensorflow subgraph to onnx dropout op
 """
 
+import numpy as np
 from tf2onnx import utils
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 from tf2onnx import logging
@@ -52,39 +53,51 @@ def rewrite_dropout(g, ops):
         for match in match_results:
             input2 = match.get_op('input2')
             input3 = match.get_op('input3')
-            if input3.is_const():
-                ratio = input3.get_tensor_value()
-            else:
-                # If the ratio isn't constant, set it to 0
-                logger.warning("Dropout node has non-constant ratio. Using ratio=0.0")
-                ratio = 0.0
-            if input2.inputs[0].type == "RealDiv":
-                data = input2.input[1]
-            else:
-                data = input2.input[0]
-            # TODO(tomwildenhain): replace dropout node with identity if ratio is 0
             outputs = match.get_op('outputs')
+
+            if not input3.is_scalar():
+                logger.warning("Dropout pattern rooted at %s does not have a "
+                               "constant ratio and cannot be replaced.", outputs.name)
+                continue
+            ratio = input3.get_tensor_value()
+
+            if input2.inputs[0].is_scalar():
+                data = input2.inputs[1]
+                scaling_constant = input2.inputs[0].get_tensor_value()
+            elif input2.inputs[1].is_scalar():
+                data = input2.inputs[0]
+                scaling_constant = input2.inputs[1].get_tensor_value()
+            else:
+                logger.warning("Could not find scaling constant for dropout pattern rooted at %s. "
+                               "The pattern will not be replaced with an ONNX dropout node.", outputs.name)
+                continue
+
+            #The scaling constant should be 1/(1-ratio), otherwise this isn't truly a dropout node
+            if not np.allclose([1], [scaling_constant * (1 - ratio)]):
+                logger.warning("Scaling constant %f for dropout pattern rooted at %s is inconsistent with dropout "
+                               "ratio %f. The pattern will not be replaced with an ONNX dropout node.",
+                               scaling_constant, outputs.name, ratio)
+                continue
+
+            nodes_to_remove = [n for n in match.get_nodes() if n.name != input3.name]
+            if not g.is_safe_to_remove_nodes(nodes_to_remove, [outputs.output[0]]):
+                logger.warning("Nodes in dropout pattern rooted at %s cannot be removed because intermediate results "
+                               "of some nodes are referenced elsewhere in graph.", outputs.name)
+                continue
+
             op_name = utils.make_name("Dropout")
             out_name = utils.port_name(op_name)
             new_node = g.make_node(
                 "Dropout",
-                [data],
+                inputs=[data.output[0]],
                 outputs=[out_name],
                 name=op_name,
                 attr={"ratio": ratio},
-                shapes=[g.get_shape(input2.input[0])],
-                dtypes=[g.get_dtype(input2.input[0])]
+                shapes=[g.get_shape(data.output[0])],
+                dtypes=[g.get_dtype(data.output[0])]
             )
             g.replace_all_inputs(ops, outputs.output[0], new_node.output[0])
-            nodes_to_remove = []
-            for node in match.get_nodes():
-                if node.name != input3.name:
-                    nodes_to_remove.append(node)
-            if g.safe_to_remove_nodes(nodes_to_remove):
-                for n in nodes_to_remove:
-                    g.remove_node(n.name)
-            else:
-                logger.warning("Nodes replaced by dropout node cannot be removed because intermediate results are "
-                               "referenced elsewhere in graph")
+            for n in nodes_to_remove:
+                g.remove_node(n.name)
 
     return ops