Implementing LSTM Models in PyTorch for Sequence Classification

Hello! Welcome to our next lesson.

In our previous lessons, we've built a strong foundation for sequence modeling. We learned how to process raw text through tokenization, and then how to convert those tokens into meaningful, dense vectors using embedding layers. You now understand the what and the why of preparing text data for a neural network.

Today, we'll take the final and most exciting step in this part of our journey: assembling these components into a complete, functional LSTM-based model for sequence classification using PyTorch. We will focus on defining the model's architecture itself.

By the end of this lesson, you will be able to:

• Structure a sequence model in PyTorch by subclassing nn.Module.
• Define the necessary layers in the __init__ method: an embedding layer, an LSTM layer, and a final linear layer for classification.
• Implement the forward pass to correctly pipe data through the network.
• Extract the appropriate LSTM output for a many-to-one classification task.

This lesson is all about constructing the blueprint of our classifier. In the next lesson, we will use this blueprint to train the model on a real dataset.

The Anatomy of an LSTM Classifier

Before we write any code, let's visualize the architecture of a typical LSTM-based text classifier. It follows a "many-to-one" pattern: it processes a sequence of many inputs (the words in a sentence) to produce a single output (a class label, like "positive" or "negative" sentiment).

This architecture consists of three main stages:

1. Embedding Layer (nn.Embedding): As we covered in the last lesson, this layer acts as a lookup table. It takes the integer IDs of the tokens in our input sentence and converts each one into a dense vector representation.
2. LSTM Layer (nn.LSTM): This is the core of our model. It processes the sequence of embedding vectors one by one, updating its internal hidden state and cell state at each step to capture contextual information from the sequence.
3. Classifier Head (nn.Linear): This is a standard fully-connected layer. It takes the final, information-rich output from the LSTM layer and projects it down to the number of output classes we want to predict.

Let's briefly review the LSTM's mechanism for handling sequences.

As the article mentions, the key is that the LSTM maintains a memory of the sequence, and we use its final state to make our prediction. Now, let's translate this architecture into PyTorch code.

Implementing the Model with nn.Module

In PyTorch, all neural network models are created by creating a class that inherits from torch.nn.Module. This base class provides all the necessary functionality for tracking parameters, moving them between CPU/GPU, etc.

To create our custom model, we need to implement two essential methods:

• __init__(): The constructor, where we define and initialize all the layers our network will use (e.g., embedding, LSTM, linear).
• forward(): This method defines the "forward pass"—how an input tensor flows through the layers we defined in __init__ to produce an output.

Let's see how this is done with a practical example from the DataCamp article.

1. The __init__ Method: Defining the Layers

In the __init__ method, we instantiate the layers we'll need and assign them as attributes of the class (e.g., self.lstm).

Here's a typical __init__ block for an LSTM classifier:

import torch
import torch.nn as nn

class LSTMClassifier(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, n_layers, dropout_prob=0.5):
        super(LSTMClassifier, self).__init__()
        
        self.hidden_dim = hidden_dim
        self.n_layers = n_layers




        # 1. Embedding Layer
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        



        # 2. LSTM Layer
        self.lstm = nn.LSTM(input_size=embedding_dim,
                            hidden_size=hidden_dim,
                            num_layers=n_layers,
                            batch_first=True,  # Crucial for intuitive tensor shapes
                            dropout=dropout_prob if n_layers > 1 else 0)
        



        # 3. Dropout Layer for regularization
        self.dropout = nn.Dropout(dropout_prob)
        



        # 4. Fully Connected Layer (Classifier Head)
        self.fc = nn.Linear(hidden_dim, output_dim)
        



        # 5. Sigmoid activation for binary classification
        self.sig = nn.Sigmoid()

Let's break down the nn.LSTM parameters:

