Merge pull request #446 from lucienwang1009/dropout_fix

fix dropout bug

Author: nbcsm

tests/backend_test_base.py

@@ -130,7 +130,7 @@ def run_test_case(self, feed_dict, input_names_with_port, output_names_with_port
         graph_def = tf_optimize(input_names_with_port, output_names_with_port,
                                 sess.graph_def, constant_fold)
 
-        if self.config.is_debug_mode and constant_fold:
+        if self.config.is_debug_mode:
             model_path = os.path.join(self.test_data_directory, self._testMethodName + "_after_tf_optimize.pb")
             utils.save_protobuf(model_path, graph_def)
             self.log.debug("created file  %s", model_path)

tests/run_pretrained_models.py

@@ -243,6 +243,8 @@ def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=No
                 inputs[k] = v
 
         graph_def = tf2onnx.tfonnx.tf_optimize(inputs.keys(), self.output_names, graph_def, fold_const)
+        if debug:
+            utils.save_protobuf(os.path.join(TEMP_DIR, name + "_after_tf_optimize.pb"), graph_def)
         shape_override = {}
         g = tf.import_graph_def(graph_def, name='')
         with tf.Session(config=tf.ConfigProto(allow_soft_placement=True), graph=g) as sess:

tests/test_backend.py

@@ -353,6 +353,23 @@ def test_dropout(self):
         output_names_with_port = ["output:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port)
 
+    def test_nn_dropout(self):
+        keep_prob = tf.placeholder_with_default(1., (), "keep_prob")
+        x_val = np.ones([1, 24, 24, 3], dtype=np.float32)
+        # Define a scope for reusing the variables
+        x = tf.placeholder(tf.float32, shape=x_val.shape, name="input_1")
+        x_ = tf.identity(x)
+
+        fc1 = tf.nn.dropout(x_, keep_prob)
+
+        _ = tf.identity(fc1, name="output")
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        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)
+
     def test_conv2d_with_input_transpose(self):
         x_shape = [2, 32, 32, 3]
         kernel_shape = [3, 3, 3, 3]

tests/test_graph.py

@@ -165,10 +165,9 @@ def test_dropout(self):
             g = process_tf_graph(sess.graph, opset=self.config.opset)
             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] Dropout__3 [op_type=Dropout] ' \
-                       'output1 [op_type=Identity] output2 [op_type=Identity] output [op_type=Identity] ' \
-                       'input1:0 -> Add input2:0 -> Add Add:0 -> Dropout__3 Dropout__3:0 -> output1 ' \
-                       'output1:0 -> output2 output2:0 -> output }'
+                       '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)
 
     def test_add(self):

tf2onnx/graph.py

@@ -687,6 +687,7 @@ def _get_unvisited_child(g, node, not_visited):
             all_input = list(filter(lambda a: a != '', all_input))
             for inp in all_input:
                 j = self.get_node_by_output(inp)
+                utils.make_sure(j is not None, "Cannot find node with output {}".format(inp))
                 if self.parent_graph and j.name not in op_name_to_index:
                     # there might be some outer-scoped inputs for an inner Graph.
                     pass
@@ -754,6 +755,7 @@ def make_graph(self, doc, graph_name="tf2onnx"):
         placeholder_default_const_ops = []
         for op in placeholder_ops:
             if op.type == "PlaceholderWithDefault":
+                utils.make_sure(op.inputs[0] is not None, "Cannot find node with output {}".format(op.input[0]))
                 utils.make_sure(op.inputs[0].is_const(),
                                 "non-const default value for PlaceholderWithDefault is not supported.")
                 # copy the tensor value, set its name to current node's output, add as initializer
@@ -1027,6 +1029,11 @@ def extract_sub_graph_nodes(self, outputs_name, input_checker=None, ignore_unuse
                 if node.is_graph_input():
                     if node not in res_set:
                         res_set.add(node)
+                    if node.type == "PlaceholderWithDefault" and \
+                            node.inputs[0].is_const() and \
+                            node.inputs[0] not in res_set:
+                        res_set.add(node.inputs[0])
+
         return list(res_set)
 
     def delete_unused_nodes(self, outputs_name):

tf2onnx/tfonnx.py

@@ -192,7 +192,7 @@ def rewrite_dropout(g, ops):
             OpTypePattern('Floor', inputs=[
                 OpTypePattern('Add', inputs=[
                     OpTypePattern(None, name="input3"),
-                    OpTypePattern('RandomUniform'),
+                    OpTypePattern('RandomUniform|RandomUniformLike'),
                 ])
             ]),
         ])
@@ -216,6 +216,12 @@ def rewrite_dropout(g, ops):
         for n in set(match.get_nodes()):
             g.remove_node(n.name)
 
+    # 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