Merge pull request #506 from lucienwang1009/multi_birnn

support multi birnn share the same input

Author: zhijxu-MS

tests/test_gru.py

@@ -482,6 +482,143 @@ 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_bigru_unknown_batch_size(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, [None, 3, 2], name="input_1")
+
+        cell1 = rnn.GRUCell(units)
+        cell2 = rnn.GRUCell(units)
+        _, cell_state = tf.nn.bidirectional_dynamic_rnn(
+            cell1,
+            cell2,
+            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, rtol=1e-06,
+                           graph_validator=lambda g: check_gru_count(g, 1))
+
+    def test_dynamic_bigru_outputs_partially_consumed(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.GRUCell(units)
+        cell2 = rnn.GRUCell(units)
+        (output_fw, _), (_, state_bw) = tf.nn.bidirectional_dynamic_rnn(
+            cell1,
+            cell2,
+            x,
+            dtype=tf.float32)
+
+        _ = tf.identity(output_fw, name="output")
+        _ = tf.identity(state_bw, name="cell_state")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        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-06,
+                           graph_validator=lambda g: check_gru_count(g, 1))
+
+    def test_dynamic_multi_bigru_with_same_input_hidden_size(self):
+        units = 5
+        batch_size = 10
+        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.GRUCell(units)
+        cell2 = rnn.GRUCell(units)
+        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)
+        cell2 = rnn.GRUCell(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))
+
+    def test_dynamic_multi_bigru_with_same_input_seq_len(self):
+        units = 5
+        batch_size = 10
+        x_val = np.array([[1., 1.], [2., 2.], [3., 3.]], dtype=np.float32)
+        x_val = np.stack([x_val] * batch_size)
+        seq_len_val = np.array([3], dtype=np.int32)
+
+        x = tf.placeholder(tf.float32, x_val.shape, name="input_1")
+
+        y1 = tf.placeholder(tf.int32, seq_len_val.shape, name="input_2")
+        seq_len1 = tf.tile(y1, [batch_size])
+        cell1 = rnn.GRUCell(units)
+        cell2 = rnn.GRUCell(units)
+        outputs_1, cell_state_1 = tf.nn.bidirectional_dynamic_rnn(
+            cell1,
+            cell2,
+            x,
+            sequence_length=seq_len1,
+            dtype=tf.float32,
+            scope="bigru_1"
+        )
+
+        y2 = tf.placeholder(tf.int32, seq_len_val.shape, name="input_3")
+        seq_len2 = tf.tile(y2, [batch_size])
+        cell1 = rnn.GRUCell(units)
+        cell2 = rnn.GRUCell(units)
+        outputs_2, cell_state_2 = tf.nn.bidirectional_dynamic_rnn(
+            cell1,
+            cell2,
+            x,
+            sequence_length=seq_len2,
+            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_2:0": seq_len_val, "input_3:0": seq_len_val}
+        input_names_with_port = ["input_1:0", "input_2:0", "input_3: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()

tests/test_lstm.py

@@ -554,6 +554,157 @@ 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_bilstm_outputs_partially_consumed(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)
+
+        # 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)
+        (output_fw, _), (_, state_bw) = tf.nn.bidirectional_dynamic_rnn(
+            cell1,
+            cell2,
+            x,
+            dtype=tf.float32)
+
+        _ = tf.identity(output_fw, name="output")
+        _ = tf.identity(state_bw, name="cell_state")
+
+        feed_dict = {"input_1:0": x_val}
+        input_names_with_port = ["input_1:0"]
+        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-06,
+                           graph_validator=lambda g: check_lstm_count(g, 1))
+
+    def test_dynamic_bilstm_unknown_batch_size(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, [None, 3, 2], 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)
+        _, cell_state = tf.nn.bidirectional_dynamic_rnn(
+            cell1,
+            cell2,
+            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, rtol=1e-06,
+                           graph_validator=lambda g: check_lstm_count(g, 1))
+
+    def test_dynamic_multi_bilstm_with_same_input_hidden_size(self):
+        units = 5
+        batch_size = 10
+        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)
+        cell2 = rnn.LSTMCell(units)
+        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)
+        cell2 = rnn.LSTMCell(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))
+
+    def test_dynamic_multi_bilstm_with_same_input_seq_len(self):
+        units = 5
+        batch_size = 10
+        x_val = np.array([[1., 1.], [2., 2.], [3., 3.]], dtype=np.float32)
+        x_val = np.stack([x_val] * batch_size)
+        seq_len_val = np.array([3], dtype=np.int32)
+
+        x = tf.placeholder(tf.float32, x_val.shape, name="input_1")
+
+        y1 = tf.placeholder(tf.int32, seq_len_val.shape, name="input_2")
+        seq_len1 = tf.tile(y1, [batch_size])
+        cell1 = rnn.LSTMCell(units)
+        cell2 = rnn.LSTMCell(units)
+        outputs_1, cell_state_1 = tf.nn.bidirectional_dynamic_rnn(
+            cell1,
+            cell2,
+            x,
+            sequence_length=seq_len1,
+            dtype=tf.float32,
+            scope="bilstm_1"
+        )
+
+        y2 = tf.placeholder(tf.int32, seq_len_val.shape, name="input_3")
+        seq_len2 = tf.tile(y2, [batch_size])
+        cell1 = rnn.LSTMCell(units)
+        cell2 = rnn.LSTMCell(units)
+        outputs_2, cell_state_2 = tf.nn.bidirectional_dynamic_rnn(
+            cell1,
+            cell2,
+            x,
+            sequence_length=seq_len2,
+            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_2:0": seq_len_val, "input_3:0": seq_len_val}
+        input_names_with_port = ["input_1:0", "input_2:0", "input_3: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()

tf2onnx/graph.py

@@ -179,17 +179,29 @@ def get_attr(self, name, default=None):
         attr = self.attr.get(name, default)
         return attr
 
+    def get_attr_value(self, name, default=None):
+        attr = self.get_attr(name)
+        if attr:
+            return helper.get_attribute_value(attr)
+        return default
+
     def get_attr_int(self, name):
         """Get attribute value as int."""
-        attr = self.get_attr(name)
-        utils.make_sure(attr is not None, "attribute %s is None", name)
-        attr = attr.i
-        return attr
+        attr_int = self.get_attr_value(name)
+        utils.make_sure(
+            attr_int is not None and isinstance(attr_int, int),
+            "attribute %s is None", name
+        )
+        return attr_int
 
     def get_attr_str(self, name, encoding="utf-8"):
         """Get attribute value as string."""
-        attr = self.get_attr(name)
-        return attr.s.decode(encoding) if attr else None
+        attr_str = self.get_attr_value(name)
+        utils.make_sure(
+            attr_str is not None and isinstance(attr_str, bytes),
+            "attribute %s is None", name
+        )
+        return attr_str.decode(encoding)
 
     def set_attr(self, name, value):
         self.attr[name] = helper.make_attribute(name, value)

tf2onnx/graph_matcher.py

@@ -20,13 +20,10 @@
 from __future__ import unicode_literals
 
 import copy
-import logging
 
 import six
 
 
-logger = logging.getLogger(__name__)
-
 
 class OpTypePattern(object):
     """A tree pattern that matches TF expressions with certain op types."""
@@ -161,12 +158,6 @@ def _match_pattern(self, pattern, op, tensor):
 
         if pattern.op_type != '*':
             if op is None or op.type not in pattern.op_type.split('|'):
-                logger.debug(
-                    "mismatched type at %s: [%s, %s]",
-                    op.name if op else "None",
-                    pattern.op_type,
-                    op.type if op else "None"
-                )
                 return False
 
         self._match_result.add(pattern, op, tensor)
@@ -177,12 +168,6 @@ def _match_pattern(self, pattern, op, tensor):
             return True
 
         if not op or len(op.inputs) != len(pattern.inputs):
-            logger.debug(
-                "mismatched input number at %s: [%s, %s]",
-                op.name if op else "None",
-                len(pattern.inputs),
-                len(op.inputs)
-            )
             return False
 
         if self._allow_reorder:
@@ -219,7 +204,6 @@ def match_op(self, op):
           Returns a `MatchResult` if `op` matches the pattern; otherwise, returns
           None.
         """
-        logger.debug("match %s against the pattern", op.name)
         self._match_result = MatchResult()
         if not self._match_pattern(self._pattern, op, tensor=None):
             return None