rewrite the graph and body graph recursively

Author: wayuanho

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

@@ -391,6 +391,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

@@ -2419,6 +2419,28 @@ def tensorflow_onnx_mapping(g, continue_on_error, custom_op_handlers):
     return mapped_op, unmapped_op
 
 
+def tensorflow_onnx_rewrite(g, rewriters):
+    try:
+        ops = g.get_nodes()
+        for rewrite in rewriters:
+            ops = rewrite(g, ops)
+            g.set_nodes(ops)
+        for node in ops:
+            body_graphs = node.get_body_graphs()
+            if body_graphs:
+                for attr, b_g in body_graphs.items():
+                    log.debug("start rewriting subgraph of %s's attribute %s", node.name, attr)
+                    tensorflow_onnx_rewrite(b_g, rewriters)
+    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
+
+
 def transpose_inputs(ctx, inputs_as_nchw):
     """Insert a transpose from NHWC to NCHW on model input on users request."""
     ops = []
@@ -2553,19 +2575,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
+    tensorflow_onnx_rewrite(g, rewriters)
 
     # some nodes may already copied into inner Graph, so remove them from main Graph.
     g.delete_unused_nodes(output_names)