Open In Colab

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]:
%%capture
!pip install watermark
!pip install torchinfo datasets
In [2]:
%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.10.0+cu128


Bidirectional Gated Recurrent Unit (Bi-GRU) Classifier on AG News¶

The network in this notebook is an implementation of a Bidirectional Gated Recurrent Unit (Bi-GRU) network [1] on AG news dataset [8, 9]. Building a Bi-GRU from scratch is the definitive way to prove understanding of temporal context. In a standard GRU, the model is "historically limited"—it only knows what came before. In a Bi-GRU, the model becomes "omnipresent" within the sequence, seeing both the past and the future for every token.


Network Architecture:

A Bi-GRU consists of two independent GRU layers (or cells in our from-scratch case). The Forward Layer captures information from preceding words, while the Backward Layer captures information from succeeding words. This is particularly effective for text classification (like AG News), where the global context of a sentence determines its category.


RNN vs LSTM vs GRU

RNN vs LSTM vs GRU. (Source)




GRU vs Bi-GRU

GRU vs Bi-GRU (Source)





Key Concepts of the Bi-GRU Architecture:¶

  1. Dual-Directional Flow: The model maintains two separate hidden states ($h_f$ and $h_b$). For a sentence like "The apple was red," the forward pass sees "The apple..." while the backward pass sees "...was red."
  2. Concatenation: At each time step, the hidden states from both directions are concatenated: $h_t = [h_{f,t} ; h_{b,t}]$. This doubles the feature space available to the final Linear layer.
  3. Information Density: Because the model sees the entire sequence from both ends, it is significantly better at resolving ambiguities in language (e.g., distinguishing "Bank" as a financial institution vs. a river bank based on surrounding words).



Quick Comparison Table: GRU vs. Bidirectional GRU¶

Feature Standard GRU (Unidirectional) Bidirectional GRU (Bi-GRU)
Data Processing Left-to-Right ($0 \rightarrow T$) Both Directions ($0 \leftrightarrow T$)
Hidden States Single ($h_t$) Dual ($h_{forward}$ and $h_{backward}$)
Context Past only Past and Future
Output Dim hidden_size hidden_size * 2
Parameters Base ($3 \times W$) $2 \times$ Base (Two separate sets of weights)
Latency Faster Slower (Must complete both passes)
Best For Real-time / Causal Prediction Sentence Classification / Translation




REFERENCES:

  1. Schuster, Mike, and Kuldip K. Paliwal. "Bidirectional recurrent neural networks." IEEE Transactions on Signal Processing, 1997. (Original Bidirectional Paper)
  2. Cho, Kyunghyun, et al. "Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation." arXiv preprint arXiv:1406.1078, 2014. (Original GRU Paper)
  3. Bahdanau, Dzmitry, Kyunghyun Cho, and Yoshua Bengio. "Neural Machine Translation by Jointly Learning to Align and Translate." arXiv preprint arXiv:1409.0473, 2014. (Bi-GRU Implementation)
  4. Zhang, Xiang, Junbo Zhao, and Yann LeCun. "Character-level Convolutional Networks for Text Classification." Advances in Neural Information Processing Systems (NeurIPS), 2015.
  5. Gulli, Antonio. "AG's Corpus of News Articles." 2004.
  6. Google. (2026). "Gemini 3 Flash" [Large language model]. Google AI

In [3]:
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 [4]:
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)

0. Model Hyperparameters¶


In [5]:
##########################
### 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

1. Data Preparation & Loading¶


In [6]:
# 1. Load Dataset (IMDB)
raw_datasets = load_dataset("ag_news")


# 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)
DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 108000
    })
    validation: Dataset({
        features: ['text', 'label'],
        num_rows: 12000
    })
    test: Dataset({
        features: ['text', 'label'],
        num_rows: 7600
    })
})
Vocab size: 20002

2. Model Creation¶


