support multi birnn share the same input

lucienwang1009 · lucienwang1009 · commit 03cf6cd93389 · 2019-05-29T07:25:27.000Z
diff --git a/tests/test_gru.py b/tests/test_gru.py
@@ -482,6 +482,56 @@ def test_dynamic_bidirectional_but_one_gru_and_state_consumed_only(self):
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06,
                            graph_validator=lambda g: check_gru_count(g, 1))
 
+    def test_dynamic_multi_bigru_with_same_input(self):
+        units = 5
+        batch_size = 1
+        x_val = np.array([[1., 1.], [2., 2.], [3., 3.]], dtype=np.float32)
+        x_val = np.stack([x_val] * batch_size)
+
+        x = tf.placeholder(tf.float32, x_val.shape, name="input_1")
+        initializer = init_ops.constant_initializer(0.5)
+
+        # bigru, no scope
+        cell1 = rnn.GRUCell(
+            units,
+            kernel_initializer=initializer)
+        cell2 = rnn.GRUCell(
+            units,
+            kernel_initializer=initializer)
+        outputs_1, cell_state_1 = tf.nn.bidirectional_dynamic_rnn(
+            cell1,
+            cell2,
+            x,
+            dtype=tf.float32,
+            scope="bigru_1"
+        )
+
+        units = 10
+        cell1 = rnn.GRUCell(
+            units,
+            kernel_initializer=initializer)
+        cell2 = rnn.GRUCell(
+            units,
+            kernel_initializer=initializer)
+        outputs_2, cell_state_2 = tf.nn.bidirectional_dynamic_rnn(
+            cell1,
+            cell2,
+            x,
+            dtype=tf.float32,
+            scope="bigru_2"
+        )
+
+        _ = tf.identity(outputs_1, name="output_1")
+        _ = tf.identity(cell_state_1, name="cell_state_1")
+        _ = tf.identity(outputs_2, name="output_2")
+        _ = tf.identity(cell_state_2, name="cell_state_2")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        output_names_with_port = ["output_1:0", "cell_state_1:0", "output_2:0", "cell_state_2:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06,
+                           graph_validator=lambda g: check_gru_count(g, 2))
+
 
 if __name__ == '__main__':
     unittest_main()
diff --git a/tests/test_grublock.py b/tests/test_grublock.py
@@ -447,6 +447,47 @@ def test_dynamic_bidirectional_but_one_gru_and_state_consumed_only(self):
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06,
                            graph_validator=lambda g: check_gru_count(g, 1))
 
+    def test_dynamic_multi_bigru_with_same_input(self):
+        units = 5
+        batch_size = 1
+        x_val = np.array([[1., 1.], [2., 2.], [3., 3.]], dtype=np.float32)
+        x_val = np.stack([x_val] * batch_size)
+
+        x = tf.placeholder(tf.float32, x_val.shape, name="input_1")
+
+        # bigru, no scope
+        cell1 = rnn.GRUBlockCell(units)
+        cell2 = rnn.GRUBlockCell(units)
+        outputs_1, cell_state_1 = tf.nn.bidirectional_dynamic_rnn(
+            cell1,
+            cell2,
+            x,
+            dtype=tf.float32,
+            scope="bigru_1"
+        )
+
+        units = 10
+        cell1 = rnn.GRUBlockCell(units)
+        cell2 = rnn.GRUBlockCell(units)
+        outputs_2, cell_state_2 = tf.nn.bidirectional_dynamic_rnn(
+            cell1,
+            cell2,
+            x,
+            dtype=tf.float32,
+            scope="bigru_2"
+        )
+
+        _ = tf.identity(outputs_1, name="output_1")
+        _ = tf.identity(cell_state_1, name="cell_state_1")
+        _ = tf.identity(outputs_2, name="output_2")
+        _ = tf.identity(cell_state_2, name="cell_state_2")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        output_names_with_port = ["output_1:0", "cell_state_1:0", "output_2:0", "cell_state_2:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06,
+                           graph_validator=lambda g: check_gru_count(g, 2))
+
 
 if __name__ == '__main__':
     unittest_main()
diff --git a/tests/test_lstm.py b/tests/test_lstm.py
@@ -554,6 +554,69 @@ def test_dynamic_bilstm_state_consumed_only(self, state_is_tuple=True):
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06,
                            graph_validator=lambda g: check_lstm_count(g, 1))
 
