Merge pull request #995 from onnx/tom/FixDropoutRewriter

Fixed dropout rewriter to properly read ratio

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,8 +5,12 @@
 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
+
+logger = logging.getLogger(__name__)
 
 
 # pylint: disable=missing-docstring
@@ -18,7 +22,7 @@ def rewrite_dropout(g, ops):
             OpTypePattern('RealDiv', name="input2"),
             OpTypePattern('Floor', inputs=[
                 OpTypePattern('Add', inputs=[
-                    OpTypePattern(None, name="input3"),
+                    OpTypePattern("*", name="input3"),
                     OpTypePattern('RandomUniform|RandomUniformLike'),
                 ])
             ]),
@@ -28,7 +32,7 @@ def rewrite_dropout(g, ops):
             OpTypePattern("Cast", inputs=[
                 OpTypePattern("GreaterEqual", inputs=[
                     OpTypePattern("RandomUniform|RandomUniformLike"),
-                    OpTypePattern(None, name="input3")
+                    OpTypePattern("*", name="input3")
                 ])
             ])
         ]),
@@ -37,7 +41,7 @@ def rewrite_dropout(g, ops):
             OpTypePattern("Cast", inputs=[
                 OpTypePattern("GreaterEqual", inputs=[
                     OpTypePattern("RandomUniform|RandomUniformLike"),
-                    OpTypePattern(None, name="input3")
+                    OpTypePattern("*", name="input3")
                 ])
             ]),
             OpTypePattern("Mul", name="input2"),
@@ -47,30 +51,53 @@ def rewrite_dropout(g, ops):
         matcher = GraphMatcher(pattern, allow_reorder=True)
         match_results = list(matcher.match_ops(ops))
         for match in match_results:
-            inputs2 = match.get_op('input2')
-            if inputs2.inputs[0].type == "RealDiv":
-                data = inputs2.input[1]
-            else:
-                data = inputs2.input[0]
+            input2 = match.get_op('input2')
+            input3 = match.get_op('input3')
             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": 1.0},
-                shapes=[g.get_shape(inputs2.input[0])],
-                dtypes=[g.get_dtype(inputs2.input[0])]
+                attr={"ratio": ratio},
+                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])
-            g.safe_remove_nodes(match.get_nodes())
-
-    # remove dropout if its ratio is 1.0
-    for node in g.get_nodes():
-        if node.type == "Dropout" and node.get_attr("ratio").f == 1.0:
-            g.replace_all_inputs(g.get_nodes(), node.output[0], node.input[0])
-            g.remove_node(node.name)
+            for n in nodes_to_remove:
+                g.remove_node(n.name)
 
     return ops