align dropout with tf 1.14

Author: lucienwang1009

tests/common.py

@@ -297,6 +297,12 @@ def validate_const_node(node, expected_val):
 def group_nodes_by_type(graph):
     res = defaultdict(list)
     for node in graph.get_nodes():
+        attr_body_graphs = node.get_body_graphs()
+        if attr_body_graphs:
+            for _, body_graph in attr_body_graphs.items():
+                body_graph_res = group_nodes_by_type(body_graph)
+                for k, v in body_graph_res.items():
+                    res[k].extend(v)
         res[node.type].append(node)
     return res
 

tests/test_backend.py

@@ -441,7 +441,9 @@ def test_dropout(self):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["output:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port,
+                           graph_validator=lambda g: (check_op_count(g, "RandomUniform", 0) and
+                                                      check_op_count(g, "RandomUniformLike", 0)))
 
     def test_nn_dropout(self):
         keep_prob = tf.placeholder_with_default(1., (), "keep_prob")
@@ -458,7 +460,10 @@ def test_nn_dropout(self):
         output_names_with_port = ["output:0"]
         # when constant_fold is enabled, PlaceholderWithDefault will be folded into either a const or a placeholder.
         # here we set it False to test PlaceholderWithDefault bug: https://github.com/onnx/tensorflow-onnx/pull/446
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, constant_fold=False)
+        # Dropout with ratio 1.0 will be optimized so that only one Identity is left
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, constant_fold=False,
+                           graph_validator=lambda g: (check_op_count(g, "RandomUniform", 0) and
+                                                      check_op_count(g, "RandomUniformLike", 0)))
 
     @check_tf_min_version("1.13")
     def test_nn_dropout_with_rate(self):
@@ -474,7 +479,9 @@ def test_nn_dropout_with_rate(self):
         feed_dict = {"input_1:0": x_val}
         input_names_with_port = ["input_1:0"]
         output_names_with_port = ["output:0"]
-        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, constant_fold=False)
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, constant_fold=False,
+                           graph_validator=lambda g: (check_op_count(g, "RandomUniform", 0) and
+                                                      check_op_count(g, "RandomUniformLike", 0)))
 
     def test_conv2d_with_input_transpose(self):
         x_shape = [2, 32, 32, 3]

tests/test_graph.py

@@ -22,7 +22,7 @@
 from tf2onnx.handler import tf_op
 
 from backend_test_base import Tf2OnnxBackendTestBase
-from common import unittest_main, check_tf_min_version, check_tf_max_version
+from common import unittest_main
 
 
 # pylint: disable=missing-docstring,unused-argument,unused-variable
@@ -139,7 +139,6 @@ def test_randomnormal(self):
                        'RandomNormal__2:0 -> output }'
             self.assertEqual(expected, actual)
 
-    @check_tf_max_version("1.12")
     def test_dropout(self):
         with tf.Session() as sess:
             x1 = tf.placeholder(tf.float32, [2, 3], name="input1")
@@ -150,35 +149,15 @@ def test_dropout(self):
             x_ = tf.identity(x_, name="output1")
             x_ = tf.identity(x_, name="output2")
             _ = tf.identity(x_, name="output")
-            g = process_tf_graph(sess.graph, opset=self.config.opset)
+            # feed output_names in order to remove unused nodes.
+            g = process_tf_graph(sess.graph, opset=self.config.opset, output_names=["output:0"])
+            utils.save_protobuf("./test.onnx", g.make_model("test"))
             actual = onnx_to_graphviz(g)
             expected = 'digraph { prob [op_type=Placeholder shape="[]"] input2 [op_type=Placeholder shape="[1, 3]"] ' \
                        'input1 [op_type=Placeholder shape="[2, 3]"] Add [op_type=Add] output1 [op_type=Identity] ' \