+    def test_dynamic_multi_bilstm_with_same_input_state_is_tuple(self):
+        self.internal_test_dynamic_multi_bilstm_with_same_input(True)
+
+    def test_dynamic_multi_bilstm_with_same_input_state_is_list(self):
+        self.internal_test_dynamic_multi_bilstm_with_same_input(False)
+
+    def internal_test_dynamic_multi_bilstm_with_same_input(self, state_is_tuple):
+        units = 5
+        batch_size = 1
+        x_val = np.array([[1., 1.], [2., 2.], [3., 3.]], dtype=np.float32)
+        x_val = np.stack([x_val] * batch_size)
+
+        x = tf.placeholder(tf.float32, x_val.shape, name="input_1")
+        initializer = init_ops.constant_initializer(0.5)
+
+        cell1 = rnn.LSTMCell(
+            units,
+            initializer=initializer,
+            state_is_tuple=state_is_tuple
+        )
+        cell2 = rnn.LSTMCell(
+            units,
+            initializer=initializer,
+            state_is_tuple=state_is_tuple
+        )
+        outputs_1, cell_state_1 = tf.nn.bidirectional_dynamic_rnn(
+            cell1,
+            cell2,
+            x,
+            dtype=tf.float32,
+            scope="bilstm_1"
+        )
+
+        units = 10
+        cell1 = rnn.LSTMCell(
+            units,
+            initializer=initializer,
+            state_is_tuple=state_is_tuple
+        )
+        cell2 = rnn.LSTMCell(
+            units,
+            initializer=initializer,
+            state_is_tuple=state_is_tuple
+        )
+        outputs_2, cell_state_2 = tf.nn.bidirectional_dynamic_rnn(
+            cell1,
+            cell2,
+            x,
+            dtype=tf.float32,
+            scope="bilstm_2"
+        )
+
+        _ = tf.identity(outputs_1, name="output_1")
+        _ = tf.identity(cell_state_1, name="cell_state_1")
+        _ = tf.identity(outputs_2, name="output_2")
+        _ = tf.identity(cell_state_2, name="cell_state_2")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        output_names_with_port = ["output_1:0", "cell_state_1:0", "output_2:0", "cell_state_2:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06,
+                           graph_validator=lambda g: check_lstm_count(g, 2))
+
 
 if __name__ == '__main__':
     unittest_main()
