fix bug of lstm

Author: zhijxu-MS

tests/test_gru.py

@@ -123,34 +123,36 @@ def test_single_dynamic_gru_seq_length_is_const(self):
                            graph_validator=lambda g: check_gru_count(g, 1))
 
     def test_single_dynamic_gru_seq_length_is_not_const(self):
-        units = 5
-        batch_size = 1
-        x_val = np.array([[1., 1.], [2., 2.], [3., 3.], [4., 4.], [5., 5.]], 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)
-
-        y_val = np.array([5], dtype=np.int32)
-        seq_length = tf.placeholder(tf.int32, y_val.shape, name="input_2")
-
-        # no scope
-        cell = rnn.GRUCell(
-            units,
-            kernel_initializer=initializer)
-        outputs, cell_state = tf.nn.dynamic_rnn(
-            cell,
-            x,
-            dtype=tf.float32,
-            sequence_length=tf.identity(seq_length))
+        for np_dtype, tf_dtype in [[np.int32, tf.int32], [np.int64, tf.int64], [np.float32, tf.float32]]:
+            tf.reset_default_graph()
+            units = 5
+            batch_size = 1
+            x_val = np.array([[1., 1.], [2., 2.], [3., 3.], [4., 4.], [5., 5.]], 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)
+
+            y_val = np.array([5], dtype=np_dtype)
+            seq_length = tf.placeholder(tf_dtype, y_val.shape, name="input_2")
+
+            # no scope
+            cell = rnn.GRUCell(
+                units,
+                kernel_initializer=initializer)
+            outputs, cell_state = tf.nn.dynamic_rnn(
+                cell,
+                x,
+                dtype=tf.float32,
+                sequence_length=tf.identity(seq_length))
 
-        _ = tf.identity(outputs, name="output")
-        _ = tf.identity(cell_state, name="cell_state")
+            _ = tf.identity(outputs, name="output")
+            _ = tf.identity(cell_state, name="cell_state")
 
-        feed_dict = {"input_1:0": x_val, "input_2:0": y_val}
-        input_names_with_port = ["input_1:0", "input_2: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-03, atol=1e-06,
-                           graph_validator=lambda g: check_gru_count(g, 1))
+            feed_dict = {"input_1:0": x_val, "input_2:0": y_val}
+            input_names_with_port = ["input_1:0", "input_2: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-03, atol=1e-06,
+                               graph_validator=lambda g: check_gru_count(g, 1))
 
     def test_single_dynamic_gru_placeholder_input(self):
         units = 5

tests/test_lstm.py

@@ -147,36 +147,38 @@ def test_single_dynamic_lstm_seq_length_is_const(self):
                            graph_validator=lambda g: check_lstm_count(g, 1))
 
     def test_single_dynamic_lstm_seq_length_is_not_const(self):
-        units = 5
-        batch_size = 6
-        x_val = np.array([[1., 1.], [2., 2.], [3., 3.], [4., 4.], [5., 5.]], dtype=np.float32)
-        x_val = np.stack([x_val] * batch_size)
-        state_is_tuple = True
-        x = tf.placeholder(tf.float32, x_val.shape, name="input_1")
-        initializer = init_ops.constant_initializer(0.5)
-
-        y_val = np.array([4, 3, 4, 5, 2, 1], dtype=np.int32)
-        seq_length = tf.placeholder(tf.int32, y_val.shape, name="input_2")
-
-        # no scope
-        cell = rnn.LSTMCell(
-            units,
-            initializer=initializer,
-            state_is_tuple=state_is_tuple)
-        outputs, cell_state = tf.nn.dynamic_rnn(
-            cell,
-            x,
-            dtype=tf.float32,
-            sequence_length=tf.identity(seq_length))
+        for np_dtype, tf_dtype in [[np.int32, tf.int32], [np.int64, tf.int64], [np.float32, tf.float32]]:
+            tf.reset_default_graph()
+            units = 5
+            batch_size = 6
+            x_val = np.array([[1., 1.], [2., 2.], [3., 3.], [4., 4.], [5., 5.]], dtype=np.float32)
+            x_val = np.stack([x_val] * batch_size)
+            state_is_tuple = True
+            x = tf.placeholder(tf.float32, x_val.shape, name="input_1")
+            initializer = init_ops.constant_initializer(0.5)
+
+            y_val = np.array([4, 3, 4, 5, 2, 1], dtype=np_dtype)
+            seq_length = tf.placeholder(tf_dtype, y_val.shape, name="input_2")
+
+            # no scope
+            cell = rnn.LSTMCell(
+                units,
+                initializer=initializer,
+                state_is_tuple=state_is_tuple)
+            outputs, cell_state = tf.nn.dynamic_rnn(
+                cell,
+                x,
+                dtype=tf.float32,
+                sequence_length=tf.identity(seq_length))
 
