fix rnn lstm gru bug

Author: hanjr92

examples/basic_tutorials/imdb_LSTM_simple.py

@@ -4,14 +4,16 @@
 # The same set of code can switch the backend with one line
 import os
 # os.environ['TL_BACKEND'] = 'tensorflow'
-os.environ['TL_BACKEND'] = 'mindspore'
-# os.environ['TL_BACKEND'] = 'paddle'
+# os.environ['TL_BACKEND'] = 'mindspore'
+os.environ['TL_BACKEND'] = 'paddle'
 # os.environ['TL_BACKEND'] = 'torch'
 import tensorlayerx as tlx
 from tensorlayerx.nn import Module
-from tensorlayerx.nn import Linear, LSTM, Embedding
+from tensorlayerx.nn import Linear, LSTM, Embedding, RNN
 from tensorlayerx.dataflow import Dataset
 import numpy as np
+prev_h = np.random.random([1, 200, 64]).astype(np.float32)
+prev_h = tlx.convert_to_tensor(prev_h)
 
 X_train, y_train, X_test, y_test = tlx.files.load_imdb_dataset('data', nb_words=20000, test_split=0.2)
 
@@ -48,7 +50,7 @@ def __init__(self):
 
     def forward(self, x):
         x = self.embedding(x)
-        x, _ = self.lstm(x)
+        x, _ = self.lstm(x, [prev_h, prev_h])
         x = tlx.reduce_mean(x, axis=1)
         x = self.linear1(x)
         x = self.linear2(x)

tensorlayerx/backend/ops/paddle_nn.py

@@ -1638,10 +1638,16 @@ def _cudnn_impl(self, inputs, initial_states, sequence_length):
         out = pd.tensor.transpose(out, [1, 0, 2]) if not self.time_major else out
         return out, tuple(state) if len(state) > 1 else state[0]
 
+    def check_hidden(self, h, batch_size):
+        expected_hidden_size = [self.num_layers * self.bidirect, batch_size, self.hidden_size]
+        if h.shape != expected_hidden_size:
+            raise ValueError('Expected hidden size {}, got {}.'.format(expected_hidden_size, h.shape))
+
     def forward(self, inputs, initial_states=None):
         batch_index = 1 if self.time_major else 0
         dtype = inputs.dtype
         sequence_length = None
+        batch_size = inputs.shape[batch_index]
         if initial_states is None:
             state_shape = (self.num_layers * self.bidirect, -1, self.hidden_size)
             if self.state_components == 1:
@@ -1655,6 +1661,15 @@ def forward(self, inputs, initial_states=None):
                         for _ in range(self.state_components)
                     ]
                 )
+        else:
+            if self.mode == 'LSTM':
+                h, c = initial_states
+                self.check_hidden(h, batch_size)
+                self.check_hidden(c, batch_size)
+            else:
+                self.check_hidden(initial_states, batch_size)
+        if not isinstance(initial_states, (tuple, list)):
+            initial_states = [initial_states,]
         if self.could_use_cudnn:
             # Add CPU kernel and dispatch in backend later
             return self._cudnn_impl(inputs, initial_states, sequence_length)

tensorlayerx/backend/ops/torch_nn.py

@@ -1934,7 +1934,7 @@ def check_input(self, input_shape):
             )
 
     def check_hidden(self, h, batch_size):
-        expected_hidden_size = (self.num_layers * self.bidirect, batch_size, self.hidden_size)
+        expected_hidden_size = (self.num_layers * self.num_directions, batch_size, self.hidden_size)
         if h.shape != expected_hidden_size:
             raise ValueError('Expected hidden size {}, got {}.'.format(expected_hidden_size, h.shape))