Fixed dropout rewriter to properly read ratio

Author: TomWildenhain-Microsoft

tf2onnx/rewriter/dropout_rewriter.py

@@ -7,6 +7,9 @@
 
 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 +21,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 +31,7 @@ def rewrite_dropout(g, ops):
             OpTypePattern("Cast", inputs=[
                 OpTypePattern("GreaterEqual", inputs=[
                     OpTypePattern("RandomUniform|RandomUniformLike"),
-                    OpTypePattern(None, name="input3")
+                    OpTypePattern("*", name="input3")
                 ])
             ])
         ]),
@@ -37,7 +40,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"),
@@ -48,10 +51,18 @@ def rewrite_dropout(g, ops):
         match_results = list(matcher.match_ops(ops))
         for match in match_results:
             inputs2 = match.get_op('input2')
+            inputs3 = match.get_op('input3')
+            if inputs3.type == "Const":
+                ratio = inputs3.get_tensor_value()
+            else:
+                # If the ratio isn't constant, set it to 0
+                logger.error("Dropout node has non-constant ratio. Using ratio=0.0")
+                ratio = 0.0
             if inputs2.inputs[0].type == "RealDiv":
                 data = inputs2.input[1]
             else:
                 data = inputs2.input[0]
+            # TODO(tomwildenhain): replace dropout node with identity if ratio is 0
             outputs = match.get_op('outputs')
             op_name = utils.make_name("Dropout")
             out_name = utils.port_name(op_name)
@@ -60,17 +71,11 @@ def rewrite_dropout(g, ops):
                 [data],
                 outputs=[out_name],
                 name=op_name,
-                attr={"ratio": 1.0},
+                attr={"ratio": ratio},
                 shapes=[g.get_shape(inputs2.input[0])],
                 dtypes=[g.get_dtype(inputs2.input[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)
-
     return ops