diff --git a/tests/test_lstmblock.py b/tests/test_lstmblock.py
@@ -488,5 +488,45 @@ def test_multi_rnn_lstm(self):
         output_names_with_port = ["output:0", "cell_state:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-03, atol=1e-06)
 
+    def test_dynamic_multi_bilstm_with_same_input(self):
+        units = 5
+        batch_size = 1
+        x_val = np.array([[1., 1.], [2., 2.], [3., 3.]], dtype=np.float32)
+        x_val = np.stack([x_val] * batch_size)
+
+        x = tf.placeholder(tf.float32, x_val.shape, name="input_1")
+
+        cell1 = rnn.LSTMBlockCell(units)
+        cell2 = rnn.LSTMBlockCell(units)
+        outputs_1, cell_state_1 = tf.nn.bidirectional_dynamic_rnn(
+            cell1,
+            cell2,
+            x,
+            dtype=tf.float32,
+            scope="bilstm_1"
+        )
+
+        units = 10
+        cell1 = rnn.LSTMBlockCell(units)
+        cell2 = rnn.LSTMBlockCell(units)
+        outputs_2, cell_state_2 = tf.nn.bidirectional_dynamic_rnn(
+            cell1,
+            cell2,
+            x,
+            dtype=tf.float32,
+            scope="bilstm_2"
+        )
+
+        _ = tf.identity(outputs_1, name="output_1")
+        _ = tf.identity(cell_state_1, name="cell_state_1")
+        _ = tf.identity(outputs_2, name="output_2")
+        _ = tf.identity(cell_state_2, name="cell_state_2")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        output_names_with_port = ["output_1:0", "cell_state_1:0", "output_2:0", "cell_state_2:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3, atol=1e-06,
+                           graph_validator=lambda g: check_lstm_count(g, 2))
+
 if __name__ == '__main__':
     unittest_main()
diff --git a/tf2onnx/graph.py b/tf2onnx/graph.py
@@ -1036,6 +1036,15 @@ def find_output_consumers(self, output_name):
                     nodes.extend(g.find_output_consumers(output_name))
         return nodes
 
+    def find_common_consumers(self, *outputs):
+        if not outputs:
+            return False
+
+        common_consumer = set.intersection(
+            *[set(self.find_output_consumers(out)) for out in outputs]
+        )
+        return list(common_consumer)
+
     @staticmethod
     def replace_all_inputs(ops, old_input, new_input):
         """Replace all inputs pointing to old_input with new_input."""
diff --git a/tf2onnx/rewriter/bigru_rewriter.py b/tf2onnx/rewriter/bigru_rewriter.py
@@ -13,24 +13,18 @@
 import logging
 import numpy as np
 from tf2onnx import utils
-from tf2onnx.utils import is_tf_reverse_op
-from tf2onnx.rewriter.bilstm_rewriter import slice_bilstm_for_original_lstm_consumers,\
-     get_reverse_nodes_after_y_output, get_np_val_for_const, _process_single_init_node
-
+from tf2onnx.rewriter.rnn_utils import find_bidirectional_rnns, ONNX_RNN_TYPE, get_np_val_for_const, \
+    process_single_init_node, slice_birnn_for_original_rnn_consumers
 
 
 logger = logging.getLogger(__name__)
 
 # pylint: disable=invalid-name,unused-argument,missing-docstring
 
 def process_bigru(g, bi_grus):
-    for fw, bw in bi_grus:
-        input_id = fw[0]
+    for gru_fw, gru_bw in bi_grus:
         logger.debug("=========================")
-        logger.debug("start handling potential bidirectional gru %s", input_id)
-
-        gru_fw = fw[1]
-        gru_bw = bw[1]
+        logger.debug("start handling potential bidirectional gru: %s, %s", gru_fw.name, gru_bw.name)
 
         w_fw = get_np_val_for_const(g, gru_fw, 1)
         w_bw = get_np_val_for_const(g, gru_bw, 1)
@@ -102,80 +96,25 @@ def process_bigru(g, bi_grus):
 
         to_remove = [gru_fw.name, gru_fw.input[1], gru_fw.input[2], gru_fw.input[3],
                      gru_bw.name, gru_bw.input[1], gru_bw.input[2], gru_bw.input[3]]
-        slice_bilstm_for_original_lstm_consumers(
+        slice_birnn_for_original_rnn_consumers(
             g, gru_fw, gru_bw, bi_gru_node, 0, all_nodes, to_remove)
-        slice_bilstm_for_original_lstm_consumers(
+        slice_birnn_for_original_rnn_consumers(
             g, gru_fw, gru_bw, bi_gru_node, 1, all_nodes, to_remove)
 
-        gru_bw_old_x = gru_bw.input[0]
-
         for n in to_remove:
             g.remove_node(n)
 
-        old_x_consumers = g.find_output_consumers(gru_bw_old_x)
-        # the transpose/reverse here must be followed by GRU if it is still useful.
-        # this is guaranteed by dynamic_rnn logic.
-        old_x_has_gru_as_consumer = [
-            n for n in old_x_consumers if n.type == "GRU"]
-        if not old_x_has_gru_as_consumer:
-            logger.debug("plan to remove useless reverse op in bw")
-            reverse_node = g.get_node_by_output(gru_bw_old_x)
-
-            if reverse_node.type == "Transpose":
-                reverse_node = reverse_node.inputs[0]
-
-            g.replace_all_inputs(
-                g.get_nodes(), reverse_node.output[0], reverse_node.input[0])
-            g.remove_node(reverse_node.name)
-        else:
-            raise ValueError(
-                "Reverse is still used by GRU as input, cannot remove")
-
     return g.get_nodes()
 
 
 def process_init_nodes(g, gru_fw, gru_bw, to_append):
-    initializer_node = _process_single_init_node(
+    initializer_node = process_single_init_node(
         g, gru_fw.input[5], gru_bw.input[5], to_append)
 
     return initializer_node
 
 
 def rewrite_bidirectional_grus(g, ops):
-    """
-        return: list of tuple, format of tuple is
-        ((fw input_id, fw onnx gru node), (bw input_id, bw onnx gru node)), and fw input_id equals to bw input_id
-    """
-    fw_gru = {}
-    bw_gru = {}
-    for n in g.get_nodes():
-        if n.type != "GRU":
-            continue
-        input_id = n.input[0]
-        temp = n.inputs[0]
-        is_backward_gru = False
-        if temp.type == "Transpose":
-            input_id = temp.input[0]
-            temp = temp.inputs[0]
-
-        if is_tf_reverse_op(temp):
-            input_id = temp.input[0]
-            is_backward_gru = True
-
-        if is_backward_gru:
-            # if output 0 is consumed, and there is no reverse after the gru output.
-            # it's not reversed gru
-            if g.find_output_consumers(n.output[0]) and not get_reverse_nodes_after_y_output(g, n):
-                continue
-            logger.debug("find bw gru %s", input_id)
-            bw_gru[input_id] = [input_id, n]
-        else:
-            logger.debug("find fw gru %s", input_id)
-            fw_gru[input_id] = [input_id, n]
-
-    # when fw_gru has same input as bw_gru, then it may be a bi gru
-    bigru_input = list(set(fw_gru.keys()).intersection(bw_gru.keys()))
-    bi_grus = [(fw_gru[input_id], bw_gru[input_id])
-               for input_id in bigru_input]
+    bi_grus = find_bidirectional_rnns(g, ops, ONNX_RNN_TYPE.GRU)
 
     return process_bigru(g, bi_grus)
diff --git a/tf2onnx/rewriter/bilstm_rewriter.py b/tf2onnx/rewriter/bilstm_rewriter.py
diff --git a/tf2onnx/rewriter/rnn_utils.py b/tf2onnx/rewriter/rnn_utils.py
