Merge pull request #581 from lucienwang1009/init_cond_loop

calculate condition initialization for Loop

Author: nbcsm

tests/test_loops.py

@@ -160,6 +160,33 @@ def b(i, out_ta):
         output_names_with_port = ["i:0", "output_ta:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
 
+    def test_while_loop_with_cond_init_false(self):
+        i = tf.placeholder(tf.int32, (), name="input_1")
+        inputs = tf.placeholder(tf.float32, (10,), name="input_2")
+
+        inputs_2 = tf.identity(inputs)
+        input_ta = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True).unstack(inputs_2)
+        output_ta = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True)
+
+        c = lambda i, *_: tf.logical_and(i < 10, i >= 0)
+
+        def b(i, out_ta):
+            new_i = tf.add(i, 1)
+            x = input_ta.read(i)
+            y = x + 3
+            out_ta_new = out_ta.write(i, y)
+            return new_i, out_ta_new
+
+        i_final, out_final = tf.while_loop(c, b, [i, output_ta])
+        _ = tf.identity(i_final, name="i")
+        _ = tf.identity(out_final.stack(), name="output_ta")
+        input_names_with_port = ["input_1:0", "input_2:0"]
+        feed_dict = {"input_1:0": np.array(20, dtype=np.int32),
+                     "input_2:0": np.array([2.0, 16.0, 5.0, 1.6, 5.0, 6.0, 7.0, 8.0, 9.0, 10.], dtype=np.float32)}
+
+        output_names_with_port = ["i:0", "output_ta:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
+
     def test_map_fn(self):
         def fn0(elem):
             res = elem + elem * elem

tf2onnx/rewriter/loop_rewriter.py

@@ -45,21 +45,19 @@ def rewrite(self, context):
             cell_g_info = context.cell_graph
             cond_g_info = context.cond_graph
 
-            # todo(pengwa): we don't check the case where loop body won't be executed at all.
+            # create a dummy loop to calculate the init condition
+            init_cond_output = self._create_subgraph_initial_cond(cond_g_info)
 
             ## create Loop body graph with existing nodes
 
-            # replace condition graph's inputs to be cell graph's outputs, because we want condition graph
-            # to consumer cell graph outputs.
-            for loop_var in cond_g_info.dependent_vars:
-                self.g.replace_all_inputs(cond_g_info.nodes, loop_var.switch_true_identity_output.id,
-                                          loop_var.next_iteration_input.id)
-
             body_nodes = set(cell_g_info.nodes + cond_g_info.nodes)
             body_outputs = cond_g_info.outputs + cell_g_info.outputs
             for out_tensor_value_info in body_outputs:
                 shape = out_tensor_value_info.shape
-                utils.make_sure(shape is not None, "Shape of {} is None".format(out_tensor_value_info.id))
+                utils.make_sure(
+                    shape is not None,
+                    "Conversion of Loop requries output shape [{}] exists".format(out_tensor_value_info.id)
+                )
                 out_tensor_value_info.shape = utils.create_vague_shape_like(shape)
 
             loop_body_g = LoopRewriterBase.construct_graph_from_nodes(self.g, body_nodes, body_outputs)
@@ -90,7 +88,7 @@ def rewrite(self, context):
                 loop_body_g.replace_all_inputs(loop_body_g.get_nodes(), input_ta.consumer.id, data_node.output[0])
 
             ## create Loop node
-            loop_node = self._create_loop_node(context, loop_props)
+            loop_node = self._create_loop_node(context, loop_props, init_cond_output)
             if not loop_node:
                 logger.error("failed to create loop node during rewrite")
                 return REWRITER_RESULT.FAIL
@@ -104,7 +102,48 @@ def rewrite(self, context):
             logger.error("loop rewrite failed, due to exception: %s, details:%s", ex, tb)
             return REWRITER_RESULT.FAIL
 
-    def _create_loop_node(self, context, loop_props):
+    def _create_subgraph_initial_cond(self, cond_graph):
+        """Create subgraph to calculate initial cond."""
+        # copy condition subgraph to parent graph
+        copied_nodes = []
+        name_scope = utils.make_name("copy")
+        for node in cond_graph.nodes:
+            new_name = "{}/{}".format(name_scope, node.name)
+            new_outputs = ["{}/{}".format(name_scope, out) for out in node.output]
+            # some inputs are out of cond_graph.nodes, keep them intact
+            new_inputs = []
+            for inp in node.input:
+                if self.g.get_node_by_output(inp) in cond_graph.nodes:
+                    new_inputs.append("{}/{}".format(name_scope, inp))
+                else:
+                    new_inputs.append(inp)
+
+            new_node = self.g.make_node(
+                node.type, new_inputs, outputs=new_outputs,
+                attr=node.attr, name=new_name,
+                shapes=node.output_shapes, dtypes=node.output_dtypes,
+                skip_conversion=node.skip_conversion, infer_shape_dtype=False
+            )
+            body_graphs = node.graph.contained_graphs.pop(node.name, None)
+            if body_graphs:
+                for attr_name, body_graph in body_graphs.items():
+                    body_graph.parent_graph = g
+                    new_node.set_body_graph_as_attr(attr_name, body_graph)
+            copied_nodes.append(new_node)
+
+        # replace all inputs of condition graph by initializer (enter_input)
+        for loop_var in cond_graph.dependent_vars:
+            self.g.replace_all_inputs(
+                copied_nodes,
+                loop_var.next_iteration_input.id,
+                loop_var.enter_input_id
+            )
+        init_cond_output = "{}/{}".format(name_scope, cond_graph.outputs[0].id)
+        self.g.set_dtype(init_cond_output, cond_graph.outputs[0].dtype)
+        self.g.set_shape(init_cond_output, cond_graph.outputs[0].shape)
+        return init_cond_output
+
+    def _create_loop_node(self, context, loop_props, init_cond_output):
         loop_outputs = []
         loop_output_shapes = []
         loop_output_dtypes = []
@@ -123,8 +162,7 @@ def _create_loop_node(self, context, loop_props):
         # trip count and cond are not used, giving them values just because bug
         # (https://github.com/Microsoft/onnxruntime/issues/255) of onnxruntime.
         trip_cnt = self.g.make_const(utils.make_name("trip_count"), np.array(sys.maxsize, dtype=np.int64))
-        cond = self.g.make_const(utils.make_name("cond"), np.array(True, dtype=np.bool))
-        loop_node = self.g.make_node("Loop", [trip_cnt.output[0]] + [cond.output[0]] +
+        loop_node = self.g.make_node("Loop", [trip_cnt.output[0]] + [init_cond_output] +
                                      loop_props.state_inputs_initial_values,  # ONNX Loop support state inputs only
                                      outputs=loop_outputs, op_name_scope="generic_loop",
                                      shapes=loop_output_shapes, dtypes=loop_output_dtypes,

tf2onnx/rewriter/loop_rewriter_base.py

@@ -295,9 +295,11 @@ def _crop_loop_condition_sub_graph(self, context):
             loop_var = context.loop_properties.all_variables[enter_node.name]
 
             # cut off connection between condition graph and Merge node.
+            # replace condition graph's inputs to be cell graph's outputs, because we want condition graph
+            # to consumer cell graph outputs.
             non_switch_consumers = [n for n in self.g.find_output_consumers(merge_node.output[0]) if n.type != "Switch"]
             self.g.replace_all_inputs(non_switch_consumers, merge_node.output[0],
-                                      loop_var.switch_true_identity_output.id)
+                                      loop_var.next_iteration_input.id)
             dependent_vars.append(loop_var)
 
         # cut off connection between condition graph and LoopCond node.