In [7]:
class GRUCellFromScratch(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.hidden_size = hidden_size

        # We need weights for 3 operations:
        # Update gate (z), Reset gate (r), and Candidate hidden state (h_tilde)
        self.gate_weights = nn.Linear(input_size + hidden_size, 2 * hidden_size)
        self.candidate_weights = nn.Linear(input_size + hidden_size, hidden_size)

    def forward(self, x, h_prev):
        # x: [batch, input_size]
        # h_prev: [batch, hidden_size]

        # 1. Calculate Update (z) and Reset (r) gates
        combined = torch.cat((x, h_prev), dim=1)
        gates = torch.sigmoid(self.gate_weights(combined))
        z_gate, r_gate = gates.chunk(2, dim=1)

        # 2. Calculate Candidate Hidden State (h_tilde)
        # We use the Reset gate to "mask" the previous hidden state
        combined_candidate = torch.cat((x, r_gate * h_prev), dim=1)
        h_tilde = torch.tanh(self.candidate_weights(combined_candidate))

        # 3. Final Hidden State update
        # This is a linear interpolation between past and candidate
        h_curr = (1 - z_gate) * h_prev + z_gate * h_tilde

        return h_curr
In [8]:
class BiGRUFromScratch(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.gru_cell_forward = GRUCellFromScratch(embed_dim, hidden_size)
        self.gru_cell_backward = GRUCellFromScratch(embed_dim, hidden_size)

        # The FC layer MUST be twice as big because we concatenate the states
        self.fc = nn.Linear(hidden_size * 2, output_size)

    def forward(self, x, lengths=None):
        # 1. Robust check for torchinfo/bundled inputs
        if isinstance(x, (list, tuple)) and lengths is None:
            x, lengths = x

        batch_size, seq_len = x.shape
        x = self.embedding(x) # [Batch, Seq, Embed]

        # 2. Initialize both hidden states
        h_f = torch.zeros(batch_size, self.hidden_size).to(x.device)
        h_b = torch.zeros(batch_size, self.hidden_size).to(x.device)

        all_h_f = []
        all_h_b = []

        # 3. FORWARD PASS (Left to Right)
        for t in range(seq_len):
            h_f = self.gru_cell_forward(x[:, t, :], h_f)
            all_h_f.append(h_f.unsqueeze(1))

        # 4. BACKWARD PASS (Right to Left)
        for t in reversed(range(seq_len)):
            h_b = self.gru_cell_backward(x[:, t, :], h_b)
            all_h_b.append(h_b.unsqueeze(1))

        # Reverse backward states back to chronological order for concatenation
        all_h_b = all_h_b[::-1]

        # 5. Concatenate across the hidden dimension: [Batch, Seq, Hidden * 2]
        all_h_f = torch.cat(all_h_f, dim=1)
        all_h_b = torch.cat(all_h_b, dim=1)

        # 6. ROBUST INDEXING
        # For Forward: We want the state at the actual end of the sentence (lengths-1)
        idx_f = (lengths - 1).view(-1, 1, 1).expand(batch_size, 1, self.hidden_size)
        last_h_f = all_h_f.gather(1, idx_f).squeeze(1)

        # For Backward: The 'last' information is at index 0
        # (because it read from end to start, so word 0 is its final summary)
        last_h_b = all_h_b[:, 0, :]

        # 7. Final Concatenation and Prediction
        combined_out = torch.cat([last_h_f, last_h_b], dim=1) # [Batch, Hidden * 2]

        return self.fc(combined_out)
In [9]:
def get_percentile_length(dataset, tokenizer, p=95):
    # Unpack the tuple directly
    lengths = []
    for text, label in dataset:
        tokens = tokenizer(text)
        lengths.append(len(tokens))

    return int(np.percentile(lengths, p))

_MAX_LEN = get_percentile_length(train_ds, tokenizer, p=95)
print(f"95% of sequences are shorter than {_MAX_LEN} tokens.")
95% of sequences are shorter than 58 tokens.
In [10]:
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)
95th Percentile: 58.00
Nearest Power of 2: 64
Final Robust MAX_LEN: 64
In [11]:
VOCAB_SIZE = len(vocab)
MAX_LEN = 64
In [12]:
# --- Configuration Switch ---
DATASET_NAME = "ag-news"

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 = BiGRUFromScratch(
    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, MAX_LEN), dtype=torch.long).to(DEVICE),
                    torch.ones((BATCH_SIZE,), dtype=torch.long).to(DEVICE) * MAX_LEN),
        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
