Repository: pytorch/pytorch
Pull request: Add rnn args check
Fixes #3851, #3259

Added a high-level check for arguments to RNNBase (these were moved from arg checks in `cudnn/rnn.py`)

### Test Plan
`python test/test_nn.py`

Problem statement:
Improve error message for RuntimeError: inconsistent tensor size in RNN
I run RNN sucessfully with the code.
```
rnn = nn.RNN(features, hidden_size, num_layers)
o, hidden = rnn(x, hidden)
```
and I run embedding example. Then, got RuntimeError.
```
RuntimeError: inconsistent tensor size, expected r_ [5 x 25], t [5 x 25] and src [2 x 25] to 
have thesame number of elements, but got 125, 125 and 50 elements respectively at 
/opt/conda/conda-bld/pytorch_1503970438496/work/torch/lib/TH/generic/THTensorMath.c:887
```
I set `batch_first=True` to use `Embedding`
```
x.size is [5, 8, 15]
hidden_size is [5, 2, 25]
```
nothing seems wrong.
And the error info seems not the dim problem?
The error is so strange.

Hints from issue discussion:
What exactly do you mean by "embedding example"? Could you provide a code snippet that causes the RuntimeError?
@zou3519 
To better reproduce the error, I use the official code example of [torch.nn.RNN](http://pytorch.org/docs/master/nn.html?highlight=rnn#recurrent-layers)
The origin runs well:
```
rnn = nn.RNN(10, 20, 2, batch_first=False)
input = Variable(torch.randn(5, 3, 10))
h0 = Variable(torch.randn(2, 3, 20))
output, hn = rnn(input, h0)
output.size(), hn.size()
```
Due to use `embedding`, I have to make batch first.
And then , it raises error.
```
rnn = nn.RNN(10, 20, 2, batch_first=True)
input = Variable(torch.randn(5, 3, 10).transpose(0, 1))
h0 = Variable(torch.randn(2, 3, 20).transpose(0, 1))
output, hn = rnn(input, h0)
output.size(), hn.size()
```

```
RuntimeError: inconsistent tensor size, expected r_ [3 x 20], t [3 x 20] and src [2 x 20] to
 have the same number of elements, but got 60, 60 and 40 elements respectively at 
/opt/conda/conda-bld/pytorch_1503970438496/work/torch/lib/TH/generic/THTensorMath.c:887
```
Thank you for the example!

The [docs](http://pytorch.org/docs/master/nn.html?highlight=rnn#torch.nn.RNN) say that setting `batch_first=True` means that the first dim of input is the batch, ie, input has size `(batch, seq_len, num_features)`. `h0` however is unchanged, so the following should work:

```
rnn = nn.RNN(10, 20, 2, batch_first=True)
input = Variable(torch.randn(5, 3, 10).transpose(0, 1))
h0 = Variable(torch.randn(2, 3, 20))
output, hn = rnn(input, h0)
output.size(), hn.size()
```

I do agree the error message could be better, though

Commit message: Add rnn args check

Patch to review:
diff --git a/torch/backends/cudnn/rnn.py b/torch/backends/cudnn/rnn.py
index 51da709d8be5..ab17d7761984 100644
--- a/torch/backends/cudnn/rnn.py
+++ b/torch/backends/cudnn/rnn.py
@@ -203,13 +203,6 @@ def forward(fn, input, hx, weight, output, hy):
         if fn.batch_first and not is_input_packed:
             input = input.transpose(0, 1)
 
-        if (not is_input_packed and input.dim() != 3) or (is_input_packed and input.dim() != 2):
-            raise RuntimeError(
-                'input must have 3 dimensions, got {}'.format(input.dim()))
-        if fn.input_size != input.size(-1):
-            raise RuntimeError('input.size(-1) must be equal to input_size. Expected {}, got {}'.format(
-                fn.input_size, input.size(-1)
-            ))
         if fn.dropout != 0 and cudnn.version() < 5103:
             raise RuntimeError('dropout supported only in cudnn v5.1 and above')
 
@@ -261,9 +254,6 @@ def forward(fn, input, hx, weight, output, hy):
             fn.w_desc = init_weight_descriptor(fn, fn.weight_buf)
             w = fn.weight_buf
 
-        if tuple(hx.size()) != hidden_size:
-            raise RuntimeError('Expected hidden size {}, got {}'.format(
-                hidden_size, tuple(hx.size())))
         if cx is not None and tuple(cx.size()) != hidden_size:
             raise RuntimeError('Expected cell size {}, got {}'.format(
                 hidden_size, tuple(cx.size())))
diff --git a/torch/nn/modules/rnn.py b/torch/nn/modules/rnn.py
index bbab30708a33..7c6b81fb7928 100644
--- a/torch/nn/modules/rnn.py
+++ b/torch/nn/modules/rnn.py
@@ -128,6 +128,33 @@ def reset_parameters(self):
         for weight in self.parameters():
             weight.data.uniform_(-stdv, stdv)
 
+    def check_forward_args(self, input, hidden, batch_sizes):
+        is_input_packed = batch_sizes is not None
+        expected_input_dim = 2 if is_input_packed else 3
+        if input.dim() != expected_input_dim:
+            raise RuntimeError(
+                'input must have {} dimensions, got {}'.format(
+                    expected_input_dim, input.dim()))
+        if self.input_size != input.size(-1):
+            raise RuntimeError(
+                'input.size(-1) must be equal to input_size. Expected {}, got {}'.format(
+                    fn.input_size, input.size(-1)))
+
+        if is_input_packed:
+            mini_batch = batch_sizes[0]
+        else:
+            mini_batch = input.size(0) if self.batch_first else input.size(1)
+
+        num_directions = 2 if self.bidirectional else 1
+        expected_hidden_size = (self.num_layers * num_directions,
+                                mini_batch, self.hidden_size)
+        if self.mode == 'LSTM':
+            hidden = hidden[0]
+        if tuple(hidden.size()) != expected_hidden_size:
+            raise RuntimeError(
+                'Expected hidden size {}, got {}'.format(
+                    expected_hidden_size, tuple(hidden.size())))
+
     def forward(self, input, hx=None):
         is_packed = isinstance(input, PackedSequence)
         if is_packed:
@@ -153,6 +180,8 @@ def forward(self, input, hx=None):
             flat_weight = first_data.new().set_(first_data.storage(), 0, torch.Size([self._param_buf_size]))
         else:
             flat_weight = None
+
+        self.check_forward_args(input, hx, batch_sizes)
         func = self._backend.RNN(
             self.mode,
             self.input_size,


Code review comments:
diff_hunk: @@ -128,6 +128,33 @@ def reset_parameters(self):
         for weight in self.parameters():
             weight.data.uniform_(-stdv, stdv)
 
+    def check_forward_args(self, input, hidden, batch_sizes):
+        is_input_packed = batch_sizes is not None
+        expected_input_dim = 2 if is_input_packed else 3
+        if input.dim() != expected_input_dim:
+            raise RuntimeError(
+                'input must have {} dimensions, got {}'.format(
+                    expected_input_dim, input.dim()))
+        if self.input_size != input.size(-1):
+            raise RuntimeError(
+                'input.size(-1) must be equal to input_size. Expected {}, got {}'.format(
+                    fn.input_size, input.size(-1)))
+
+        if is_input_packed:
+            mini_batch = batch_sizes[0]
+        else:
+            mini_batch = input.size(0) if self.batch_first else input.size(1)
+
+        num_directions = 2 if self.bidirectional else 1
+        expected_hidden_size = (self.num_layers * num_directions,
+                                mini_batch, self.hidden_size)
+        if self.mode == 'LSTM':
+            hidden = hidden[0]
+        if tuple(hidden.size()) != expected_hidden_size:
path: torch/nn/modules/rnn.py
text: @user1:
You should test both in case of LSTM (they have the same size)