Merge pull request #337 from lucienwang1009/cond_inside_loop

rewrite the graph and body graph recursively

Author: nbcsm

tests/test_cond.py

@@ -35,7 +35,6 @@ def test_simple_cond(self):
         output_names_with_port = ["output:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port)
 
-    @unittest.skip("known issue about onnxruntime that initilizer is subgraph input")
     def test_cond_with_const_branch(self):
         x_val = np.array([1, 2, 3], dtype=np.float32)
         y_val = np.array([4, 5, 6], dtype=np.float32)
@@ -132,8 +131,7 @@ def cond_graph2():
         output_names_with_port = ["output:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port)
 
-    @unittest.skip("not support for now")
-    def test_cond_with_while_loop(self):
+    def test_while_loop_between_conds(self):
         x_val = np.array([1, 2, 3], dtype=np.float32)
         y_val = np.array([4, 5, 6], dtype=np.float32)
         x = tf.placeholder(tf.float32, x_val.shape, name="input_1")
@@ -156,7 +154,6 @@ def cond_graph():
         output_names_with_port = ["output:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port)
 
-    @unittest.skip("not support for now")
     def test_cond_in_while_loop(self):
         i = tf.placeholder(tf.int32, (), name="input_1")
         inputs = tf.placeholder(tf.float32, (10,), name="input_2")
@@ -170,7 +167,7 @@ def test_cond_in_while_loop(self):
         def b(i, out_ta):
             new_i = tf.add(i, 1)
             x = input_ta.read(i)
-            x = tf.cond(x >= 0, lambda: x - 1, lambda: x + 3)
+            x = tf.cond(x > 0, lambda: x - 1, lambda: x + 3)
             out_ta_new = out_ta.write(i, x)
             return new_i, out_ta_new
 

tf2onnx/graph.py

@@ -390,6 +390,8 @@ def make_node(self, op_type, inputs, attr=None, output_count=1, outputs=None, sk
         if outputs is None:
             outputs = [name + ":" + str(i) for i in range(output_count)]
 
+        output_count = len(outputs)
+
         raw_attr = {}
         onnx_attrs = []
         for a, v in attr.items():

tf2onnx/rewriter/cond_rewriter.py

@@ -104,28 +104,63 @@ def run(self):
 
         return self.g.get_nodes()
 
+    def _get_output_shape_dtype(self, cond_context):
+        output_shapes = []
+        output_dtypes = []
+        for i, _ in enumerate(cond_context.true_branch_context.output):
+            true_output = cond_context.true_branch_context.output[i]
+            false_output = cond_context.false_branch_context.output[i]
+            true_shape = self.g.get_shape(true_output)
+            true_dtype = self.g.get_dtype(true_output)
+            false_shape = self.g.get_shape(false_output)
+            false_dtype = self.g.get_dtype(false_output)
+            if true_shape != false_shape:
+                raise RuntimeError(
+                    "the shape of outputs {} and {} mismatch: {}, {}".format(
+                        true_output,
+                        false_output,
+                        true_shape,
+                        false_shape
+                    )
+                )
+            if true_dtype != false_dtype:
+                raise RuntimeError(
+                    "the shape of outputs {} and {} mismatch: {}, {}".format(
+                        true_output,
+                        false_output,
+                        true_dtype,
+                        false_dtype
+                    )
+                )
+            output_shapes.append(true_shape)
+            output_dtypes.append(true_dtype)
+        return output_shapes, output_dtypes
+
     def _create_if_node(self, cond_context):
+        output_shapes, output_dtypes = self._get_output_shape_dtype(cond_context)
         if_node = self.g.make_node(
             "If",
             [cond_context.pred_input],
             op_name_scope=cond_context.cond_scope,
             outputs=[m.output[0] for m in cond_context.merges],
+            shapes=output_shapes,
+            dtypes=output_dtypes,
             skip_conversion=False
         )
         log.debug("set graph for if branchs")
         true_graph = utils.construct_graph_from_nodes(
             self.g,
             list(cond_context.true_branch_context.nodes),
             cond_context.true_branch_context.output,
-            [self.g.get_shape(out) for out in cond_context.true_branch_context.output],
-            [self.g.get_dtype(out) for out in cond_context.true_branch_context.output]
+            output_shapes,
+            output_dtypes
         )
         false_graph = utils.construct_graph_from_nodes(
             self.g,
             list(cond_context.false_branch_context.nodes),
             cond_context.false_branch_context.output,
-            [self.g.get_shape(out) for out in cond_context.false_branch_context.output],
-            [self.g.get_dtype(out) for out in cond_context.false_branch_context.output]
+            output_shapes,
+            output_dtypes
         )
         if_node.set_body_graph_as_attr("then_branch", true_graph)
         if_node.set_body_graph_as_attr("else_branch", false_graph)
@@ -268,7 +303,7 @@ def _branch_type(self, branch_output, nodes):
         if branch == BranchType.UNKNOWN:
             log.debug(
                 "branch only contains const node: [%s]",
-                ",".join(n for n in nodes)
+                ",".join(n.name for n in nodes)
             )
         return branch
 

tf2onnx/tfonnx.py

@@ -2482,16 +2482,28 @@ def topological_sort(g, continue_on_error):
             pass
 
 
-def run_late_rewriters(g, funcs, continue_on_error):
+def run_rewriters(g, funcs, continue_on_error):
+    """Rewrite the original graph and body graphs of nodes"""
+    # NOTE(wayuanho):
+    # 1. we don't sort graph here, rewriter is expected to do it on its own.
+    # 2. the graph here may have circles, current topological_sort cannot handle it.
+    for func in funcs:
+        try:
+            ops = func(g, g.get_nodes())
+            g.set_nodes(ops)
+        except Exception as ex:
+            type_, value_, traceback_ = sys.exc_info()
+            log.error("rewriter %s: exception %s", func, ex)
+            ex_ext = traceback.format_exception(type_, value_, traceback_)
+            if continue_on_error:
+                log.info(ex_ext)
+            else:
+                raise ex
+
     if g.contained_graphs:
         for dict_val in g.contained_graphs.values():
             for attr_name, b_g in dict_val.items():
-                run_late_rewriters(b_g, funcs, attr_name)
-
-    topological_sort(g, continue_on_error)
-    for func in funcs:
-        ops = func(g, g.get_nodes())
-        g.set_nodes(ops)
+                run_rewriters(b_g, funcs, attr_name)
 
 
 def process_tf_graph(tf_graph, continue_on_error=False, verbose=False, target=None,
@@ -2561,19 +2573,7 @@ def process_tf_graph(tf_graph, continue_on_error=False, verbose=False, target=No
     if custom_rewriter is not None:
         rewriters.extend(custom_rewriter)
 
-    try:
-        ops = g.get_nodes()
-        for rewrite in rewriters:
-            ops = rewrite(g, ops)
-            g.set_nodes(ops)
-    except Exception as ex:
-        type_, value_, traceback_ = sys.exc_info()
-        log.error("node %s: exception %s" % (rewrite, ex))
-        ex_ext = traceback.format_exception(type_, value_, traceback_)
-        if continue_on_error:
-            log.info(ex_ext)
-        else:
-            raise ex
+    run_rewriters(g, rewriters, continue_on_error)
 
     # some nodes may already copied into inner Graph, so remove them from main Graph.
     g.delete_unused_nodes(output_names)
@@ -2590,7 +2590,7 @@ def process_tf_graph(tf_graph, continue_on_error=False, verbose=False, target=No
     if TARGET_RS6 in target:
         late_rewriters.append(rewrite_incomplete_type_support_rs6)
     if late_rewriters:
-        run_late_rewriters(g, late_rewriters, continue_on_error)
+        run_rewriters(g, late_rewriters, continue_on_error)
 
     # onnx requires topological sorting
     topological_sort(g, continue_on_error)