-        _ = tf.identity(outputs, name="output")
-        _ = tf.identity(cell_state, name="cell_state")
+            _ = tf.identity(outputs, name="output")
+            _ = tf.identity(cell_state, name="cell_state")
 
-        feed_dict = {"input_1:0": x_val, "input_2:0": y_val}
-        input_names_with_port = ["input_1:0", "input_2: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))
+            feed_dict = {"input_1:0": x_val, "input_2:0": y_val}
+            input_names_with_port = ["input_1:0", "input_2: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_single_dynamic_lstm_placeholder_input(self):
         units = 5

tf2onnx/rewriter/unit_rewriter_base.py

@@ -215,21 +215,21 @@ def find_sequence_length_node(self, rnn_scope_name):
         seq_len_node_cnt = len(seq_len_nodes)
         if seq_len_node_cnt == 0:
             return None
+
         if seq_len_node_cnt == 1:
             seq_len_node = seq_len_nodes[0]
             if seq_len_node.is_const():
                 return seq_len_node
-            # input of the "identity" node may be a "cast"
-            # if so, then we have to keep it
-            # sentence "math_ops.to_int32(sequence_length)" in tf results in the "cast" op
-            if seq_len_node.inputs[0].type == "Cast":
-                cast_node = seq_len_node.inputs[0]
-                if not cast_node.inputs[0].name.startswith(rnn_scope_name):
-                    return seq_len_node.inputs[0]
-                raise ValueError("sequence length node should be outside of rnn scope")
+
             if not seq_len_node.inputs[0].name.startswith(rnn_scope_name):
                 return seq_len_node.inputs[0]
-            raise ValueError("sequence length node should be outside of rnn scope")
+
+            # input of the "identity" node may be a "cast" op generated by "math_ops.to_int32(sequence_length)" in tf
+            # if so, then we have to find cast input as sequence node.
+            node = seq_len_node.inputs[0]
+            if node.type == "Cast" and not node.inputs[0].name.startswith(rnn_scope_name):
+                return node.inputs[0]
+
         raise ValueError("there are more sequence length nodes than expected")
 
     def get_rnn_input_blacklist(self, rnn_weights, rnn_props):
@@ -342,6 +342,13 @@ def process_seq_length(self, rnn_props, seq_length_node):
                                                attr={"to": onnx_pb.TensorProto.INT32})
 
             self.all_nodes.extend([timestep_node, tile_node, seq_length_node])
+        else:
+            # LSTM sequence_lens needs to be int32
+            ori_seq_dtype = self.g.get_dtype(seq_length_node.name)
+            if ori_seq_dtype != onnx_pb.TensorProto.INT32:
+                seq_length_node = self.g.make_node('Cast', [seq_length_node.output[0]],
+                                                   attr={"to": onnx_pb.TensorProto.INT32})
+                self.all_nodes.append(seq_length_node)
 
         rnn_props.onnx_input_ids["sequence_lens"] = seq_length_node.output[0]
         return seq_length_node, batchsize_node