RNNs and LSTMs in PyTorch

Real-World Scenario

A fintech company needs to classify transaction descriptions into categories (groceries, travel, entertainment) using short text sequences. An LSTM reads the tokens in order, carrying context forward so that “Apple Store” can be read differently from “apple”. The same machinery forecasts time series: read the recent past one step at a time, then predict the next value.

LSTM Inputs and Outputs

Most RNN bugs are shape bugs, so start with what goes in and what comes out.

import torch
import torch.nn as nn

torch.manual_seed(0)
lstm = nn.LSTM(input_size=10, hidden_size=64, num_layers=2, batch_first=True, dropout=0.2)
lstm.eval()                                   # dropout between layers only applies in training

x = torch.randn(32, 20, 10)                   # (batch, seq_len, features)
output, (h_n, c_n) = lstm(x)

print("output:", tuple(output.shape))         # every time step, last layer only
print("h_n:   ", tuple(h_n.shape))            # last time step, every layer
print("c_n:   ", tuple(c_n.shape))
print(torch.allclose(output[:, -1, :], h_n[-1]))

$ python lstm_shapes.py
output: (32, 20, 64)
h_n:    (2, 32, 64)
c_n:    (2, 32, 64)
True

With batch_first=True, input is (batch, seq_len, features). The LSTM returns two different summaries: output is the top layer’s hidden state at every one of the 20 steps, while h_n and c_n hold the hidden and cell state after the last step for each of the 2 layers. The final True confirms that the last step of output and the top layer of h_n are the same numbers. Without batch_first=True, PyTorch expects (seq_len, batch, features), and a (32, 20, 10) batch would be read as 32 time steps of 20 sequences — with no error.

RNN, GRU, and LSTM Side by Side

import torch
import torch.nn as nn

for cls in (nn.RNN, nn.GRU, nn.LSTM):
    layer = cls(input_size=10, hidden_size=64, batch_first=True)
    n = sum(p.numel() for p in layer.parameters())
    print(f"{cls.__name__:4s} {n:>6,} parameters  weight_ih {tuple(layer.weight_ih_l0.shape)}")

gru = nn.GRU(10, 64, num_layers=2, batch_first=True, bidirectional=True)
output, h_n = gru(torch.randn(32, 20, 10))    # no cell state
print("GRU output:", tuple(output.shape), " h_n:", tuple(h_n.shape))

$ python param_counts.py
RNN   4,864 parameters  weight_ih (64, 10)
GRU  14,592 parameters  weight_ih (192, 10)
LSTM 19,456 parameters  weight_ih (256, 10)
GRU output: (32, 20, 128)  h_n: (4, 32, 64)

The input weight matrix has 64 rows for a plain RNN, 192 for a GRU, and 256 for an LSTM: the GRU computes three gate-sized blocks at once and the LSTM four, so their parameter counts are exactly 3× and 4× the RNN’s. A GRU returns only (output, h_n) — it has no cell state. With bidirectional=True the output width doubles to 128 (forward and backward concatenated) and h_n has 4 entries: 2 layers × 2 directions.

Padding Changes the Final Hidden State

Batches need equal-length sequences, so shorter ones get padded with zeros. The LSTM doesn’t know those zeros are padding.

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

torch.manual_seed(0)
lstm = nn.LSTM(input_size=4, hidden_size=3, batch_first=True)

real = torch.randn(1, 2, 4)                        # a sequence of length 2
padded = torch.cat([real, torch.zeros(1, 3, 4)], dim=1)   # padded to length 5

_, (h_real, _) = lstm(real)
out_padded, (h_padded, _) = lstm(padded)
packed = pack_padded_sequence(padded, torch.tensor([2]), batch_first=True)
_, (h_packed, _) = lstm(packed)

print("real sequence:   ", [round(v, 4) for v in h_real[-1].squeeze().tolist()])
print("padded, h_n:     ", [round(v, 4) for v in h_padded[-1].squeeze().tolist()])
print("padded, step 2:  ", [round(v, 4) for v in out_padded[0, 1].tolist()])
print("packed, h_n:     ", [round(v, 4) for v in h_packed[-1].squeeze().tolist()])

