Properly shaping LSTM input/output for text sequence prediction?

Viewed 146

I'm trying to learn how to use an LSTM to predict the next character based on the previous character. I've successfully created a function that batches a long string of text (after one-hot encoding) into shape (batch_size, seq_len, one_hot_features). one_hot_features is simply the number of unique characters in the text.

Since my batching function gives me both training sequences and 'ground truth' sequences, I need the output of the LSTM network to be the same shape as the 'ground truth' batches so that I can plug both output and labels into the loss function.

So my questions are:

  1. How do I correctly define my network architecture to meet the above requirements.
  2. How do I shape my data as it passes through the various layers in the forward function?

This is my attempt, and the code will run, but the output of the network is incorrectly shaped for comparison to the 'ground truth' sequence batches:

class charNN(nn.Module):
    def __init__(self, vocab, hidden_size, n_layers, dropout=0.5):
        super().__init__()

        self.vocab_length = len(vocab)
        self.hidden_size = hidden_size
        self.lstm = nn.LSTM(self.vocab_length, hidden_size, n_layers, batch_first=True)
        self.fc   = nn.Linear(hidden_size, self.vocab_length)
        
    def forward(self, x, hidden=None):
        x, hidden = self.lstm(x, hidden)     # -> (n_batches, seq_len, hidden_size)
        x = x.reshape(-1, self.hidden_size)  # -> (n_batches * seq_len, hidden_size)
        x = self.fc(x)                       # -> (n_batches * seq_len, vocab_length)
        
        return x, hidden    

Thank you in advance. I've been struggling with this one for a couple days. I will also happily provide the code for training the network, if it is required. It does not throw any errors except to tell me that the output size and the true label sizes do not match.

1 Answers

Your input_size: len(vocab) and you have also stated adequate amount information for __init__ method.

I'm also a beginner on the lstm.

Answer 1:

According to me, you have already correctly defined your architecture. Maybe you could add the followings:

  • token_size variable: your output size

  • encoder and decoder structure.

  • Method for initializing weights

For instance:

class charNN(nn.Module):
    def __init__(self, vocab, hidden_size, n_layers, dropout=0.5, token_size):
        super(charNN, self).__init__()
        self.token_size = token_size
        self.drop = nn.Dropout(dropout)
        self.vocab_length = len(vocab)
        self.hidden_size = hidden_size
        self.encoder = nn.Embedding(token_size, vocab_length)
        self.lstm = nn.LSTM(self.vocab_length,   
                            hidden_size, 
                            n_layers, 
                            batch_first=True, 
                            dropout=dropout)
        self.fc   = nn.Linear(hidden_size,  
                              self.vocab_length)
        self.decoder = nn.Linear(hidden_size, token_size)
        self.init_weights()
        self.hidden_size = hidden_size
        self.num_layers = num_layers

    def init_weights(self):
        init_range = 0.1
        uniform_(self.encoder.weight, -init_range, init_range)
        zeros_(self.decoder.weight)
        uniform_(self.decoder.weight, -init_range, init_range)
        

    def init_hidden(self, batch_size):
        weight = next(self.parameters())
        return (weight.new_zeros(self.num_layers, batch_size,
                                 self.hidden_size),
                weight.new_zeros(self.num_layers, batch_size,
                                 self.hidden_size))

Answer 2:

You could use the following structure for the forward function:

    def forward(self, input_, hidden_):
        embedded = self.drop(self.encoder(input_))
        output, hidden_ = self.rnn(input_, hidden_)
        output = self.drop(output)
        decoded = self.decoder(output)
        decoded = decoded.view(-1, self.token_size)
        output = log_softmax(decoded, dim=1)
        return output, hidden_

Hope it helps, good luck!

Related