-                       'output2 [op_type=Identity] output [op_type=Identity] input1:0 -> Add input2:0 -> Add ' \
-                       'Add:0 -> output1 output1:0 -> output2 output2:0 -> output }'
-            self.assertEqual(expected, actual)
-
-    @check_tf_min_version("1.13")
-    def test_dropout_2(self):
-        with tf.Session() as sess:
-            x1 = tf.placeholder(tf.float32, [2, 3], name="input1")
-            x2 = tf.placeholder(tf.float32, [1, 3], name="input2")
-            prop = tf.placeholder(tf.float32, (), name="prob")
-            x_ = tf.add(x1, x2)
-            x_ = tf.nn.dropout(x_, prop)
-            x_ = tf.identity(x_, name="output1")
-            x_ = tf.identity(x_, name="output2")
-            _ = tf.identity(x_, name="output")
-            g = process_tf_graph(sess.graph, opset=self.config.opset)
-            actual = onnx_to_graphviz(g)
-            expected = 'digraph { "sub/x" [op_type=Const] prob [op_type=Placeholder shape="[]"] ' \
-                       'sub [op_type=Sub] input2 [op_type=Placeholder shape="[1, 3]"] ' \
-                       'input1 [op_type=Placeholder shape="[2, 3]"] "dropout/sub/x" [op_type=Const] ' \
-                       '"dropout/sub" [op_type=Sub] Add [op_type=Add] output1 [op_type=Identity] ' \
-                       'output2 [op_type=Identity] output [op_type=Identity] "sub/x":0 -> sub ' \
-                       'prob:0 -> sub "dropout/sub/x":0 -> "dropout/sub" sub:0 -> "dropout/sub" ' \
-                       'input1:0 -> Add input2:0 -> Add Add:0 -> output1 output1:0 -> output2 ' \
-                       'output2:0 -> output }'
+                       'output2 [op_type=Identity] output [op_type=Identity] output_graph_outputs_Identity__3 ' \
+                       '[op_type=Identity] input1:0 -> Add input2:0 -> Add Add:0 -> output1 output1:0 -> output2 ' \
+                       'output2:0 -> output output_raw_output___2:0 -> output_graph_outputs_Identity__3 }'
             self.assertEqual(expected, actual)
 
     def test_add(self):

tf2onnx/tfonnx.py

@@ -179,7 +179,7 @@ def rewrite_random_normal(g, ops):
 
 
 def rewrite_dropout(g, ops):
-    pattern = \
+    patterns = [
         OpTypePattern('Mul', name='outputs', inputs=[
             OpTypePattern('RealDiv', name="input2"),
             OpTypePattern('Floor', inputs=[
@@ -188,25 +188,36 @@ def rewrite_dropout(g, ops):
                     OpTypePattern('RandomUniform|RandomUniformLike'),
                 ])
             ]),
+        ]),
+        OpTypePattern("Mul", name="outputs", inputs=[
+            OpTypePattern("Mul", name="input2"),
+            OpTypePattern("Cast", inputs=[
+                OpTypePattern("GreaterEqual", inputs=[
+                    OpTypePattern("RandomUniform|RandomUniformLike"),
+                    OpTypePattern(None, name="input3")
+                ])
+            ])
         ])
-    matcher = GraphMatcher(pattern)
-    match_results = list(matcher.match_ops(ops))
-    for match in match_results:
-        inputs2 = match.get_op('input2')
-        outputs = match.get_op('outputs')
-        op_name = utils.make_name("Dropout")
-        out_name = port_name(op_name)
-        new_node = g.make_node(
-            "Dropout",
-            [inputs2.input[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])]
-        )
-        g.replace_all_inputs(ops, outputs.output[0], new_node.output[0])
-        g.safe_remove_nodes(match.get_nodes())
+    ]
+    for pattern in patterns:
+        matcher = GraphMatcher(pattern)
+        match_results = list(matcher.match_ops(ops))
+        for match in match_results:
+            inputs2 = match.get_op('input2')
+            outputs = match.get_op('outputs')
+            op_name = utils.make_name("Dropout")
+            out_name = port_name(op_name)
+            new_node = g.make_node(
+                "Dropout",
+                [inputs2.input[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])]
+            )
+            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():