fix bug of bigru/bilstm

only "state" tensor of bigru is consmued, then output of bigru doesn't have
a reverse op.

can be merged

Author: zhijxu-MS

tests/run_pretrained_models.py

@@ -23,14 +23,14 @@
 import tensorflow as tf
 from tensorflow.core.framework import graph_pb2
 from tensorflow.python.framework.graph_util import convert_variables_to_constants
+from tensorflow.contrib.rnn import GRUBlockCell  # pylint: disable=unused-import
 import yaml
 import PIL.Image
 
 import tf2onnx
 from tf2onnx import utils
 from tf2onnx.graph import GraphUtil
 from tf2onnx.tfonnx import process_tf_graph
-from tensorflow.contrib.rnn import GRUBlockCell  # pylint: disable=unused-import
 
 # pylint: disable=broad-except,logging-not-lazy,unused-argument,unnecessary-lambda
 

tests/test_gru.py

@@ -255,6 +255,54 @@ def test_single_dynamic_gru_random_weights2(self):
         output_names_with_port = ["output:0", "cell_state:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.01)
 
+    def test_dynamic_gru_output_consumed_only(self):
+        units = 5
+        batch_size = 6
+        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 = tf.random_uniform_initializer(-1.0, 1.0)
+        cell1 = rnn.GRUCell(
+            units,
+            kernel_initializer=initializer)
+
+        outputs, _ = tf.nn.dynamic_rnn(
+            cell1,
+            x,
+            dtype=tf.float32)
+
+        _ = tf.identity(outputs, name="output")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        output_names_with_port = ["output:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.0001)
+
+    def test_dynamic_gru_state_consumed_only(self):
+        units = 5
+        batch_size = 6
+        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 = tf.random_uniform_initializer(-1.0, 1.0)
+        cell1 = rnn.GRUCell(
+            units,
+            kernel_initializer=initializer)
+
+        _, cell_state = tf.nn.dynamic_rnn(
+            cell1,
+            x,
+            dtype=tf.float32)
+
+        _ = tf.identity(cell_state, name="cell_state")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        output_names_with_port = ["cell_state:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.0001)
+
     def test_dynamic_bigru(self):
         units = 5
         batch_size = 1
@@ -320,6 +368,38 @@ def test_dynamic_bigru_output_consumed_only(self):
         output_names_with_port = ["output:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3)
 
+    def test_dynamic_bigru_state_consumed_only(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)
+
+        gru_list = []
+        if True:
+            # bigru, no scope
+            cell1 = rnn.GRUCell(
+                units,
+                kernel_initializer=initializer)
+            cell2 = rnn.GRUCell(
+                units,
+                kernel_initializer=initializer)
+            outputs, cell_state = tf.nn.bidirectional_dynamic_rnn(
+                cell1,
+                cell2,
+                x,
+                dtype=tf.float32)
+            gru_list.append(outputs)
+
+            _ = tf.identity(cell_state, name="cell_state")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        output_names_with_port = ["cell_state:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3)
+
     def test_dynamic_bidirectional_but_one_gru(self):
         units = 5
         batch_size = 1
@@ -377,6 +457,33 @@ def test_dynamic_bidirectional_but_one_gru_and_output_consumed_only(self):
         output_names_with_port = ["output:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3)
 
+    def test_dynamic_bidirectional_but_one_gru_and_state_consumed_only(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")
+
+        gru_list = []
+        if True:
+            # bigru, no scope
+            cell = rnn.GRUCell(
+                units)
+            outputs, cell_state = tf.nn.bidirectional_dynamic_rnn(
+                cell,
+                cell,
+                x,
+                dtype=tf.float32)
+            gru_list.append(outputs)
+
+        _ = tf.identity(cell_state, name="cell_state")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        output_names_with_port = ["cell_state:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3)
+
 
 if __name__ == '__main__':
     Tf2OnnxBackendTestBase.trigger(GRUTests)

tests/test_grublock.py

@@ -238,6 +238,50 @@ def test_single_dynamic_gru_random_weights2(self):
         output_names_with_port = ["output:0", "cell_state:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.01)
 
+    def test_dynamic_gru_output_consumed_only(self):
+        units = 5
+        batch_size = 6
+        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.GRUBlockCell(
+            units)
+
+        outputs, _ = tf.nn.dynamic_rnn(
+            cell1,
+            x,
+            dtype=tf.float32)
+
+        _ = tf.identity(outputs, name="output")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        output_names_with_port = ["output:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.0001)
+
+    def test_dynamic_gru_state_consumed_only(self):
+        units = 5
+        batch_size = 6
+        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.GRUBlockCell(
+            units)
+
+        _, cell_state = tf.nn.dynamic_rnn(
+            cell1,
+            x,
+            dtype=tf.float32)
+
+        _ = tf.identity(cell_state, name="cell_state")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        output_names_with_port = ["cell_state:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.0001)
+
     def test_dynamic_bigru(self):
         units = 5
         batch_size = 1
@@ -297,6 +341,35 @@ def test_dynamic_bigru_output_consumed_only(self):
         output_names_with_port = ["output:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3)
 
+    def test_dynamic_bigru_state_consumed_only(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")
+
+        gru_list = []
+        if True:
+            # bigru, no scope
+            cell1 = rnn.GRUBlockCell(
+                units)
+            cell2 = rnn.GRUBlockCell(
+                units)
+            _, cell_state = tf.nn.bidirectional_dynamic_rnn(
+                cell1,
+                cell2,
+                x,
+                dtype=tf.float32)
+            gru_list.append(cell_state)
+
+        _ = tf.identity(cell_state, name="cell_state")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        output_names_with_port = ["cell_state:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3)
+
     def test_dynamic_bidirectional_but_one_gru(self):
         units = 5
         batch_size = 1
@@ -352,6 +425,33 @@ def test_dynamic_bidirectional_but_one_gru_and_output_consumed_only(self):
         output_names_with_port = ["output:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3)
 
+    def test_dynamic_bidirectional_but_one_gru_and_state_consumed_only(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")
+
+        gru_list = []
+        if True:
+            # bigru, no scope
+            cell = rnn.GRUBlockCell(
+                units)
+            _, cell_state = tf.nn.bidirectional_dynamic_rnn(
+                cell,
+                cell,
+                x,
+                dtype=tf.float32)
+            gru_list.append(cell_state)
+
+        _ = tf.identity(cell_state, name="cell_state")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        output_names_with_port = ["cell_state:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-3)
+
 
 if __name__ == '__main__':
     Tf2OnnxBackendTestBase.trigger(GRUBlockTests)

tests/test_lstm.py

@@ -19,6 +19,7 @@
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test
 
+
 class LSTMTests(Tf2OnnxBackendTestBase):
     def test_test_single_dynamic_lstm_state_is_tuple(self):
         self.internal_test_single_dynamic_lstm(True)
@@ -334,6 +335,52 @@ def test_dynamic_basiclstm(self):
         output_names_with_port = ["output:0", "cell_state:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.0001)
 
+    def test_dynamic_lstm_output_consumed_only(self):
+        units = 5
+        batch_size = 6
+        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.LSTMCell(
+            units,
+            state_is_tuple=True)
+
+        outputs, _ = tf.nn.dynamic_rnn(
+            cell1,
+            x,
+            dtype=tf.float32)
+
+        _ = tf.identity(outputs, name="output")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        output_names_with_port = ["output:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.0001)
+
+    def test_dynamic_lstm_state_consumed_only(self):
+        units = 5
+        batch_size = 6
+        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.LSTMCell(
+            units,
+            state_is_tuple=True)
+
+        _, cell_state = tf.nn.dynamic_rnn(
+            cell1,
+            x,
+            dtype=tf.float32)
+
+        _ = tf.identity(cell_state, name="cell_state")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        output_names_with_port = ["cell_state:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, 0.0001)
+
     def test_dynamic_bilstm_state_is_tuple(self):
         self.internal_test_dynamic_bilstm_with_parameters(True)
 
@@ -409,6 +456,40 @@ def test_dynamic_bilstm_output_consumed_only(self, state_is_tuple=True):
         output_names_with_port = ["output:0"]
         self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
 
+    def test_dynamic_bilstm_state_consumed_only(self, state_is_tuple=True):
+        units = 5
+        batch_size = 6
+        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)
+
+        lstm_list = []
+        if True:
+            # bilstm, no scope
+            cell1 = rnn.LSTMCell(
+                units,
+                initializer=initializer,
+                state_is_tuple=state_is_tuple)  # state_is_tuple will impact Pack node (for cell_state)'s usage pattern
+            cell2 = rnn.LSTMCell(
+                units,
+                initializer=initializer,
+                state_is_tuple=state_is_tuple)
+            outputs, cell_state = tf.nn.bidirectional_dynamic_rnn(
+                cell1,
+                cell2,
+                x,
+                dtype=tf.float32)
+            lstm_list.append(outputs)
+
+        _ = tf.identity(cell_state, name="cell_state")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        output_names_with_port = ["cell_state:0"]
+        self.run_test_case(feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
+
 
 if __name__ == '__main__':
     Tf2OnnxBackendTestBase.trigger(LSTMTests)

tf2onnx/rewriter/bigru_rewriter.py

@@ -153,10 +153,12 @@ def rewrite_bidirectional_grus(g, ops):
             is_backward_gru = True
 
         if is_backward_gru:
-            # make sure reverse gru output will be reversed back
-            if get_reverse_nodes_after_y_output(g, n):
-                log.debug("find bw gru %s", input_id)
-                bw_gru[input_id] = [input_id, n]
+            # 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
+            log.debug("find bw gru %s", input_id)
+            bw_gru[input_id] = [input_id, n]
         else:
             log.debug("find fw gru %s", input_id)
             fw_gru[input_id] = [input_id, n]