============================================================================================================================================
BiGRUFromScratch                         [128, 64]                 [128, 4]                  --                        --
├─Embedding: 1-1                         [128, 64]                 [128, 64, 128]            2,560,256                 327,712,768
├─GRUCellFromScratch: 1-2                [128, 128]                [128, 256]                --                        --
│    └─Linear: 2-1                       [128, 384]                [128, 512]                197,120                   25,231,360
│    └─Linear: 2-2                       [128, 384]                [128, 256]                98,560                    12,615,680
├─GRUCellFromScratch: 1-3                [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-3                       [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-4                       [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-4                [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-5                       [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-6                       [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-5                [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-7                       [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-8                       [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-6                [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-9                       [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-10                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-7                [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-11                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-12                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-8                [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-13                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-14                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-9                [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-15                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-16                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-10               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-17                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-18                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-11               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-19                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-20                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-12               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-21                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-22                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-13               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-23                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-24                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-14               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-25                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-26                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-15               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-27                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-28                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-16               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-29                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-30                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-17               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-31                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-32                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-18               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-33                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-34                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-19               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-35                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-36                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-20               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-37                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-38                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-21               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-39                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-40                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-22               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-41                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-42                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-23               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-43                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-44                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-24               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-45                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-46                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-25               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-47                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-48                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-26               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-49                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-50                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-27               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-51                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-52                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-28               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-53                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-54                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-29               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-55                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-56                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-30               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-57                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-58                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-31               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-59                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-60                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-32               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-61                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-62                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-33               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-63                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-64                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-34               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-65                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-66                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-35               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-67                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-68                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-36               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-69                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-70                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-37               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-71                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-72                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-38               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-73                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-74                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-39               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-75                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-76                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-40               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-77                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-78                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-41               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-79                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-80                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-42               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-81                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-82                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-43               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-83                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-84                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-44               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-85                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-86                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-45               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-87                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-88                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-46               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-89                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-90                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-47               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-91                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-92                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-48               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-93                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-94                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-49               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-95                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-96                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-50               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-97                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-98                      [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-51               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-99                      [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-100                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-52               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-101                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-102                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-53               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-103                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-104                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-54               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-105                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-106                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-55               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-107                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-108                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-56               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-109                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-110                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-57               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-111                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-112                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-58               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-113                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-114                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-59               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-115                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-116                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-60               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-117                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-118                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-61               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-119                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-120                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-62               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-121                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-122                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-63               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-123                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-124                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-64               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-125                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-126                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-65               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-127                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-128                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-66               [128, 128]                [128, 256]                --                        --
│    └─Linear: 2-129                     [128, 384]                [128, 512]                197,120                   25,231,360
│    └─Linear: 2-130                     [128, 384]                [128, 256]                98,560                    12,615,680
├─GRUCellFromScratch: 1-67               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-131                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-132                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-68               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-133                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-134                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-69               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-135                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-136                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-70               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-137                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-138                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-71               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-139                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-140                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-72               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-141                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-142                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-73               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-143                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-144                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-74               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-145                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-146                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-75               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-147                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-148                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-76               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-149                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-150                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-77               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-151                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-152                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-78               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-153                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-154                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-79               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-155                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-156                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-80               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-157                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-158                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-81               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-159                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-160                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-82               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-161                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-162                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-83               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-163                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-164                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-84               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-165                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-166                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-85               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-167                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-168                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-86               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-169                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-170                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-87               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-171                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-172                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-88               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-173                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-174                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-89               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-175                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-176                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-90               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-177                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-178                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-91               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-179                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-180                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-92               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-181                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-182                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-93               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-183                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-184                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-94               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-185                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-186                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-95               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-187                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-188                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-96               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-189                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-190                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-97               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-191                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-192                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-98               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-193                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-194                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-99               [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-195                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-196                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-100              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-197                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-198                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-101              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-199                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-200                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-102              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-201                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-202                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-103              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-203                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-204                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-104              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-205                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-206                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-105              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-207                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-208                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-106              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-209                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-210                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-107              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-211                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-212                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-108              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-213                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-214                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-109              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-215                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-216                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-110              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-217                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-218                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-111              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-219                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-220                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-112              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-221                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-222                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-113              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-223                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-224                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-114              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-225                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-226                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-115              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-227                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-228                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-116              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-229                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-230                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-117              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-231                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-232                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-118              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-233                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-234                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-119              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-235                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-236                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-120              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-237                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-238                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-121              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-239                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-240                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-122              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-241                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-242                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-123              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-243                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-244                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-124              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-245                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-246                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-125              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-247                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-248                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-126              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-249                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-250                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-127              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-251                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-252                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-128              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-253                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-254                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─GRUCellFromScratch: 1-129              [128, 128]                [128, 256]                (recursive)               --
│    └─Linear: 2-255                     [128, 384]                [128, 512]                (recursive)               25,231,360
│    └─Linear: 2-256                     [128, 384]                [128, 256]                (recursive)               12,615,680
├─Linear: 1-130                          [128, 512]                [128, 4]                  2,052                     262,656
============================================================================================================================================
Total params: 3,153,668
Trainable params: 3,153,668
Non-trainable params: 0
Total mult-adds (Units.GIGABYTES): 5.17
============================================================================================================================================
Input size (MB): 0.07
Forward/backward pass size (MB): 109.06
Params size (MB): 12.61
Estimated Total Size (MB): 121.74
============================================================================================================================================

3. Model Training, Testing & Evaluation¶


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]:
(False, 4)
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.4773 / 0.3194 | Acc (train/val): 91.35% / 88.97% | LR: 0.000994
Epoch 2/20 | Loss (train/val): 0.2435 / 0.2735 | Acc (train/val): 94.70% / 90.91% | LR: 0.000976
Epoch 3/20 | Loss (train/val): 0.1660 / 0.2870 | Acc (train/val): 96.48% / 90.47% | LR: 0.000946
Epoch 4/20 | Loss (train/val): 0.1039 / 0.3061 | Acc (train/val): 98.56% / 91.03% | LR: 0.000905
Epoch 5/20 | Loss (train/val): 0.0548 / 0.3624 | Acc (train/val): 99.19% / 90.59% | LR: 0.000854
Epoch 6/20 | Loss (train/val): 0.0310 / 0.4508 | Acc (train/val): 99.52% / 90.75% | LR: 0.000794
Epoch 7/20 | Loss (train/val): 0.0200 / 0.4799 | Acc (train/val): 99.71% / 90.82% | LR: 0.000727
Epoch 8/20 | Loss (train/val): 0.0138 / 0.5174 | Acc (train/val): 99.78% / 91.13% | LR: 0.000655
Epoch 9/20 | Loss (train/val): 0.0104 / 0.5384 | Acc (train/val): 99.89% / 91.14% | LR: 0.000578
Epoch 10/20 | Loss (train/val): 0.0073 / 0.5522 | Acc (train/val): 99.86% / 91.09% | LR: 0.000500
Epoch 11/20 | Loss (train/val): 0.0067 / 0.5674 | Acc (train/val): 99.92% / 91.28% | LR: 0.000422
Epoch 12/20 | Loss (train/val): 0.0040 / 0.5918 | Acc (train/val): 99.95% / 91.22% | LR: 0.000345
Epoch 13/20 | Loss (train/val): 0.0028 / 0.5816 | Acc (train/val): 99.95% / 91.05% | LR: 0.000273
Epoch 14/20 | Loss (train/val): 0.0025 / 0.5980 | Acc (train/val): 99.97% / 91.18% | LR: 0.000206
Epoch 15/20 | Loss (train/val): 0.0016 / 0.6273 | Acc (train/val): 99.97% / 91.27% | LR: 0.000146
Epoch 16/20 | Loss (train/val): 0.0013 / 0.6439 | Acc (train/val): 99.97% / 91.22% | LR: 0.000095
Epoch 17/20 | Loss (train/val): 0.0010 / 0.6569 | Acc (train/val): 99.97% / 91.14% | LR: 0.000054
Epoch 18/20 | Loss (train/val): 0.0008 / 0.6723 | Acc (train/val): 99.98% / 91.16% | LR: 0.000024
Epoch 19/20 | Loss (train/val): 0.0006 / 0.6755 | Acc (train/val): 99.97% / 91.17% | LR: 0.000006
Epoch 20/20 | Loss (train/val): 0.0005 / 0.6769 | Acc (train/val): 99.98% / 91.19% | LR: 0.000000

Total Training Time: 72.96 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: 91.22%
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)
No description has been provided for this image

4. Model Inference and Sentiment Prediction¶


In [21]:
# 1. Define the Human Labels
ag_news_labels = {0: "World", 1: "Sports", 2: "Business", 3: "Sci/Tech"}

# 2. Re-use your Training Tokenizer
# This ensures "Don't" becomes ["don", "t"] consistently in both phases
tokenizer = lambda text: re.findall(r'\w+', text.lower())


# Adapted the below code from:
# https://github.com/bentrevett/pytorch-sentiment-analysis/blob/experimental/2%20-%20Upgraded%20Sentiment%20Analysis.ipynb
def predict_topic(model, sentence):
    model.eval()

    # Step A: Tokenize exactly like training
    tokens = tokenizer(sentence)

    # Step B: Numericalize using the <UNK> index (1)
    indexed = [vocab.get(t, 1) for t in tokens]

    # Step C: Prepare Tensors [Batch, 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: Get the highest score
        # Using softmax to see the confidence %
        probs = torch.softmax(logits, dim=1)
        prediction_idx = torch.argmax(probs, dim=1).item()
        confidence = probs[0][prediction_idx].item()

    return ag_news_labels[prediction_idx], confidence

# --- Test ----------------------------------------
headline = "The Olympic games were postponed due to global health concerns."
topic, score = predict_topic(model, headline)
print(f"Result: {topic} ({score:.2%})")
Result: World (59.77%)
In [22]:
# Create a list of test cases
ag_news_tests = [
    "The Federal Reserve is expected to keep interest rates steady as inflation cools.", # Business
    "The star striker scored a hat-trick in the final minutes of the world cup qualifier.", # Sports
    "Astronomers have discovered a new Earth-like planet orbiting a distant star system.", # Sci/Tech
    "Protests erupted in the capital city following the controversial election results.",    # World
    "New software update fixes critical security vulnerabilities in mobile operating systems." # Sci/Tech
]

print(f"{'Headline':<60} | {'Predicted Topic':<12} | {'Confidence'}")
print("-" * 85)

for text in ag_news_tests:
    topic, conf = predict_topic(model, text)
    print(f"{text[:58]:<60} | {topic:<12} | {conf:.2%}")
Headline                                                     | Predicted Topic | Confidence
-------------------------------------------------------------------------------------
The Federal Reserve is expected to keep interest rates ste   | Business     | 100.00%
The star striker scored a hat-trick in the final minutes o   | Sports       | 100.00%
Astronomers have discovered a new Earth-like planet orbiti   | Sci/Tech     | 100.00%
Protests erupted in the capital city following the controv   | World        | 100.00%
New software update fixes critical security vulnerabilitie   | Sci/Tech     | 100.00%
In [23]:
%watermark -iv
datasets   : 4.0.0
matplotlib : 3.10.0
numpy      : 2.0.2
pandas     : 2.2.2
re         : 2.2.1
sklearn    : 1.6.1
torch      : 2.10.0+cu128
torchinfo  : 1.8.0
torchvision: 0.25.0+cu128