$ python padding.py
real sequence:    [-0.0866, -0.0265, 0.006]
padded, h_n:      [-0.0287, 0.026, -0.0285]
padded, step 2:   [-0.0866, -0.0265, 0.006]
packed, h_n:      [-0.0866, -0.0265, 0.006]

Run on its own, the length-2 sequence ends in state [-0.0866, -0.0265, 0.006]. Padded to length 5, the LSTM keeps stepping through three rows of zeros and h_n drifts to a different vector. The correct state is still in output at step 2, and pack_padded_sequence gets it into h_n directly by telling the LSTM each sequence’s true length. Without packing, a classifier built on h_n reads the padding, not the sentence.

Text Classification with an LSTM

First, turn text into integer ids. The vocabulary is built from the training sentences; anything else maps to <unk>. Saved as vocab.py:

import re
from collections import Counter

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

def tokenize(text):
    return re.sub(r"[^\w\s]", "", text.lower()).split()

counts = Counter(t for text in TEXTS for t in tokenize(text))
vocab = {"<pad>": 0, "<unk>": 1} | {w: i + 2 for i, (w, _) in enumerate(counts.most_common())}

def encode(text, max_len=10):
    ids = [vocab.get(t, 1) for t in tokenize(text)][:max_len]
    return ids + [0] * (max_len - len(ids))

if __name__ == "__main__":
    print(len(vocab), "tokens in vocabulary")
    print(encode("the acting was fantastic"))
    print(encode("a truly dreadful sequel"))

$ python vocab.py
44 tokens in vocabulary
[4, 17, 7, 9, 0, 0, 0, 0, 0, 0]
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0]

Frequent words get small ids, 0 is reserved for padding and 1 for unknown words. Every word of "a truly dreadful sequel" is missing from the eight training sentences, so it encodes as four 1s — the model will see no actual words.

import torch
import torch.nn as nn
from vocab import TEXTS, LABELS, vocab, encode

class LSTMClassifier(nn.Module):
    def __init__(self, vocab_size, embed_dim=32, hidden=32):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(embed_dim, hidden, batch_first=True, bidirectional=True)
        self.fc = nn.Linear(hidden * 2, 2)

    def forward(self, ids):
        _, (h_n, _) = self.lstm(self.embedding(ids))
        return self.fc(torch.cat([h_n[-2], h_n[-1]], dim=1))   # forward + backward

torch.manual_seed(0)
X = torch.tensor([encode(t) for t in TEXTS])
y = torch.tensor(LABELS)
model = LSTMClassifier(len(vocab))
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)

for epoch in range(1, 31):
    optimizer.zero_grad()
    loss = nn.functional.cross_entropy(model(X), y)
    loss.backward()
    optimizer.step()
    if epoch % 10 == 0:
        print(f"epoch {epoch}  loss {loss.item():.4f}")

model.eval()
for text in ["brilliant and amazing", "awful waste of time", "a truly dreadful sequel"]:
    probs = model(torch.tensor([encode(text)])).softmax(1)[0]
    print(f"{text!r:26} positive={probs[1].item():.2f}")

$ python sentiment.py
epoch 10  loss 0.0033
epoch 20  loss 0.0000
epoch 30  loss 0.0000
'brilliant and amazing'    positive=1.00
'awful waste of time'      positive=0.00
'a truly dreadful sequel'  positive=0.00

nn.Embedding turns each id into a learned 32-dimensional vector, with padding_idx=0 keeping the padding vector at zero. The classifier concatenates the final forward and backward hidden states, h_n[-2] and h_n[-1]. With only eight sentences the loss rounds to zero by epoch 20 — the model has memorized its training set, and phrases built from training words score confidently. So does "a truly dreadful sequel", although the model saw nothing but <unk> tokens: a confident score is not evidence the model understood the input. Real text classification needs thousands of examples, a held-out validation set, and packed sequences as above.

Time-Series Forecasting

Predict the next value of a noisy sine wave from the previous 30. First the data, saved as series_data.py:

import numpy as np
import torch

rng = np.random.default_rng(0)
t = np.linspace(0, 100, 2000)
series = np.sin(t) + rng.normal(0, 0.1, len(t))

