Get the last hidden state in stacked LSTM

macc_n: <p>Hello everyone,</p>
<p>I’m developing a classifier based on LSTM and I defined the model in this way:</p>
<pre><code class="lang-auto">class LSTMClassifier(nn.Module):

    def __init__(self, input_size, seq_len, n_classes, hidden_size, device, dim_feedforward=1024, num_layers=1, dropout=0, bidirectional=False, batch_first=True):
        super(LSTMClassifier, self).__init__()

        self.hidden_size = hidden_size
        self.seq_len = seq_len
        self.input_size = input_size
        self.num_layers = num_layers
        self.device = device

        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers, dropout=dropout, bidirectional=bidirectional, batch_first=batch_first)
        self.classifier = nn.Linear(hidden_size, n_classes)
    
    def init_hidden_state(self, x):
        return (
            torch.zeros(self.num_layers, x.size()[0], self.hidden_size).to(self.device),
            torch.zeros(self.num_layers, x.size()[0], self.hidden_size).to(self.device)
        )
    
    def forward(self, x):
        self.hidden = self.init_hidden_state(x)

        out, (hn, cn) = self.lstm(x, self.hidden)
        hn = hn.view(-1, self.hidden_size)
        y = self.classifier(hn)
        
        return y
</code></pre>
<p>It should be correct. If I create a model with hidden_size=64 and num_layers=1, the size of hn (before view) is [1,1,64] when I give in input 1 sample.<br>
If I change num_layers=2, the size of hn (before view) is [2,1,64] because it concatenates the hidden state of the two stacked layers.</p>
<p>My doubt is: what is the order of hidden states in hn? The first element in hn is the hidden state of the first or of the last LSTM layer?<br>
Is it correct to consider the last hidden state of the last LSTM layer to classify the sequences?</p>
<p>Thanks for your help.</p>

achingacham: <p>Hi there!</p>
<p>I am not sure if this response is still useful. But sharing what I understand about getting the last hidden state from stacked LSTM for future references. If I get the question right, you are interested in using the last hidden states of each element in a batch for a classification task. The variable <code>out</code> in the given code already has that information. In the case of a uni-directional LSTM, it is straightforward. Using <code>out[:, -1]</code> when <code>batch_first=True</code> should work. <code>-1</code> index refers to the last item in the chain of hidden states. <code>hn[-1]</code> before updating it with <em>view()</em> should also give the same tensor.</p>
<p>However, in the case of bidirectional, follow the note given in the PyTorch documentation:</p>
<blockquote>
<p>For bidirectional LSTMs, forward and backward are directions 0 and 1 respectively. Example of splitting the output layers when <code>batch_first=False</code>: <code>output.view(seq_len, batch, num_directions, hidden_size)</code>.</p>
</blockquote>
<p>Also, for reshaping <code>hn</code>, I will recommend <code>hn.view(num_layers, num_directions, batch, hidden_size)</code> and use <code>-1</code> index for the last layer and <code>0/1</code> for selecting the direction.</p>
<p>Finally, about the question of the correctness of choosing the last layer for classification, I believe this can only be answered through empirical analysis. It depends heavily on the task at hand.</p>
<p>Hope this helps.</p>

vdw: <p>You can have a look a <a href="https://github.com/chrisvdweth/ml-toolkit/blob/master/pytorch/models/text/classifier/rnn.py" rel="noopener nofollow ugc">this code</a>; it handles multiple layers as well as 1 or 2 directions. The key snippet is:</p>
<pre><code class="lang-auto"># Extract last hidden state
if self.params.rnn_type == RnnType.GRU:
    final_state = self.hidden.view(self.params.num_layers, self.num_directions, batch_size, self.params.rnn_hidden_dim)[-1]
elif self.params.rnn_type == RnnType.LSTM:
    final_state = self.hidden[0].view(self.params.num_layers, self.num_directions, batch_size, self.params.rnn_hidden_dim)[-1]
# Handle directions
final_hidden_state = None
if self.num_directions == 1:
    final_hidden_state = final_state.squeeze(0)
elif self.num_directions == 2:
    h_1, h_2 = final_state[0], final_state[1]
    # final_hidden_state = h_1 + h_2               # Add both states (requires changes to the input size of first linear layer + attention layer)
    final_hidden_state = torch.cat((h_1, h_2), 1)  # Concatenate both states
</code></pre>