Hidden state tensors have a different order than the returned tensors

As part of GRU training, I want to retrieve the hidden state tensors.


I have defined a GRU with two layers:

self.lstm = nn.GRU(params.vid_embedding_dim, params.hidden_dim , 2)


The forward function is defined as follows (the following is just a part of the implementation):

    def forward(self, s, order, batch_size, where, anchor_is_phrase = False):
    """
    Forward prop. 
    """
      # s is of shape [128 , 1 , 300] , 128 is batch size
      output, (a,b) = self.lstm(s.cuda())
      output.data.contiguous()


And out is of shape: [128 , 400] (128 is the number of samples which each one is embedded in 400 dimensional vector).

I understand that out is the output of the last hidden state and thus I expect it to be equal to b. However, after I checked the values I saw that it's indeed equal but b contains the tensor in a different order, that is for example output[0] is b[49]. Am I missing something here ?

Thanks.

I understand your confusion. Have a look the example bellow and the comments:

# [Batch size, Sequence length, Embedding size]
inputs = torch.rand(128, 5, 300)
gru = nn.GRU(input_size=300, hidden_size=400, num_layers=2, batch_first=True)

with torch.no_grad():
    # output is all hidden states, for each element in the batch of the last layer in the RNN
    # a is the last hidden state of the first layer
    # b is the last hidden state of the second (last) layer
    output, (a, b) = gru(inputs)


If we print out the shapes, they will confirm our understanding:

print(output.shape) # torch.Size([128, 5, 400])
print(a.shape) # torch.Size([128, 400])
print(b.shape) # torch.Size([128, 400])


Also, we can test whether the last hidden state, for each element in the batch, of the last layer, obtained from output is equal to b:

np.testing.assert_almost_equal(b.numpy(), output[:,:-1,:].numpy())


Finally, we can create an RNN with 3 layers, and run the same tests:

gru = nn.GRU(input_size=300, hidden_size=400, num_layers=3, batch_first=True)
with torch.no_grad():
    output, (a, b, c) = gru(inputs)

np.testing.assert_almost_equal(c.numpy(), output[:,-1,:].numpy())


Again, the assertion passes but only if we do it for c, which is now the last layer of the RNN. Otherwise:

np.testing.assert_almost_equal(b.numpy(), output[:,-1,:].numpy())


Raises an error:


  AssertionError:  Arrays are not almost equal to 7 decimals


I hope that this makes things clear for you.
