Deep Learning Models: A collection of various deep learning architectures, models, & their implementations (PyTorch or TensorFlow) in Jupyter Notebooks.
- Author: Chizoba Obasi
- Github Repo: https://github.com/chizkidd/deep-learning-models
In [1]:
!pip install watermark
Collecting watermark Downloading watermark-2.6.0-py3-none-any.whl.metadata (15 kB) Requirement already satisfied: ipython>=6.0 in /usr/local/lib/python3.12/dist-packages (from watermark) (7.34.0) Requirement already satisfied: importlib-metadata>=1.4 in /usr/local/lib/python3.12/dist-packages (from watermark) (8.7.1) Requirement already satisfied: setuptools in /usr/local/lib/python3.12/dist-packages (from watermark) (75.2.0) Requirement already satisfied: zipp>=3.20 in /usr/local/lib/python3.12/dist-packages (from importlib-metadata>=1.4->watermark) (3.23.0) Requirement already satisfied: jedi>=0.16 in /usr/local/lib/python3.12/dist-packages (from ipython>=6.0->watermark) (0.19.2) Requirement already satisfied: decorator in /usr/local/lib/python3.12/dist-packages (from ipython>=6.0->watermark) (4.4.2) Requirement already satisfied: pickleshare in /usr/local/lib/python3.12/dist-packages (from ipython>=6.0->watermark) (0.7.5) Requirement already satisfied: traitlets>=4.2 in /usr/local/lib/python3.12/dist-packages (from ipython>=6.0->watermark) (5.7.1) Requirement already satisfied: prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0 in /usr/local/lib/python3.12/dist-packages (from ipython>=6.0->watermark) (3.0.52) Requirement already satisfied: pygments in /usr/local/lib/python3.12/dist-packages (from ipython>=6.0->watermark) (2.19.2) Requirement already satisfied: backcall in /usr/local/lib/python3.12/dist-packages (from ipython>=6.0->watermark) (0.2.0) Requirement already satisfied: matplotlib-inline in /usr/local/lib/python3.12/dist-packages (from ipython>=6.0->watermark) (0.2.1) Requirement already satisfied: pexpect>4.3 in /usr/local/lib/python3.12/dist-packages (from ipython>=6.0->watermark) (4.9.0) Requirement already satisfied: parso<0.9.0,>=0.8.4 in /usr/local/lib/python3.12/dist-packages (from jedi>=0.16->ipython>=6.0->watermark) (0.8.5) Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.12/dist-packages (from pexpect>4.3->ipython>=6.0->watermark) (0.7.0) Requirement already satisfied: wcwidth in /usr/local/lib/python3.12/dist-packages (from prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0->ipython>=6.0->watermark) (0.6.0) Downloading watermark-2.6.0-py3-none-any.whl (13 kB) Installing collected packages: watermark Successfully installed watermark-2.6.0
In [2]:
!pip install torchinfo datasets
Requirement already satisfied: torchinfo in /usr/local/lib/python3.12/dist-packages (1.8.0) Requirement already satisfied: datasets in /usr/local/lib/python3.12/dist-packages (4.0.0) Requirement already satisfied: filelock in /usr/local/lib/python3.12/dist-packages (from datasets) (3.20.3) Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.12/dist-packages (from datasets) (2.0.2) Requirement already satisfied: pyarrow>=15.0.0 in /usr/local/lib/python3.12/dist-packages (from datasets) (18.1.0) Requirement already satisfied: dill<0.3.9,>=0.3.0 in /usr/local/lib/python3.12/dist-packages (from datasets) (0.3.8) Requirement already satisfied: pandas in /usr/local/lib/python3.12/dist-packages (from datasets) (2.3.3) Requirement already satisfied: requests>=2.32.2 in /usr/local/lib/python3.12/dist-packages (from datasets) (2.32.4) Requirement already satisfied: tqdm>=4.66.3 in /usr/local/lib/python3.12/dist-packages (from datasets) (4.67.1) Requirement already satisfied: xxhash in /usr/local/lib/python3.12/dist-packages (from datasets) (3.6.0) Requirement already satisfied: multiprocess<0.70.17 in /usr/local/lib/python3.12/dist-packages (from datasets) (0.70.16) Requirement already satisfied: fsspec<=2025.3.0,>=2023.1.0 in /usr/local/lib/python3.12/dist-packages (from fsspec[http]<=2025.3.0,>=2023.1.0->datasets) (2025.3.0) Requirement already satisfied: huggingface-hub>=0.24.0 in /usr/local/lib/python3.12/dist-packages (from datasets) (1.4.1) Requirement already satisfied: packaging in /usr/local/lib/python3.12/dist-packages (from datasets) (25.0) Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.12/dist-packages (from datasets) (6.0.3) Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /usr/local/lib/python3.12/dist-packages (from fsspec[http]<=2025.3.0,>=2023.1.0->datasets) (3.13.3) Requirement already satisfied: hf-xet<2.0.0,>=1.2.0 in /usr/local/lib/python3.12/dist-packages (from huggingface-hub>=0.24.0->datasets) (1.2.0) Requirement already satisfied: httpx<1,>=0.23.0 in /usr/local/lib/python3.12/dist-packages (from huggingface-hub>=0.24.0->datasets) (0.28.1) Requirement already satisfied: shellingham in /usr/local/lib/python3.12/dist-packages (from huggingface-hub>=0.24.0->datasets) (1.5.4) Requirement already satisfied: typer-slim in /usr/local/lib/python3.12/dist-packages (from huggingface-hub>=0.24.0->datasets) (0.21.1) Requirement already satisfied: typing-extensions>=4.1.0 in /usr/local/lib/python3.12/dist-packages (from huggingface-hub>=0.24.0->datasets) (4.15.0) Requirement already satisfied: charset_normalizer<4,>=2 in /usr/local/lib/python3.12/dist-packages (from requests>=2.32.2->datasets) (3.4.4) Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.12/dist-packages (from requests>=2.32.2->datasets) (3.11) Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.12/dist-packages (from requests>=2.32.2->datasets) (2.5.0) Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.12/dist-packages (from requests>=2.32.2->datasets) (2026.1.4) Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.12/dist-packages (from pandas->datasets) (2.9.0.post0) Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.12/dist-packages (from pandas->datasets) (2025.2) Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.12/dist-packages (from pandas->datasets) (2025.3) Requirement already satisfied: aiohappyeyeballs>=2.5.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets) (2.6.1) Requirement already satisfied: aiosignal>=1.4.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets) (1.4.0) Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets) (25.4.0) Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets) (1.8.0) Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets) (6.7.0) Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets) (0.4.1) Requirement already satisfied: yarl<2.0,>=1.17.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2025.3.0,>=2023.1.0->datasets) (1.22.0) Requirement already satisfied: anyio in /usr/local/lib/python3.12/dist-packages (from httpx<1,>=0.23.0->huggingface-hub>=0.24.0->datasets) (4.12.1) Requirement already satisfied: httpcore==1.* in /usr/local/lib/python3.12/dist-packages (from httpx<1,>=0.23.0->huggingface-hub>=0.24.0->datasets) (1.0.9) Requirement already satisfied: h11>=0.16 in /usr/local/lib/python3.12/dist-packages (from httpcore==1.*->httpx<1,>=0.23.0->huggingface-hub>=0.24.0->datasets) (0.16.0) Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.12/dist-packages (from python-dateutil>=2.8.2->pandas->datasets) (1.17.0) Requirement already satisfied: click>=8.0.0 in /usr/local/lib/python3.12/dist-packages (from typer-slim->huggingface-hub>=0.24.0->datasets) (8.3.1)
In [3]:
%load_ext watermark
# %reload_ext watermark
%watermark -a 'Chizoba Obasi' -v -p torch
Author: Chizoba Obasi Python implementation: CPython Python version : 3.12.12 IPython version : 7.34.0 torch: 2.9.0+cu126
Long Short-Term Memory (LSTM) Classifier on IMDB Reviews¶
The network in this notebook is an implementation of a Long Short-Term Memory (LSTM) network [1] on IMDB reviews [7]. The model used is a Single-Layer LSTM implemented from scratch. It utilizes a gated recurrent unit structure to maintain a long-term cell state, distinguishing it from a standard Vanilla RNN.
Network Architecture:
The LSTM evolved from the simple RNN [2] to solve the Vanishing Gradient problem. While a standard RNN has a single hidden state, an LSTM introduces a Cell State ($c_t$)—a "long-term memory" track—and Gating Mechanisms that regulate what information is added, removed, or passed forward.
The repeating module in a standard RNN contains a single layer. (Source, [5])
The repeating module in an LSTM contains four interacting layers. (Source, [5])
Key Concepts of the LSTM Architecture:¶
The Cell State ($c_t$): Often called the "constant error carousel," the cell state runs through the entire sequence with only minor linear interactions. This allows gradients to flow over long distances without vanishing, enabling the model to remember context from the beginning of a long news article.
Gating Mechanisms: LSTMs use three specialized gates to control information flow:
- Forget Gate ($f_t$): Decides which parts of the past memory are no longer relevant.
- Input Gate ($i_t$): Determines which new information from the current word should be stored in the cell state.
- Output Gate ($o_t$): Filters the cell state to produce the hidden state ($h_t$) used for the current prediction.
Dual State Recurrence: At each time step $t$, the LSTM maintains both a Hidden State ($h_t$) (short-term memory/output) and a Cell State ($c_t$) (long-term memory). This dual-track system is what makes LSTMs significantly more powerful than the "from-scratch" simple RNN.
Improvements of LSTM over Simple RNN:¶
Solving the Vanishing Gradient: In a simple RNN, the repeated multiplication of weights ($W_{hh}$) causes gradients to shrink exponentially as they travel back in time. LSTMs use additive updates for the cell state, which protects the gradient and allows the model to learn dependencies across 100+ tokens.
Handling Long-Range Dependencies: A simple RNN often forgets the first few words of a sentence by the time it reaches the tenth. An LSTM can "lock" important information (like the subject of a sentence) into the cell state and keep it there until it is needed much later.
Selective Filtering: Because of the gates, LSTMs are better at ignoring "noise" (like stop words or irrelevant punctuation) and focusing only on the "signal" (key verbs or nouns) that determines the news category.
Quick Comparison Table: RNN vs. LSTM Architecture¶
| Feature | Simple RNN | Long Short-Term Memory (LSTM) |
|---|---|---|
| Memory Track | Single (Hidden State $h_t$) | Dual (Hidden State $h_t$ + Cell State $c_t$) |
| Gradient Flow | Multiplication (Vanishing Gradients) | Addition (Stable Gradient Flow) |
| Dependency Range | Short-term (Approx. 10–20 steps) | Long-term (Approx. 100+ steps) |
| Gating Control | None (Static updates) | 3 Gates (Forget, Input, Output) |
| Complexity | Low (Fastest to train) | High (4x more parameters than RNN) |
| Use Case | Simple patterns / Toy tasks | Complex NLP (Translation, Topic Modeling) |
| Information Flow | Aggressive (Overwrites memory) | Selective (Remembers/Forgets strategically) |
REFERENCES:
- Hochreiter, Sepp, and Jürgen Schmidhuber. "Long Short-Term Memory." Neural Computation, 1997.
- Elman, Jeffrey L. "Finding Structure in Time." Cognitive Science, 1990.
- PyTorch. "Recurrent Layers: nn.LSTM." PyTorch Documentation, 2026.
- Zhang, Aston, et al. "Long Short-Term Memory (Chapter 10.1)." Dive into Deep Learning. Cambridge University Press, 2023.
- Olah, Christopher. "Understanding LSTM Networks." colah's blog, 2015.
- Raschka, Sebastian. "L15: Introduction to recurrent neural networks." Introduction to Deep Learning, 2021. [slides]
- Maas, Andrew L., et al. "Learning Word Vectors for Sentiment Analysis." Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, 2011.
- Lhoest, Quentin, et al. "Datasets: A Community Library for Natural Language Processing." Hugging Face Datasets, 2021.
- Google. (2026). "Gemini 3 Flash" [Large language model]. Google AI.
In [4]:
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader
from torch.nn.utils.rnn import pad_sequence
from datasets import load_dataset, DatasetDict
from collections import Counter
# from torch.cuda.amp import GradScaler, autocast
from torch.amp import GradScaler, autocast
from torchinfo import summary
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
%matplotlib inline
import pandas as pd
import numpy as np
import math
import time
import os
import random
import re
if torch.cuda.is_available():
torch.backends.cudnn.deterministic = True
In [5]:
def set_all_seeds(seed):
os.environ["PL_GLOBAL_SEED"] = str(seed)
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
In [6]:
##########################
### SETTINGS
##########################
# Hyperparameters
RANDOM_SEED = 123
BATCH_SIZE = 128
LEARNING_RATE = 0.001
NUM_EPOCHS = 20
NUM_CLASSES = 1
VOCABULARY_SIZE = 20000
EMBEDDING_DIM = 128
HIDDEN_DIM = 256
# OUTPUT_DIM = 1
# Set device
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(DEVICE)
cuda
In [7]:
# 1. Load Dataset (IMDB)
raw_datasets = load_dataset("imdb")
# 2. Split 10% of 'train' to create a 'validation' set (seed for reproducibility)
train_val_split = raw_datasets["train"].train_test_split(test_size=0.1, seed=RANDOM_SEED)
ds_splits = DatasetDict({
'train': train_val_split['train'],
'validation': train_val_split['test'], # The 10% we just carved out
'test': raw_datasets['test']
})
print(ds_splits)
# 2. Build Vocabulary (From Scratch Logic)
tokenizer = lambda text: re.findall(r'\w+', text.lower())
word_counts = Counter()
# Important: We only iterate over ds_splits['train']
for text in ds_splits['train']['text']:
word_counts.update(tokenizer(text))
# Filter rare words and create mapping: : <PAD>=0, <UNK>=1, then top 20k words
vocab = {word: i + 2 for i, (word, count) in enumerate(word_counts.most_common(VOCABULARY_SIZE))}
vocab['<PAD>'] = 0
vocab['<UNK>'] = 1
print(f"Vocab size: {len(vocab)}")
# 3. Custom Dataset Class
class TextDataset(torch.utils.data.Dataset):
def __init__(self, dataset_split):
self.data = dataset_split
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
# Return raw text and label. Let collate_fn handle the rest.
item = self.data[idx]
return item['text'], item['label']
# 4. Robust Text Pipeline
def text_pipeline(text, max_len=128):
tokens = tokenizer(text)
# Use .get(t, 1) because 1 is your <UNK> index
return [vocab.get(t, 1) for t in tokens][:max_len]
# 5. Robust Collate Batch
def collate_batch(batch):
label_list, text_list, lengths = [], [], []
for _text, _label in batch:
# 1. Labels: Use LONG for AG News (Multi-class)
label_list.append(_label)
# 2. Numericalize and Truncate
processed_text = torch.tensor(text_pipeline(_text), dtype=torch.int64)
text_list.append(processed_text)
lengths.append(processed_text.size(0))
# 3. Padding logic
# We use <PAD> index 0. We use vocab.get just in case names change.
pad_idx = vocab.get('<PAD>', 0)
# This dynamically pads to the longest sequence in THIS batch (more efficient)
text_list_padded = pad_sequence(text_list, batch_first=True, padding_value=pad_idx)
label_tensor = torch.tensor(label_list, dtype=torch.long)
length_tensor = torch.tensor(lengths, dtype=torch.int64)
return text_list_padded.to(DEVICE), label_tensor.to(DEVICE), length_tensor.to(DEVICE)
# 6. Initialize DataLoaders
train_ds = TextDataset(ds_splits['train'])
val_ds = TextDataset(ds_splits['validation'])
test_ds = TextDataset(ds_splits['test'])
train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_batch)
val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, collate_fn=collate_batch)
test_loader = DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=False, collate_fn=collate_batch)
Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
DatasetDict({
train: Dataset({
features: ['text', 'label'],
num_rows: 22500
})
validation: Dataset({
features: ['text', 'label'],
num_rows: 2500
})
test: Dataset({
features: ['text', 'label'],
num_rows: 25000
})
})
Vocab size: 20002
How do we figure out a good max_len? see code below
def get_robust_maxlen(dataset, tokenizer, p=95, absolute_max=512):
# 1. Get the statistical length
lengths = [len(tokenizer(ex[0])) for ex in dataset]
p_length = np.percentile(lengths, p)
# 2. Find the next power of 2
# math.log2 gives the exponent, ceil rounds up, 2** raises it back
# We use max(16, ...) to ensure we don't get a tiny length like 4 or 8
robust_len = 2 ** math.ceil(math.log2(max(16, p_length)))
# 3. Cap at absolute_max
final_len = min(robust_len, absolute_max)
print(f"{p}th Percentile: {p_length:.2f}")
print(f"Nearest Power of 2: {robust_len}")
print(f"Final Robust MAX_LEN: {final_len}")
return final_len
# Usage
__MAX_LEN = get_robust_maxlen(train_ds, tokenizer, p=95)
2. Model Creation¶
2.1. LSTM Unit Cell¶
LSTMCellFromScratch: The Mathematical Breakdown¶
The goal of the LSTM cell is to maintain two states: the Cell State ($C_t$) (long-term memory) and the Hidden State ($h_t$) (short-term output).
1. The Linear Bottleneck¶
self.gates = nn.Linear(input_size + hidden_size, 4 * hidden_size)
combined = torch.cat((x, h_prev), dim=1)
gates = self.gates(combined)
- What it is: Instead of 4 separate layers, we use one large matrix.
- The Math: $z = W \cdot [x_t, h_{t-1}] + b$
- Dimensions: $W$ has a shape of $[4H, I+H]$. It contains the weights for all four gates stacked vertically.
2. The Gating Logic (Unpacking)¶
We use .chunk(4) to split the linear output into the four specific control signals.
| Code Variable | Math Symbol | Activation | Logical Purpose |
|---|---|---|---|
f_gate |
$f_t$ | Sigmoid | Forget: What percentage of old memory should we delete? (0 = Delete, 1 = Keep) |
i_gate |
$i_t$ | Sigmoid | Input: How much of the new word should we let into the memory? |
g_gate |
$\tilde{C}_t$ | Tanh | Candidate: The actual "content" of the new word, scaled between -1 and 1. |
o_gate |
$o_t$ | Sigmoid | Output: How much of our internal memory should we show to the next layer? |
3. The Memory Update (The "Highway")¶
This is where the LSTM solves the vanishing gradient problem.
c_curr = (f_gate * c_prev) + (i_gate * g_gate)
- The Math: $C_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t$
- Why it works: Because we use addition ($+$) to update the cell state, the gradient can travel backward through the "memory highway" without being multiplied by a small number and disappearing.
4. The Hidden State (The "Filter")¶
The hidden state is just a filtered version of the long-term memory.
h_curr = o_gate * torch.tanh(c_curr)
- The Math: $h_t = o_t \odot \tanh(C_t)$
- Why it works: We squash the memory ($C_t$) into a -1 to 1 range and then use the Output Gate to decide which parts of that memory are relevant for the current prediction (e.g., predicting "Business" vs "Sports").
Summary of Dimensions¶
If input_size = 128 and hidden_size = 256:
- Input ($x$): $[Batch, 128]$
- Hidden ($h$): $[Batch, 256]$
- Weight Matrix ($W$): $[1024, 384]$ (Note: $1024 = 4 \times 256$ and $384 = 128 + 256$)
- Bias ($b$): $[1024]$
2.2. The Reason for Gate Combination leveraging nn.Linear¶
In the context of the LSTMCellFromScratch, self.gates is an nn.Linear layer. In PyTorch, an nn.Linear layer performs the operation $y = xA^T + b$.
When we use it for an LSTM, we are essentially "packing" four separate sub-layers into one big matrix to make the computation faster.
What is $W$ (the Weight)?¶
In the line self.gates = nn.Linear(input_size + hidden_size, 4 * hidden_size), the weight matrix $W$ (referred to as self.gates.weight) has the dimensions:
$$(4 \times \text{hidden\_size}, \text{input\_size} + \text{hidden\_size})$$
This matrix is actually a vertical stack of four smaller weight matrices, one for each "gate" ($i, f, g, o$):
- $W_i$: Input gate weights
- $W_f$: Forget gate weights
- $W_g$: Candidate gate weights (often called $W_c$)
- $W_o$: Output gate weights
What is $b$ (the Bias)?¶
The bias vector $b$ (self.gates.bias) has the dimension:
$$(4 \times \text{hidden\_size})$$
Just like the weights, it is a concatenation of four bias vectors: $[b_i, b_f, b_g, b_o]$.
How the Math matches the Code¶
When you call gates = self.gates(combined), PyTorch does this under the hood:
$$\text{gates} = \begin{bmatrix} x_t \\ h_{t-1} \end{bmatrix} \cdot \begin{bmatrix} W_i & W_f & W_g & W_o \\ \end{bmatrix}^T + \begin{bmatrix} b_i \\ b_f \\ b_g \\ b_o \end{bmatrix}$$
After this operation, the variable gates contains a long vector. The .chunk(4, dim=1) line in your code is what "unpacks" this result into the specific values for each gate:
# The 'chunk' operation splits the result of (Wx + b)
f_gate, i_gate, g_gate, o_gate = gates.chunk(4, dim=1)
| Gate Name | Code Variable | Math Activation | Equation |
|---|---|---|---|
| Forget | f_gate |
Sigmoid ($\sigma$) | $f_t = \sigma(W_f \cdot [x_t, h_{t-1}] + b_f)$ |
| Input | i_gate |
Sigmoid ($\sigma$) | $i_t = \sigma(W_i \cdot [x_t, h_{t-1}] + b_i)$ |
| Candidate | g_gate |
Tanh | $\tilde{C}_t = \tanh(W_g \cdot [x_t, h_{t-1}] + b_g)$ |
| Output | o_gate |
Sigmoid ($\sigma$) | $o_t = \sigma(W_o \cdot [x_t, h_{t-1}] + b_o)$ |
A Note on "Robust" Initialization¶
In professional implementations, $b_f$ (the forget gate bias) is often initialized to 1.0 instead of 0. This "robust" tweak ensures that at the start of training, the model defaults to "remembering" everything and only learns to forget over time, which prevents the vanishing gradient problem from the very first epoch.
Video explainer on gate combination for LSTM (Linear Layers)¶
This video provides a deep dive into how linear layers in PyTorch handle multi-dimensional weights and biases, which is essential for understanding how the LSTM's four gates are calculated simultaneously.
In [9]:
class LSTMCellFromScratch(nn.Module):
def __init__(self, input_size, hidden_size):
super().__init__()
self.hidden_size = hidden_size
# We need weights for all 4 internal operations: Forget, Input, Candidate Cell, and Output gates
# Each gate takes (input + previous_hidden)
self.gates = nn.Linear(input_size + hidden_size, 4 * hidden_size)
def forward(self, x, states):
# states is a tuple: (h_prev, c_prev)
h_prev, c_prev = states
# Combined linear pass for efficiency
combined = torch.cat((x, h_prev), dim=1)
gates = self.gates(combined)
# Split the gates
f_gate, i_gate, g_gate, o_gate = gates.chunk(4, dim=1)
# Apply activations
f_gate = torch.sigmoid(f_gate) # Forget
i_gate = torch.sigmoid(i_gate) # Input
g_gate = torch.tanh(g_gate) # New candidate memory
o_gate = torch.sigmoid(o_gate) # Output
# The Magic: Update Cell State (Long-term memory)
c_curr = (f_gate * c_prev) + (i_gate * g_gate)
# Update Hidden State (Short-term/Output)
h_curr = o_gate * torch.tanh(c_curr)
return h_curr, c_curr
In [10]:
class LSTMFromScratch(nn.Module):
def __init__(self, vocab_size, embed_dim, hidden_size, output_size):
super().__init__()
self.hidden_size = hidden_size
self.embedding = nn.Embedding(vocab_size, embed_dim)
self.cell = LSTMCellFromScratch(embed_dim, hidden_size)
self.fc = nn.Linear(hidden_size, output_size)
def forward(self, x, lengths):
batch_size, seq_len = x.shape
x = self.embedding(x) # [batch, seq, embed]
# Initialize h_0 and c_0 as zeros
h = torch.zeros(batch_size, self.hidden_size).to(x.device)
c = torch.zeros(batch_size, self.hidden_size).to(x.device)
# To handle padding robustly, we store ALL hidden states
all_h = []
for t in range(seq_len):
h, c = self.cell(x[:, t, :], (h, c))
all_h.append(h.unsqueeze(1)) # [Batch, 1, Hidden]
# Stack them: [Batch, Seq_Len, Hidden]
all_h = torch.cat(all_h, dim=1)
# ROBUST TRICK: Select the hidden state at the ACTUAL end of each sentence
# We subtract 1 because lengths are 1-indexed, but indices are 0-indexed
idx = (lengths - 1).view(-1, 1, 1).expand(batch_size, 1, self.hidden_size) # [batch, 1, hidden]
last_h = all_h.gather(1, idx).squeeze(1) #[batch, hidden]
return self.fc(last_h)
In [11]:
VOCAB_SIZE = len(vocab)
In [12]:
# --- Configuration Switch ---
DATASET_NAME = "IMDB"
if DATASET_NAME == "IMDB":
IS_BINARY = True
NUM_CLASSES = 1
# criterion = nn.BCEWithLogitsLoss()
else:
IS_BINARY = False
NUM_CLASSES = 4 # AG News has 4 classes
# criterion = nn.CrossEntropyLoss()
In [13]:
set_all_seeds(RANDOM_SEED)
model = LSTMFromScratch(
vocab_size=VOCAB_SIZE,
embed_dim=EMBEDDING_DIM,
hidden_size=HIDDEN_DIM,
output_size=NUM_CLASSES
).to(DEVICE)
summary(model,
input_data=(torch.zeros((BATCH_SIZE, EMBEDDING_DIM), dtype=torch.long).to(DEVICE),
torch.ones((BATCH_SIZE,), dtype=torch.long).to(DEVICE) * EMBEDDING_DIM),
col_names=["input_size", "output_size", "num_params", "mult_adds"],
device=DEVICE)
Out[13]:
============================================================================================================================================ Layer (type:depth-idx) Input Shape Output Shape Param # Mult-Adds ============================================================================================================================================ LSTMFromScratch [128, 128] [128, 1] -- -- ├─Embedding: 1-1 [128, 128] [128, 128, 128] 2,560,256 327,712,768 ├─LSTMCellFromScratch: 1-2 [128, 128] [128, 256] -- -- │ └─Linear: 2-1 [128, 384] [128, 1024] 394,240 50,462,720 ├─LSTMCellFromScratch: 1-3 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-2 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-4 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-3 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-5 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-4 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-6 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-5 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-7 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-6 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-8 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-7 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-9 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-8 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-10 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-9 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-11 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-10 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-12 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-11 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-13 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-12 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-14 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-13 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-15 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-14 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-16 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-15 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-17 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-16 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-18 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-17 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-19 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-18 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-20 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-19 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-21 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-20 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-22 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-21 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-23 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-22 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-24 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-23 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-25 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-24 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-26 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-25 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-27 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-26 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-28 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-27 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-29 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-28 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-30 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-29 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-31 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-30 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-32 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-31 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-33 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-32 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-34 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-33 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-35 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-34 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-36 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-35 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-37 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-36 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-38 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-37 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-39 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-38 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-40 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-39 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-41 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-40 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-42 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-41 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-43 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-42 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-44 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-43 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-45 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-44 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-46 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-45 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-47 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-46 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-48 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-47 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-49 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-48 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-50 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-49 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-51 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-50 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-52 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-51 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-53 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-52 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-54 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-53 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-55 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-54 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-56 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-55 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-57 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-56 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-58 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-57 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-59 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-58 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-60 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-59 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-61 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-60 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-62 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-61 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-63 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-62 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-64 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-63 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-65 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-64 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-66 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-65 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-67 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-66 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-68 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-67 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-69 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-68 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-70 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-69 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-71 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-70 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-72 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-71 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-73 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-72 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-74 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-73 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-75 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-74 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-76 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-75 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-77 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-76 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-78 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-77 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-79 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-78 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-80 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-79 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-81 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-80 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-82 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-81 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-83 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-82 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-84 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-83 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-85 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-84 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-86 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-85 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-87 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-86 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-88 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-87 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-89 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-88 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-90 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-89 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-91 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-90 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-92 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-91 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-93 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-92 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-94 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-93 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-95 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-94 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-96 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-95 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-97 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-96 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-98 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-97 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-99 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-98 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-100 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-99 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-101 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-100 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-102 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-101 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-103 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-102 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-104 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-103 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-105 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-104 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-106 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-105 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-107 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-106 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-108 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-107 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-109 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-108 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-110 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-109 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-111 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-110 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-112 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-111 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-113 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-112 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-114 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-113 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-115 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-114 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-116 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-115 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-117 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-116 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-118 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-117 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-119 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-118 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-120 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-119 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-121 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-120 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-122 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-121 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-123 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-122 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-124 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-123 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-125 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-124 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-126 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-125 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-127 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-126 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-128 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-127 [128, 384] [128, 1024] (recursive) 50,462,720 ├─LSTMCellFromScratch: 1-129 [128, 128] [128, 256] (recursive) -- │ └─Linear: 2-128 [128, 384] [128, 1024] (recursive) 50,462,720 ├─Linear: 1-130 [128, 256] [128, 1] 257 32,896 ============================================================================================================================================ Total params: 2,954,753 Trainable params: 2,954,753 Non-trainable params: 0 Total mult-adds (Units.GIGABYTES): 6.79 ============================================================================================================================================ Input size (MB): 0.13 Forward/backward pass size (MB): 151.00 Params size (MB): 11.82 Estimated Total Size (MB): 162.95 ============================================================================================================================================
In [14]:
def calculate_accuracy(loader, model, device, is_binary=True):
model.eval()
correct, total = 0, 0
with torch.no_grad():
for inputs, labels, lengths in loader:
inputs, labels, lengths = inputs.to(device), labels.to(device), lengths.to(device)
logits = model(inputs, lengths)
if is_binary:
predicted = (torch.sigmoid(logits.view(-1)) > 0.5).long()
else:
_, predicted = torch.max(logits.data, 1)
total += labels.size(0)
correct += (predicted == labels.long()).sum().item()
return 100 * correct / total
In [15]:
IS_BINARY, NUM_CLASSES
Out[15]:
(True, 1)
In [16]:
# 1. Update Optimizer and Scheduler
# Binary uses BCEWithLogitsLoss; Multi-class uses CrossEntropy
criterion = nn.BCEWithLogitsLoss() if IS_BINARY else nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
# Robust Scheduler: Cosine Annealing
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=NUM_EPOCHS)
In [17]:
start_time = time.time()
history = {'train_loss': [], 'val_loss': [], 'train_acc': [], 'val_acc': [], 'lrs': []}
print(f"Starting training on {DEVICE}...")
for epoch in range(NUM_EPOCHS):
model.train()
running_loss = 0.0
# 1. Training Phase
for i, (inputs, labels, lengths) in enumerate(train_loader):
inputs, labels, lengths = inputs.to(DEVICE), labels.to(DEVICE), lengths.to(DEVICE)
optimizer.zero_grad()
# Forward pass (passing lengths for masking logic)
logits = model(inputs, lengths)
if IS_BINARY:
loss = criterion(logits.squeeze(), labels.float())
else:
loss = criterion(logits, labels.long())
loss.backward()
# --- ROBUST RNN FIX: Gradient Clipping ---
# Prevents gradients from "exploding" during the backprop through time (BPTT)
nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
running_loss += loss.item() * inputs.size(0)
# 2. Metrics Calculation
# Note: Using the updated calculate_accuracy that handles (inputs, lengths)
train_loss = running_loss / len(train_ds)
train_acc = calculate_accuracy(train_loader, model, DEVICE, is_binary=IS_BINARY)
# 3. Validation Phase
model.eval()
val_loss = 0.0
with torch.no_grad():
for inputs, labels, lengths in val_loader:
inputs, labels, lengths = inputs.to(DEVICE), labels.to(DEVICE), lengths.to(DEVICE)
logits = model(inputs, lengths)
if IS_BINARY:
loss = criterion(logits.squeeze(), labels.float())
else:
loss = criterion(logits, labels.long())
val_loss += loss.item() * inputs.size(0)
val_loss /= len(val_ds)
val_acc = calculate_accuracy(val_loader, model, DEVICE, is_binary=IS_BINARY)
scheduler.step()
current_lr = optimizer.param_groups[0]['lr']
# 4. Logging
history['train_loss'].append(train_loss)
history['val_loss'].append(val_loss)
history['train_acc'].append(train_acc)
history['val_acc'].append(val_acc)
history['lrs'].append(current_lr)
print(f'Epoch {epoch + 1}/{NUM_EPOCHS} | '
f'Loss (train/val): {train_loss:.4f} / {val_loss:.4f} | '
f'Acc (train/val): {train_acc:.2f}% / {val_acc:.2f}% | '
f'LR: {current_lr:.6f}')
print('\nTotal Training Time: %.2f min' % ((time.time() - start_time)/60))
Starting training on cuda... Epoch 1/20 | Loss (train/val): 0.6594 / 0.6341 | Acc (train/val): 68.66% / 65.72% | LR: 0.000994 Epoch 2/20 | Loss (train/val): 0.5893 / 0.5818 | Acc (train/val): 76.12% / 72.64% | LR: 0.000976 Epoch 3/20 | Loss (train/val): 0.5271 / 0.6038 | Acc (train/val): 73.51% / 68.24% | LR: 0.000946 Epoch 4/20 | Loss (train/val): 0.4578 / 0.5171 | Acc (train/val): 83.43% / 78.16% | LR: 0.000905 Epoch 5/20 | Loss (train/val): 0.3701 / 0.4854 | Acc (train/val): 86.46% / 76.72% | LR: 0.000854 Epoch 6/20 | Loss (train/val): 0.2992 / 0.4453 | Acc (train/val): 91.72% / 81.52% | LR: 0.000794 Epoch 7/20 | Loss (train/val): 0.2326 / 0.4921 | Acc (train/val): 92.50% / 80.64% | LR: 0.000727 Epoch 8/20 | Loss (train/val): 0.1755 / 0.4936 | Acc (train/val): 96.26% / 81.16% | LR: 0.000655 Epoch 9/20 | Loss (train/val): 0.1205 / 0.5896 | Acc (train/val): 97.31% / 81.96% | LR: 0.000578 Epoch 10/20 | Loss (train/val): 0.0729 / 0.6241 | Acc (train/val): 99.13% / 81.88% | LR: 0.000500 Epoch 11/20 | Loss (train/val): 0.0390 / 0.7155 | Acc (train/val): 99.71% / 80.68% | LR: 0.000422 Epoch 12/20 | Loss (train/val): 0.0197 / 0.8559 | Acc (train/val): 99.92% / 81.24% | LR: 0.000345 Epoch 13/20 | Loss (train/val): 0.0108 / 0.9757 | Acc (train/val): 99.98% / 81.72% | LR: 0.000273 Epoch 14/20 | Loss (train/val): 0.0047 / 1.0680 | Acc (train/val): 100.00% / 81.40% | LR: 0.000206 Epoch 15/20 | Loss (train/val): 0.0028 / 1.1675 | Acc (train/val): 100.00% / 82.00% | LR: 0.000146 Epoch 16/20 | Loss (train/val): 0.0019 / 1.2085 | Acc (train/val): 100.00% / 81.64% | LR: 0.000095 Epoch 17/20 | Loss (train/val): 0.0015 / 1.2613 | Acc (train/val): 100.00% / 81.80% | LR: 0.000054 Epoch 18/20 | Loss (train/val): 0.0013 / 1.2911 | Acc (train/val): 100.00% / 81.76% | LR: 0.000024 Epoch 19/20 | Loss (train/val): 0.0012 / 1.3057 | Acc (train/val): 100.00% / 81.84% | LR: 0.000006 Epoch 20/20 | Loss (train/val): 0.0012 / 1.3068 | Acc (train/val): 100.00% / 81.76% | LR: 0.000000 Total Training Time: 9.41 min
In [18]:
# --- Final Evaluation ---
test_accuracy = calculate_accuracy(test_loader, model, DEVICE, is_binary=IS_BINARY)
print(f'\nFinal Test Accuracy: {test_accuracy:.2f}%')
Final Test Accuracy: 80.10%
In [19]:
def plot_history(history):
"""
Plots the training and validation loss, accuracy, and learning rate
from a PyTorch history dictionary.
Args:
history (dict): Dictionary containing training metrics.
Expected keys: 'train_loss', 'val_loss', 'train_acc', 'val_acc', 'lrs'.
"""
epochs = range(len(history['train_loss']))
plt.figure(figsize=(18, 5))
# Plot 1: Loss
plt.subplot(1, 3, 1)
plt.plot(epochs, history['train_loss'], label='Training Loss')
plt.plot(epochs, history['val_loss'], label='Validation Loss')
plt.title('Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
# plt.grid(True)
# Plot 2: Accuracy
plt.subplot(1, 3, 2)
plt.plot(epochs, history['train_acc'], label='Training Accuracy')
plt.plot(epochs, history['val_acc'], label='Validation Accuracy')
plt.title('Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()
# plt.grid(True)
# Plot 3: Learning Rate
plt.subplot(1, 3, 3)
plt.plot(epochs, history['lrs'], label='Learning Rate', color='orange')
plt.title('Learning Rate')
plt.xlabel('Epochs')
plt.ylabel('Learning Rate')
plt.legend()
# plt.grid(True)
plt.tight_layout()
plt.show()
# plt.savefig()
plot_history(history)
In [23]:
# 1. Reuse your exact training tokenizer
tokenizer = lambda text: re.findall(r'\w+', text.lower())
def predict_sentiment(model, sentence):
model.eval()
# Step A: Tokenize exactly like training
tokens = tokenizer(sentence)
# Step B: Numericalize using the <UNK> index (1)
# We use .get(t, 1) to handle words not in your top 20k vocab
indexed = [vocab.get(t, 1) for t in tokens]
# Step C: Prepare Tensors [Batch Size 1, Seq_Len]
tensor = torch.LongTensor(indexed).unsqueeze(0).to(DEVICE)
length = torch.LongTensor([len(indexed)]).to(DEVICE)
with torch.no_grad():
# Step D: Forward Pass
logits = model(tensor, length)
# Step E: Convert Logit to Probability (0.0 to 1.0)
# Using .view(-1) for safety with batch size 1
prob = torch.sigmoid(logits.view(-1)).item()
return prob
# --- Test ---
review = "This film was an absolute masterpiece of modern cinema."
score = predict_sentiment(model, review)
print(f'Review: "{review}"')
print(f'Positive Probability: {score:.4f}')
print(f'Final Sentiment: {"Positive" if score > 0.5 else "Negative"}')
Review: "This film was an absolute masterpiece of modern cinema." Positive Probability: 0.9287 Final Sentiment: Positive
In [28]:
imdb_tests = [
"I wanted to like this movie, but the acting was wooden and the plot was a mess.", # Negative
"An absolute triumph of cinema that I would recommend to everyone I know.", # Positive
"Not the worst thing I've ever seen, but definitely not the best either.", # Neutral/Leaning Neg
"The special effects were stunning, though the story felt a bit rushed.", # Mixed
"I fell asleep halfway through; avoid this waste of time at all costs." # Highly Negative
]
print(f"{'Review':<60} | {'Prob':<6} | {'Sentiment'}")
print("-" * 80)
for text in imdb_tests:
prob = predict_sentiment(model, text)
sentiment = "Positive" if prob > 0.5 else "Negative"
print(f"{text[:58]:<60} | {prob:.3f} | {sentiment}")
Review | Prob | Sentiment -------------------------------------------------------------------------------- I wanted to like this movie, but the acting was wooden and | 0.000 | Negative An absolute triumph of cinema that I would recommend to ev | 0.988 | Positive Not the worst thing I've ever seen, but definitely not the | 0.009 | Negative The special effects were stunning, though the story felt a | 1.000 | Positive I fell asleep halfway through; avoid this waste of time at | 0.000 | Negative
In [29]:
predict_sentiment(model, "This film is quite interesting")
Out[29]:
0.8147006630897522
In [30]:
predict_sentiment(model, "This film is terrible")
Out[30]:
0.02001946046948433
In [31]:
%watermark -iv
datasets : 4.0.0 matplotlib : 3.10.0 numpy : 2.0.2 pandas : 2.3.3 re : 2.2.1 sklearn : 1.6.1 torch : 2.9.0+cu126 torchinfo : 1.8.0 torchvision: 0.24.0+cu126