Skip to main content
PyTorch intermediate Lesson 6 of 11

RNNs and LSTMs in PyTorch

Build sequence models from scratch — RNNs, LSTMs, GRUs — and apply them to text classification and time-series forecasting.

Real-World Scenario

A fintech company needs to classify transaction descriptions into categories (groceries, travel, entertainment, etc.) using short text sequences. An LSTM processes each word token sequentially, capturing the context that “Apple Store” means electronics, not groceries. With 50,000 labeled transactions, the LSTM achieves 89% accuracy.

Building an LSTM from Scratch

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import numpy as np

# Understanding LSTM inputs and outputs
# Input shape:  (seq_len, batch_size, input_size)   — time-first by default
#               or (batch_size, seq_len, input_size) — with batch_first=True
# Output:       (output, (h_n, c_n))
#   output: (seq_len, batch, hidden) — all hidden states
#   h_n:   (num_layers, batch, hidden) — last hidden state
#   c_n:   (num_layers, batch, hidden) — last cell state

lstm = nn.LSTM(
    input_size=10,     # features per time step
    hidden_size=64,    # LSTM units
    num_layers=2,      # stacked LSTMs
    batch_first=True,  # (batch, seq, feature) instead of (seq, batch, feature)
    dropout=0.2,       # applied between layers (not on the last layer)
    bidirectional=False,
)

batch, seq_len, features = 32, 20, 10
x = torch.randn(batch, seq_len, features)

output, (h_n, c_n) = lstm(x)
print(f"output shape: {output.shape}")    # (32, 20, 64) — all time steps
print(f"h_n shape:    {h_n.shape}")       # (2, 32, 64)  — last hidden state, both layers
print(f"c_n shape:    {c_n.shape}")       # (2, 32, 64)  — last cell state
print(f"final step:   {output[:, -1, :].shape}")  # (32, 64) — last time step output

Text Classification with LSTM

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from collections import Counter
import re

# Toy sentiment dataset
TEXTS = [
    "the movie was absolutely fantastic and amazing",
    "terrible film waste of time and money",
    "brilliant acting and stunning visuals",
    "boring slow paced and poorly written",
    "loved every minute of this masterpiece",
    "awful direction and bad script",
    "one of the best films I have ever seen",
    "completely disappointed by this disaster",
]
LABELS = [1, 0, 1, 0, 1, 0, 1, 0]  # 1=positive, 0=negative

# Build vocabulary
def tokenize(text: str) -> list[str]:
    return re.sub(r'[^\w\s]', '', text.lower()).split()

all_tokens = [t for text in TEXTS for t in tokenize(text)]
vocab = {word: i + 2 for i, (word, _) in enumerate(Counter(all_tokens).most_common())}
vocab["<PAD>"] = 0
vocab["<UNK>"] = 1
VOCAB_SIZE = len(vocab)

def encode(text: str, max_len: int = 15) -> list[int]:
    tokens = tokenize(text)[:max_len]
    ids    = [vocab.get(t, 1) for t in tokens]
    return ids + [0] * (max_len - len(ids))  # pad to max_len