• input_size: The size of the features at each time step. For us, this is the embedding_dim, since we feed embedding vectors into the LSTM.
• hidden_size: The number of features in the hidden state h. This is a hyperparameter you choose.
• num_layers: The number of LSTM layers to stack on top of each other.
• batch_first=True: This is a very important argument. It tells PyTorch that our input tensors will have the batch dimension first, i.e., their shape will be (batch_size, sequence_length, feature_size). This is generally more intuitive to work with than the default (sequence_length, batch_size, feature_size).

2. The init_hidden Helper Method

The LSTM requires an initial hidden state (h_0) and cell state (c_0) to start its sequential processing. It's a common practice to create a helper method to initialize these states as zero-tensors. The SentimentRNN class from the DataCamp article does this cleanly.

    def init_hidden(self, batch_size):



        # Use the current device of the model
        device = next(self.parameters()).device
        



        # Create zero tensors for the hidden and cell states
        h0 = torch.zeros((self.n_layers, batch_size, self.hidden_dim)).to(device)
        c0 = torch.zeros((self.n_layers, batch_size, self.hidden_dim)).to(device)
        



        # The LSTM layer expects the hidden state as a tuple (h_0, c_0)
        hidden = (h0, c0)
        return hidden

This method takes the batch_size as input because the hidden state dimensions depend on it.

3. The forward Method: Defining the Data Flow

This method orchestrates the entire process, from input tokens to final prediction scores. The video below provides a fantastic walkthrough of a similar model's forward pass, highlighting how to handle the shapes and extract the correct output.

Synthesizing what you saw in the video and the SentimentRNN example, a forward pass for our classifier looks like this:

    def forward(self, x, hidden):
        batch_size = x.size(0)




        # 1. Pass through embedding layer
        # Input x shape: (batch_size, seq_length)
        # Output embeds shape: (batch_size, seq_length, embedding_dim)
        embeds = self.embedding(x)




        # 2. Pass through LSTM layer
        # lstm_out shape: (batch_size, seq_length, hidden_dim)
        # hidden is a tuple of the final hidden and cell states
        lstm_out, hidden = self.lstm(embeds, hidden)
        



        # 3. Stack up LSTM outputs and apply dropout
        # We only need the output of the last time step
        # lstm_out shape becomes: (batch_size, hidden_dim)
        lstm_out = lstm_out[:, -1, :]
        out = self.dropout(lstm_out)
        



        # 4. Pass through the fully-connected layer
        # out shape: (batch_size, output_dim)
        out = self.fc(out)
        



        # 5. Pass through the sigmoid activation function
        # sig_out shape: (batch_size, 1) for binary classification
        sig_out = self.sig(out)
        
        return sig_out, hidden

The most critical step here is lstm_out = lstm_out[:, -1, :]. This selects the hidden state outputs from the very last time step for all sequences in the batch, which is exactly what we need for our "many-to-one" classification.

A Complete Example

To see all these pieces work together in an end-to-end text classification context, let's look at one more practical example.

The video effectively shows how to instantiate the model and highlights a key practical aspect: the nn.Embedding layer can account for millions of parameters, making vocabulary size an important consideration.

Conclusion

Congratulations! You have now learned how to define a complete LSTM-based neural network for sequence classification in PyTorch. You've assembled the building blocks from our previous lessons into a coherent and functional model architecture.

Key Takeaways:

• PyTorch models are defined as classes inheriting from nn.Module.
• The __init__ method is where you declare and instantiate all the layers your model will use (e.g., nn.Embedding, nn.LSTM, nn.Linear).
• The forward method dictates the flow of data through these layers, from input to output.
• The nn.LSTM layer requires an initial hidden state and cell state, typically provided as a tuple of zero-tensors.
• For many-to-one sequence classification, you typically use the LSTM's output from the final time step to feed into the classification layer.

Preview of the Next Lesson

We've built the car, but we haven't started the engine yet. In our next lesson, we will do just that. We will take the LSTMClassifier class we designed today, write a complete training loop, and train it on a real sentiment analysis dataset to see it learn and make predictions.