LOOKBACK = 30
X = np.stack([series[i:i + LOOKBACK] for i in range(len(series) - LOOKBACK)])
X = torch.tensor(X, dtype=torch.float32).unsqueeze(-1)          # (N, 30, 1)
y = torch.tensor(series[LOOKBACK:], dtype=torch.float32)
split = int(len(X) * 0.8)                                       # time order: no shuffle

Each example is a 30-step window, and its target is the value right after it. A forecast only means something next to a baseline, so the script also scores the naive “next value equals the last one” prediction:

import torch
import torch.nn as nn
from series_data import X, y, split

class Forecaster(nn.Module):
    def __init__(self, rnn_cls=nn.LSTM):
        super().__init__()
        self.rnn = rnn_cls(1, 32, batch_first=True)
        self.fc = nn.Linear(32, 1)
    def forward(self, x):
        return self.fc(self.rnn(x)[0][:, -1]).squeeze(-1)      # last step's output

def train_and_score(rnn_cls):
    torch.manual_seed(0)
    model = Forecaster(rnn_cls)
    opt = torch.optim.Adam(model.parameters(), lr=1e-2)
    for epoch in range(200):
        opt.zero_grad()
        nn.functional.mse_loss(model(X[:split]), y[:split]).backward()
        opt.step()
    with torch.no_grad():
        return nn.functional.mse_loss(model(X[split:]), y[split:]).item()

if __name__ == "__main__":
    print(f"LSTM test MSE:  {train_and_score(nn.LSTM):.4f}")
    naive = nn.functional.mse_loss(X[split:, -1, 0], y[split:]).item()   # "same as last step"
    print(f"naive test MSE: {naive:.4f}")

$ python forecast.py
LSTM test MSE:  0.0139
naive test MSE: 0.0243

The split keeps time order, so the model is tested on the future of the series it trained on — shuffling before splitting would leak the future into training. The LSTM’s error is about 40% lower than the naive forecast’s. It cannot get much lower: the added noise has standard deviation 0.1, so even a perfect predictor of the underlying sine wave scores about 0.01 — the third practice task measures it.

Practice

import torch
import torch.nn as nn

torch.manual_seed(0)
lstm = nn.LSTM(input_size=10, hidden_size=64, batch_first=True, bidirectional=True)
output, (h_n, _) = lstm(torch.randn(8, 20, 10))

print(tuple(output.shape), tuple(h_n.shape))
print(torch.allclose(output[:, -1, :64], h_n[0]))   # forward direction ends at the last step
print(torch.allclose(output[:, 0, 64:], h_n[1]))    # backward direction ends at step 0
print(torch.allclose(output[:, -1, 64:], h_n[1]))

$ python p1_bidirectional.py
(8, 20, 128) (2, 8, 64)
True
True
False

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

torch.manual_seed(0)
lstm = nn.LSTM(input_size=4, hidden_size=3, batch_first=True)
batch = torch.randn(3, 6, 4)
lengths = torch.tensor([4, 6, 2])                 # not sorted: enforce_sorted=False

packed = pack_padded_sequence(batch, lengths, batch_first=True, enforce_sorted=False)
packed_out, (h_n, _) = lstm(packed)
out, out_lengths = pad_packed_sequence(packed_out, batch_first=True)

print(out_lengths.tolist())
print(out[0, 4:].abs().sum().item())              # positions past the end are zeros
last = out[torch.arange(3), lengths - 1]          # each sequence's real last step
print(torch.allclose(last, h_n[-1]))

$ python p2_packed_batch.py
[4, 6, 2]
0.0
True

import numpy as np

rng = np.random.default_rng(0)
t = np.linspace(0, 100, 2000)
clean = np.sin(t)
series = clean + rng.normal(0, 0.1, len(t))

split = int((len(t) - 30) * 0.8) + 30              # first test target, as in series_data.py
floor = np.mean((clean[split:] - series[split:]) ** 2)
print(f"perfect-knowledge MSE: {floor:.4f}")

$ python p3_noise_floor.py
perfect-knowledge MSE: 0.0111

import torch.nn as nn
from forecast import train_and_score

print(f"GRU test MSE:   {train_and_score(nn.GRU):.4f}")
print(f"LSTM test MSE:  {train_and_score(nn.LSTM):.4f}")

$ python p4_gru.py
GRU test MSE:   0.0128
LSTM test MSE:  0.0139