class LSTMClassifier(nn.Module):
    def __init__(self, vocab_size: int, embed_dim: int, hidden: int, n_classes: int):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(embed_dim, hidden, num_layers=2,
                            batch_first=True, dropout=0.3, bidirectional=True)
        self.dropout = nn.Dropout(0.3)
        # bidirectional doubles the hidden size
        self.classifier = nn.Linear(hidden * 2, n_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        emb = self.dropout(self.embedding(x))        # (B, L, E)
        out, (h_n, _) = self.lstm(emb)
        # Concatenate last hidden state from both directions
        h_fwd = h_n[-2]   # last layer, forward direction
        h_bck = h_n[-1]   # last layer, backward direction
        h = torch.cat([h_fwd, h_bck], dim=1)         # (B, H*2)
        return self.classifier(self.dropout(h))       # (B, n_classes)


# Prepare data
X = torch.tensor([encode(t) for t in TEXTS], dtype=torch.long)
y = torch.tensor(LABELS, dtype=torch.long)

model     = LSTMClassifier(VOCAB_SIZE, embed_dim=32, hidden=64, n_classes=2)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

loader = DataLoader(TensorDataset(X, y), batch_size=4, shuffle=True)

for epoch in range(20):
    model.train()
    total_loss = 0
    for X_b, y_b in loader:
        optimizer.zero_grad()
        loss = criterion(model(X_b), y_b)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        total_loss += loss.item()
    if (epoch + 1) % 5 == 0:
        print(f"Epoch {epoch+1:3d}  loss={total_loss/len(loader):.4f}")

Time-Series Forecasting with LSTM

import torch
import torch.nn as nn
import numpy as np

# Create synthetic time series: sin wave + noise
t = np.linspace(0, 100, 2000)
series = np.sin(t) + np.random.normal(0, 0.1, len(t))
series = (series - series.mean()) / series.std()

def create_sequences(data: np.ndarray, lookback: int) -> tuple[torch.Tensor, torch.Tensor]:
    X, y = [], []
    for i in range(len(data) - lookback):
        X.append(data[i:i + lookback])
        y.append(data[i + lookback])
    X = torch.tensor(np.array(X), dtype=torch.float32).unsqueeze(-1)  # (N, L, 1)
    y = torch.tensor(np.array(y), dtype=torch.float32)
    return X, y

LOOKBACK = 30
X, y = create_sequences(series, LOOKBACK)

split = int(len(X) * 0.8)
X_train, X_test = X[:split], X[split:]
y_train, y_test = y[:split], y[split:]


class TimeSeriesLSTM(nn.Module):
    def __init__(self, input_size=1, hidden=64, n_layers=2):
        super().__init__()
        self.lstm = nn.LSTM(input_size, hidden, n_layers,
                            batch_first=True, dropout=0.2)
        self.fc   = nn.Linear(hidden, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out, _ = self.lstm(x)
        return self.fc(out[:, -1, :]).squeeze()   # use last time step


model     = TimeSeriesLSTM()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.MSELoss()

from torch.utils.data import DataLoader, TensorDataset
loader = DataLoader(TensorDataset(X_train, y_train), batch_size=64, shuffle=True)

for epoch in range(30):
    model.train()
    for X_b, y_b in loader:
        optimizer.zero_grad()
        pred = model(X_b)
        loss = criterion(pred, y_b)
        loss.backward()
        optimizer.step()

model.eval()
with torch.no_grad():
    test_pred = model(X_test)
    test_mse  = criterion(test_pred, y_test)
    print(f"Test MSE: {test_mse.item():.4f}")

GRU: Simpler Alternative to LSTM

import torch
import torch.nn as nn

# GRU has no cell state — only hidden state
# Often matches LSTM performance with fewer parameters and faster training
gru = nn.GRU(
    input_size=10,
    hidden_size=64,
    num_layers=2,
    batch_first=True,
    dropout=0.2,
    bidirectional=True,
)

x = torch.randn(32, 20, 10)
output, h_n = gru(x)   # GRU returns (output, h_n) — no cell state

print(f"output: {output.shape}")  # (32, 20, 128) — bidirectional doubles hidden
print(f"h_n:    {h_n.shape}")     # (4, 32, 64)   — 2 layers × 2 directions

Packed Sequences for Variable-Length Input

import torch
import torch.nn as nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

# Without packing, the LSTM wastes computation on padding tokens
# With packing, it skips padding and processes only real tokens

class EfficientLSTMClassifier(nn.Module):
    def __init__(self, vocab_size: int, embed_dim: int, hidden: int, n_classes: int):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm      = nn.LSTM(embed_dim, hidden, batch_first=True)
        self.fc        = nn.Linear(hidden, n_classes)

    def forward(self, x: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor:
        emb    = self.embedding(x)                     # (B, L, E)
        packed = pack_padded_sequence(emb, lengths.cpu(),
                                      batch_first=True,
                                      enforce_sorted=False)
        out, (h_n, _) = self.lstm(packed)
        # h_n[-1] is the last hidden state at each sequence's actual end
        return self.fc(h_n[-1])


# Example usage
vocab_size = 1000
model = EfficientLSTMClassifier(vocab_size, 32, 64, 2)

# Batch of 3 sequences, padded to length 10
seqs = torch.tensor([
    [5, 12, 3, 7, 0, 0, 0, 0, 0, 0],
    [8, 2, 9, 4, 11, 6, 0, 0, 0, 0],
    [1, 15, 0, 0, 0, 0, 0, 0, 0, 0],
], dtype=torch.long)

lengths = torch.tensor([4, 6, 2])  # actual lengths without padding
logits  = model(seqs, lengths)
print(f"Logits shape: {logits.shape}")  # (3, 2)

Frequently Asked Questions

What problem do LSTMs solve that vanilla RNNs can't?
Vanilla RNNs suffer from vanishing gradients when backpropagating through long sequences — gradients shrink toward zero and early time steps stop learning. LSTMs introduce a cell state with gating mechanisms (forget, input, output gates) that allow gradients to flow unchanged over hundreds of steps. GRUs are a simpler alternative with similar performance.
When should I use an LSTM vs a Transformer?
LSTMs are still competitive for: (1) small datasets where Transformers overfit, (2) real-time streaming where you process tokens one at a time, (3) resource-constrained environments. Transformers are better when you have enough data, need attention over long contexts, or can afford parallelism during training.