LSTM.weight_ih_l[k] dimensions with proj_size

user_10797659 0
Viewed 476

According to Pytorch LSTM documentation :-

  • ~LSTM.weight_ih_l[k] – the learnable input-hidden weights of the kth layer (W_ii|W_if|W_ig|W_io), of shape (4*hidden_size, input_size) for k = 0. Otherwise, the shape is (4 * hidden_size, num_directions * hidden_size)

My doubt is, why for k > 0 the shape for each weight is (hidden_size, num_directions * hidden_size), according to me, shouldn't be (hidden_size, num_directions * proj_size) because the layer above the lowest layer is receiving the input which is the output of the lowest layer which have the shape of (L, N, num_directions*proj_size)

UPDATE

Indeed, the documentation have mentioned wrong shapes. It will be fixed soon.

Github Issue

lstm pytorch python
1 Answers

This is now fixed. The OP opened the Issue #65053 that was fixed by PR #65102 (commit 83878e1).

It just happens that the documentation isn't providing all the details in this case, and you're correct. In fact, you can check in the source code that the shape of W_ih is (4*hidden_size, num_directions * proj_size) when proj_size > 0 for k > 0:

# [...]
if mode == 'LSTM':
    gate_size = 4 * hidden_size
elif mode == 'GRU':
    gate_size = 3 * hidden_size
elif mode == 'RNN_TANH':
    gate_size = hidden_size
elif mode == 'RNN_RELU':
    gate_size = hidden_size
else:
    raise ValueError("Unrecognized RNN mode: " + mode)

self._flat_weights_names = []
self._all_weights = []
for layer in range(num_layers):
    for direction in range(num_directions):
        real_hidden_size = proj_size if proj_size > 0 else hidden_size
        layer_input_size = input_size if layer == 0 else real_hidden_size * num_directions

        w_ih = Parameter(torch.empty((gate_size, layer_input_size), **factory_kwargs))
        # [...]

As you can see, w_ih has the shape (gate_size, layer_input_size), where:

  • gate_size is 4 * hidden_size for LSTM, and
  • layer_input_size is
    • input_size if layer == 0 (layer is equivalent to k in the docs), else
    • real_hidden_size * num_directions for k > 0, and
  • real_hidden_size = proj_size if proj_size > 0, else it is hidden_size.

That is: if proj_size > 0 and layer > 0, layer_input_size = proj_size * num_directions, and the shape of w_ih will be equal to (4 * hidden_size, proj_size * num_directions.

It is worth noting that they do have the following in the documentation:

If proj_size > 0 is specified, LSTM with projections will be used. This changes the LSTM cell in the following way. First, the dimension of h_t will be changed from hidden_size to proj_size (dimensions of W_hi will be changed accordingly). Second, the output hidden state of each layer will be multiplied by a learnable projection matrix: h_t = W_hr * h_t. Note that as a consequence of this, the output of LSTM network will be of different shape as well.

user_42282750
Related