code/11-recurrent-neural-networks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "📝 **Author:** Amirhossein Heydari - 📧 **Email:** <amirhosseinheydari78@gmail.com> - 📍 **Origin:** [mr-pylin/pytorch-workshop](https://github.com/mr-pylin/pytorch-workshop)\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Table of contents**<a id='toc0_'></a>    \n",
    "- [Dependencies](#toc1_)    \n",
    "- [Dataset](#toc2_)    \n",
    "  - [Regular Dataset](#toc2_1_)    \n",
    "  - [Sequential Dataset](#toc2_2_)    \n",
    "- [Types of sequence-to-sequence modeling configurations](#toc3_)    \n",
    "- [Network Structure: Recurrent Neural Networks](#toc4_)    \n",
    "  - [Simple Vanilla RNN](#toc4_1_)    \n",
    "  - [Combined Weights and Concatenated Input and Hidden](#toc4_2_)    \n",
    "  - [Deep RNN](#toc4_3_)    \n",
    "  - [RNN using PyTorch](#toc4_4_)    \n",
    "  - [Long Short-Term Memory (LSTM)](#toc4_5_)    \n",
    "  - [Gated Recurrent Units (GRU)](#toc4_6_)    \n",
    "\n",
    "<!-- vscode-jupyter-toc-config\n",
    "\tnumbering=false\n",
    "\tanchor=true\n",
    "\tflat=false\n",
    "\tminLevel=1\n",
    "\tmaxLevel=6\n",
    "\t/vscode-jupyter-toc-config -->\n",
    "<!-- THIS CELL WILL BE REPLACED ON TOC UPDATE. DO NOT WRITE YOUR TEXT IN THIS CELL -->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc1_'></a>[Dependencies](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader, Dataset, TensorDataset\n",
    "from torchinfo import summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set a seed for deterministic results\n",
    "seed = 42\n",
    "torch.manual_seed(seed)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check if cuda is available\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# log\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc2_'></a>[Dataset](#toc0_)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc2_1_'></a>[Regular Dataset](#toc0_)\n",
    "\n",
    "- Regular datasets typically used in CNNs & MLPs are composed of independent data points\n",
    "- Each data point is usually represented as a fixed-size vector (or tensor for images)\n",
    "\n",
    "🧾 **Notations**:\n",
    "\n",
    "- $N$: Number of samples in the dataset.\n",
    "- $\\mathbf{x}_i$: Input data point $i$, where $i \\in \\{1, 2, \\ldots, N\\}$.\n",
    "- $\\mathbf{y}_i$: Label or target associated with input data $i$.\n",
    "\n",
    "🔬 **Formulations**:\n",
    "\n",
    "- Dataset: $D=\\{(\\mathbf{x}_i, \\mathbf{y}_i)\\mid i = 1, 2, \\ldots, N\\}$\n",
    "- Each $\\mathbf{x}_i \\in ℝ^M$, where $M$ is the dimensionality of the input feature vector\n",
    "\n",
    "🌟 **Example**: $D = \\{ (\\mathbf{x}_1, \\mathbf{y}_1), (\\mathbf{x}_2, \\mathbf{y}_2), (\\mathbf{x}_3, \\mathbf{y}_3) \\}$\n",
    "\n",
    "- $\\mathbf{x}_1 = [1.0, 2.0], \\quad \\mathbf{y}_1 = 0$\n",
    "- $\\mathbf{x}_2 = [2.5, 3.5], \\quad \\mathbf{y}_2 = 1$\n",
    "- $\\mathbf{x}_3 = [0.5, 1.5], \\quad \\mathbf{y}_3 = 0$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RegularDataset(Dataset):\n",
    "    def __init__(self):\n",
    "        self.data = torch.tensor([[1.1, 2.1], [2.5, 3.5], [0.5, 1.5]], dtype=torch.float32)\n",
    "        self.labels = torch.tensor([0, 1, 0], dtype=torch.int64)\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:\n",
    "        return self.data[idx], self.labels[idx]\n",
    "\n",
    "\n",
    "# create dataset and dataloader\n",
    "dataset = RegularDataset()\n",
    "dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0)\n",
    "\n",
    "# iterate through the dataset\n",
    "for data, label in dataloader:\n",
    "    print(f\"data: {data}, label: {label}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc2_2_'></a>[Sequential Dataset](#toc0_)\n",
    "\n",
    "- Sequential datasets used in RNNs are composed of sequences of data points.\n",
    "- Each sequence represents a temporal or sequential relationship among the data points.\n",
    "\n",
    "🧾 **Notations**:\n",
    "\n",
    "- $N$: Number of sequences in the dataset.\n",
    "- $T$: Length of each sequence.\n",
    "- $\\mathbf{x}^t_i$: Input data point at time step $t$ in the sequence $i$, where $t \\in \\{1, 2, \\ldots, T\\}$ and $i \\in \\{1, 2, \\ldots, N\\}$\n",
    "- $\\mathbf{y}_i$: Label or target associated with sequence $i$.\n",
    "\n",
    "🔬 **Formulations**:\n",
    "\n",
    "- Dataset: $D = \\{ (\\mathbf{x}_i^1, \\mathbf{x}_i^2, \\ldots, \\mathbf{x}_i^T, \\mathbf{y}_i) \\mid i = 1, 2, \\ldots, N \\}$\n",
    "- Each $\\mathbf{x}^t_i \\in ℝ^M$, where $M$ is the dimensionality of the input feature vector at each time step.\n",
    "\n",
    "🌟 **Example**: $D = \\{ (\\mathbf{x}_1, \\mathbf{x}_2, \\mathbf{x}_3, \\mathbf{y}_1), (\\mathbf{x}_2, \\mathbf{x}_3, \\mathbf{x}_4, \\mathbf{y}_2), (\\mathbf{x}_3, \\mathbf{x}_4, \\mathbf{x}_5, \\mathbf{y}_3) \\}$\n",
    "\n",
    "- $\\mathbf{x}_1 = [1.0, 0.0]$\n",
    "- $\\mathbf{x}_2 = [0.5, 1.5]$\n",
    "- $\\mathbf{x}_3 = [1.0, 2.0]$\n",
    "- $\\mathbf{x}_4 = [2.0, 1.0]$\n",
    "- $\\mathbf{x}_5 = [1.5, 0.5]$\n",
    "- $\\mathbf{y}_1 = 0$\n",
    "- $\\mathbf{y}_2 = 1$\n",
    "- $\\mathbf{y}_3 = 0$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SequentialDatasetWithoutOverlap(Dataset):\n",
    "    def __init__(self):\n",
    "        # original data points\n",
    "        self.data = torch.tensor([[1.0, 0.0], [0.5, 1.5], [1.0, 2.0], [2.0, 1.0], [1.5, 0.5], [2.5, 1.5]], dtype=torch.float32)\n",
    "\n",
    "        # labels for each sequence\n",
    "        self.labels = torch.tensor([0, 1], dtype=torch.int64)\n",
    "\n",
    "        # sequence length\n",
    "        self.seq_length = 3\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        # number of sequences without overlap\n",
    "        return len(self.data) // self.seq_length\n",
    "\n",
    "    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:\n",
    "        # calculate the start index of the sequence\n",
    "        start_idx = idx * self.seq_length\n",
    "\n",
    "        # create a sequence of length seq_length\n",
    "        sequence = self.data[start_idx : start_idx + self.seq_length]\n",
    "        label = self.labels[idx]\n",
    "        return sequence, label\n",
    "\n",
    "\n",
    "# create dataset and dataloader\n",
    "dataset = SequentialDatasetWithoutOverlap()\n",
    "dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0)\n",
    "\n",
    "# iterate through the dataset\n",
    "for sequence, label in dataloader:\n",
    "    print(f\"sequence:\\n{sequence}\\nlabel: {label}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SequentialDatasetWithOverlap(Dataset):\n",
    "    def __init__(self):\n",
    "        # original data points\n",
    "        self.data = torch.tensor([[1.0, 0.0], [0.5, 1.5], [1.0, 2.0], [2.0, 1.0], [1.5, 0.5]], dtype=torch.float32)\n",
    "\n",
    "        # labels for each sequence\n",
    "        self.labels = torch.tensor([0, 1, 0], dtype=torch.int64)\n",
    "\n",
    "        # sequence length\n",
    "        self.seq_length = 3\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self.data) - self.seq_length + 1\n",
    "\n",
    "    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:\n",
    "        # create a sequence of length seq_length\n",
    "        sequence = self.data[idx : idx + self.seq_length]\n",
    "        label = self.labels[idx]\n",
    "        return sequence, label\n",
    "\n",
    "\n",
    "# create dataset and dataloader\n",
    "dataset = SequentialDatasetWithOverlap()\n",
    "dataloader = DataLoader(dataset, batch_size=1, shuffle=True)\n",
    "\n",
    "# iterate through the dataset\n",
    "for sequence, label in dataloader:\n",
    "    print(f\"sequence:\\n{sequence}\\nlabel: {label}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc3_'></a>[Types of sequence-to-sequence modeling configurations](#toc0_)\n",
    "\n",
    "1. **One-to-One** (Single Input to Single Output):\n",
    "    - Simplest form of neural network where a single input is mapped to a single output\n",
    "    - Used in a standard feed-forward neural network (e.g. MLP or CNN based architectures)\n",
    "    - e.g. Image classification\n",
    "1. **One-to-Many** (Single Input to Sequence Output):\n",
    "    - A single input is processed by the RNN, which then produces a sequence of outputs over time.\n",
    "    - e.g. Image captioning (an image input resulting in a sequence of words).\n",
    "1. **Many-to-One** (Sequence Input to Single Output):\n",
    "    - The RNN processes each input in the sequence, and the final hidden state is used to produce the output\n",
    "    - e.g. Sentiment analysis (a sequence of words leading to a single sentiment label)\n",
    "1. **Many-to-Many** (Sequence Input to Sequence Output):\n",
    "    - A sequence of inputs leads to a sequence of outputs. This can be further divided into two subcategories:\n",
    "      - **Synchronized** Many-to-Many\n",
    "        - Each input in the sequence has a corresponding output\n",
    "        - The RNN processes a sequence of inputs, producing a corresponding output at each time step\n",
    "        - e.g. Video classification (each frame in a video results in a corresponding label)\n",
    "      - **Asynchronized** Many-to-Many\n",
    "        - The lengths of the input and output sequences can differ\n",
    "        - The RNN processes a sequence of inputs and generates a sequence of outputs which may have different lengths\n",
    "        - e.g. Machine translation (a sequence of words in one language translates to a sequence of words in another language)\n",
    "\n",
    "<figure style=\"text-align: center;\">\n",
    "  <img src=\"../assets/images/original/rnn/seq-to-seq-modeling.svg\" alt=\"seq-to-seq-modeling.svg\" style=\"width: 100%;\">\n",
    "  <figcaption style=\"text-align: center;\">sequence-to-sequence modeling</figcaption>\n",
    "</figure>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc4_'></a>[Network Structure: Recurrent Neural Networks](#toc0_)\n",
    "\n",
    "- Recurrent Neural Networks (RNNs) are a class of neural networks designed for processing sequential data, such as time series, natural language, or speech.\n",
    "- They are characterized by their ability to use information from previous time steps, enabling them to model temporal dependencies effectively.\n",
    "- Unlike feedforward neural networks, RNNs possess a \"memory\" component to process information from previous inputs, influencing the current output.\n",
    "- The same weights are used across all time steps, which reduces the number of parameters and allows learning to generalize better.\n",
    "- RNNs can suffer from vanishing and exploding gradients, making training difficult for long sequences.\n",
    "\n",
    "🧬 **RNN Variants**:\n",
    "\n",
    "- Vanilla RNN\n",
    "- Long Short-Term Memory (LSTM)\n",
    "  - Improves upon the vanilla RNN by introducing gates to control information flow\n",
    "- Gated Recurrent Units (GRU)\n",
    "  - Simplifies the LSTM architecture while maintaining performance\n",
    "\n",
    "🔗 **Usefull Links**:\n",
    "\n",
    "- [karpathy.github.io/2015/05/21/rnn-effectiveness](https://karpathy.github.io/2015/05/21/rnn-effectiveness/)\n",
    "- [stanford.edu/~shervine/teaching/cs-230/cheatsheet-recurrent-neural-networks](https://stanford.edu/~shervine/teaching/cs-230/cheatsheet-recurrent-neural-networks)\n",
    "- [d2l.ai/chapter_recurrent-modern/deep-rnn.html](https://d2l.ai/chapter_recurrent-modern/deep-rnn.html)\n",
    "- [towardsdatascience.com/illustrated-guide-to-lstms-and-gru-s-a-step-by-step-explanation-44e9eb85bf21](https://towardsdatascience.com/illustrated-guide-to-lstms-and-gru-s-a-step-by-step-explanation-44e9eb85bf21)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc4_1_'></a>[Simple Vanilla RNN](#toc0_)\n",
    "\n",
    "🧾 **Notations**:\n",
    "\n",
    "- $\\mathbf{x}_t$: input at time step $t$.\n",
    "- $\\mathbf{h}_t$: Hidden state at time step $t$.\n",
    "- $\\mathbf{y}_t$: Output at time step $t$.\n",
    "- $\\mathbf{W}_{ih}$: Weight matrix for input to hidden\n",
    "- $\\mathbf{W}_{hh}$: Weight matrix for hidden to hidden\n",
    "- $\\mathbf{W}_{ho}$: Weight matrix for hidden to output\n",
    "- $\\mathbf{b}_{ih}$: Bias for input to hidden\n",
    "- $\\mathbf{b}_{hh}$: Bias for hidden to hidden\n",
    "- $\\mathbf{b}_{ho}$: Bias for hidden to output\n",
    "- $\\mathbf{\\sigma}$: Activation function (e.g., Tanh, Sigmoid, ReLU)\n",
    "- $\\mathbf{g}$: Activation function for output (e.g., Softmax for classification)\n",
    "\n",
    "🔬 **Formulations**:\n",
    "\n",
    "- **Hidden State Calculation**:\n",
    "   $$\\mathbf{h}_t = \\sigma(\\mathbf{W}_{ih} \\mathbf{x}_t + \\mathbf{b}_{ih} + \\mathbf{W}_{hh} \\mathbf{h}_{t-1} + \\mathbf{b}_{hh}), \\quad \\mathbf{h}_0 = \\mathbf{0}$$\n",
    "- **Output Calculation**:\n",
    "   $$\\mathbf{y}_t = g(\\mathbf{W}_{ho} \\mathbf{h}_t + \\mathbf{b}_{ho})$$\n",
    "\n",
    "<figure style=\"text-align: center;\">\n",
    "  <img src=\"../assets/images/original/rnn/vanilla-rnn.svg\" alt=\"vanilla-rnn.svg\" style=\"width: 100%;\">\n",
    "  <figcaption style=\"text-align: center;\">Vanilla Recurrent Neural Networks</figcaption>\n",
    "</figure>\n",
    "\n",
    "<figure style=\"text-align: center;\">\n",
    "  <img src=\"../assets/images/original/rnn/calculation.svg\" alt=\"calculation.svg\" style=\"width: 100%;\">\n",
    "  <figcaption style=\"text-align: center;\">Calculations</figcaption>\n",
    "</figure>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VanillaRNN(nn.Module):\n",
    "    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):\n",
    "        super().__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "\n",
    "        # input to hidden connection weights\n",
    "        self.W_ih = nn.Parameter(torch.randn(hidden_dim, input_dim))\n",
    "        # input to hidden connection biases\n",
    "        self.b_ih = nn.Parameter(torch.randn(hidden_dim))\n",
    "\n",
    "        # hidden to hidden connection weights\n",
    "        self.W_hh = nn.Parameter(torch.randn(hidden_dim, hidden_dim))\n",
    "        # hidden to hidden connection biases\n",
    "        self.b_hh = nn.Parameter(torch.randn(hidden_dim))\n",
    "\n",
    "        # weights for hidden to output connection\n",
    "        self.W_ho = nn.Parameter(torch.randn(output_dim, hidden_dim))\n",
    "        # bias for output layer\n",
    "        self.b_ho = nn.Parameter(torch.randn(output_dim))\n",
    "\n",
    "    def forward(self, input: torch.Tensor, hidden: torch.Tensor) -> torch.Tensor:\n",
    "        hidden = torch.tanh(input @ self.W_ih.T + self.b_ih + hidden @ self.W_hh.T + self.b_hh)\n",
    "        output = hidden @ self.W_ho.T + self.b_ho\n",
    "        return output, hidden\n",
    "\n",
    "    def init_hidden(self, batch_size: int) -> torch.Tensor:\n",
    "        # initialize the hidden state with zeros (h_0)\n",
    "        return torch.zeros(batch_size, self.hidden_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "input_dim = 10\n",
    "hidden_dim = 20\n",
    "output_dim = 5\n",
    "num_data = 128\n",
    "sequence_length = 5\n",
    "batch_size = 32\n",
    "\n",
    "# generate synthetic dataset\n",
    "x = torch.randn(num_data, sequence_length, input_dim)\n",
    "y = torch.randn(num_data)\n",
    "\n",
    "# create dataset and dataloader\n",
    "dataset = TensorDataset(x, y)\n",
    "trainsetloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)\n",
    "\n",
    "# log\n",
    "print(f\"x.size()               : {x.size()}\")\n",
    "print(f\"y.size()               : {y.size()}\")\n",
    "print(f\"x.size() [first batch] : {next(iter(trainsetloader))[0].size()}\")\n",
    "print(f\"y.size() [first batch] : {next(iter(trainsetloader))[1].size()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize model\n",
    "rnn_1 = VanillaRNN(input_dim, hidden_dim, output_dim)\n",
    "rnn_1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "summary(rnn_1, input_size=((batch_size, input_dim), (batch_size, hidden_dim)), device=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forward pass through the RNN\n",
    "for c, (x, y_true) in enumerate(trainsetloader):\n",
    "    # initialize hidden state\n",
    "    hidden = rnn_1.init_hidden(batch_size)\n",
    "\n",
    "    for i in range(sequence_length):\n",
    "        y_pred, hidden = rnn_1(x[:, i, :], hidden)\n",
    "        print(f\"batch: {c+1}/{len(trainsetloader)} | time step: {i+1} | hidden.size(): {hidden.size()} | output.size(): {y_pred.size()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc4_2_'></a>[Combined Weights and Concatenated Input and Hidden](#toc0_)\n",
    "\n",
    "- Reformulate the Vanilla RNN by:\n",
    "  - Combining the input-to-hidden and hidden-to-hidden weights into a single weight matrix\n",
    "  - Concatenating the input and hidden states together\n",
    "\n",
    "🧾 **Notations**:\n",
    "\n",
    "- $\\mathbf{x}_t$: Input at time step $t$.\n",
    "- $\\mathbf{h}_t$: Hidden state at time step $t$.\n",
    "- $\\mathbf{y}_t$: Output at time step $t$.\n",
    "- $\\mathbf{W}$: Combined weight matrix\n",
    "- $\\mathbf{b}$: Combined bias vector\n",
    "- $\\mathbf{W}_{ho}$: Weight matrix for hidden to output\n",
    "- $\\mathbf{b}_{ho}$: Bias for hidden to output\n",
    "- $\\mathbf{\\sigma}$: Activation function (e.g., Tanh, Sigmoid, ReLU)\n",
    "- $\\mathbf{g}$: Activation function for output (e.g., Softmax for classification)\n",
    "\n",
    "🔬 **Formulations**:\n",
    "\n",
    "- **Concatenation of Input and Hidden State**:\n",
    "   $$\\mathbf{z}_t = [\\mathbf{x}_t; \\mathbf{h}_{t-1}]$$\n",
    "- **Hidden State Calculation**:\n",
    "   $$\\mathbf{h}_t = \\sigma(\\mathbf{W} \\mathbf{z}_t + \\mathbf{b})$$\n",
    "- **Output Calculation**:\n",
    "   $$\\mathbf{y}_t = g(\\mathbf{W}_{ho} \\mathbf{h}_t + \\mathbf{b}_{ho})$$\n",
    "\n",
    "<figure style=\"text-align: center;\">\n",
    "  <img src=\"../assets/images/original/rnn/combine-weights.svg\" alt=\"combine-weights.svg\" style=\"width: 100%;\">\n",
    "  <figcaption style=\"text-align: center;\">Combining Weights</figcaption>\n",
    "</figure>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VanillaRNN2(nn.Module):\n",
    "    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):\n",
    "        super().__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "\n",
    "        # combined weight matrix for input to hidden and hidden to hidden\n",
    "        self.W = nn.Parameter(torch.randn(hidden_dim, input_dim + hidden_dim))\n",
    "        self.b = nn.Parameter(torch.randn(hidden_dim))\n",
    "\n",
    "        # weights for hidden to output connection\n",
    "        self.W_ho = nn.Parameter(torch.randn(output_dim, hidden_dim))\n",
    "        self.b_ho = nn.Parameter(torch.randn(output_dim))\n",
    "\n",
    "    def forward(self, input: torch.Tensor, hidden: torch.Tensor) -> torch.Tensor:\n",
    "        combined = torch.cat((input, hidden), dim=1)  # concatenate input and hidden state\n",
    "        hidden = torch.tanh(combined @ self.W.T + self.b)\n",
    "        output = hidden @ self.W_ho.T + self.b_ho\n",
    "        return output, hidden\n",
    "\n",
    "    def init_hidden(self, batch_size: int) -> torch.Tensor:\n",
    "        # initialize the hidden state with zeros (h_0)\n",
    "        return torch.zeros(batch_size, self.hidden_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "input_dim = 10\n",
    "hidden_dim = 20\n",
    "output_dim = 5\n",
    "num_data = 128\n",
    "sequence_length = 5\n",
    "batch_size = 32\n",
    "\n",
    "# generate synthetic dataset\n",
    "x = torch.randn(num_data, sequence_length, input_dim)\n",
    "y = torch.randn(num_data)\n",
    "\n",
    "# create dataset and dataloader\n",
    "dataset = TensorDataset(x, y)\n",
    "trainsetloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)\n",
    "\n",
    "# log\n",
    "print(f\"x.size()               : {x.size()}\")\n",
    "print(f\"y.size()               : {y.size()}\")\n",
    "print(f\"x.size() [first batch] : {next(iter(trainsetloader))[0].size()}\")\n",
    "print(f\"y.size() [first batch] : {next(iter(trainsetloader))[1].size()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize model\n",
    "rnn_2 = VanillaRNN2(input_dim, hidden_dim, output_dim)\n",
    "rnn_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "summary(rnn_2, input_size=((batch_size, input_dim), hidden.size()), device=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forward pass through the RNN\n",
    "for c, (x, y_true) in enumerate(trainsetloader):\n",
    "    # initialize hidden state\n",
    "    hidden = rnn_2.init_hidden(batch_size)\n",
    "\n",
    "    for i in range(sequence_length):\n",
    "        y_pred, hidden = rnn_2(x[:, i, :], hidden)\n",
    "        print(f\"batch: {c+1}/{len(trainsetloader)} | time step: {i+1} | hidden.size(): {hidden.size()} | output.size(): {y_pred.size()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc4_3_'></a>[Deep RNN](#toc0_)\n",
    "\n",
    "- A **Deep RNN** consists of **multiple** layers of RNN cells stacked on top of each other.\n",
    "- Each layer processes the **hidden states** of the layer below as its **input**.\n",
    "- The **output** of one layer is used as the **input** to the **next layer**.\n",
    "\n",
    "🧾 **Notations**:\n",
    "\n",
    "- $\\mathbf{x}_t$: Input at time step $t$.\n",
    "- $\\mathbf{h}^l_t$: Hidden state at time step $t$ in layer $l$.\n",
    "- $\\mathbf{y}_t$: Output at time step $t$.\n",
    "- $\\mathbf{W}^l$: Combined weight matrix for layer $l$\n",
    "- $\\mathbf{b}^l$: Bias vector for layer $l$\n",
    "- $\\mathbf{W}_{ho}$: Weight matrix for hidden to output\n",
    "- $\\mathbf{b}_{ho}$: Bias for hidden to output\n",
    "- $\\mathbf{\\sigma}$: Activation function (e.g., Tanh, Sigmoid, ReLU)\n",
    "- $\\mathbf{g}$: Activation function for output (e.g., Softmax for classification)\n",
    "- $\\mathbf{L}$: Number of layers\n",
    "\n",
    "🔬 **Formulations**:\n",
    "\n",
    "- **Concatenation of Input and Hidden State for Layer 1**:\n",
    "   $$\\mathbf{z}_t^1 = [\\mathbf{x}_t; \\mathbf{h}_{t-1}^1]$$\n",
    "- **Hidden State Calculation for Layer 1**:\n",
    "   $$\\mathbf{h}_t^1 = \\sigma(\\mathbf{W}^1 \\mathbf{z}_t^1 + \\mathbf{b}^1)$$\n",
    "- **Concatenation of Hidden States for Subsequent Layers**:\n",
    "   $$\\mathbf{z}_t^l = [\\mathbf{h}_t^{l-1}; \\mathbf{h}_{t-1}^l] \\quad \\text{for} \\quad l = 2, \\ldots, L$$\n",
    "- **Hidden State Calculation for Subsequent Layers**:\n",
    "   $$\\mathbf{h}_t^l = \\sigma(\\mathbf{W}^l \\mathbf{z}_t^l + \\mathbf{b}^l) \\quad \\text{for} \\quad l = 2, \\ldots, L$$\n",
    "- **Output Calculation**:\n",
    "   $$\\mathbf{y}_t = g(\\mathbf{W}_{ho} \\mathbf{h}_t^L + \\mathbf{b}_{ho})$$\n",
    "\n",
    "<figure style=\"text-align: center;\">\n",
    "  <img src=\"../assets/images/original/rnn/deep-rnn.svg\" alt=\"deep-rnn.svg\" style=\"width: 100%;\">\n",
    "  <figcaption style=\"text-align: center;\">Deep Recurrent Neural Networks</figcaption>\n",
    "</figure>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DeepRNN(nn.Module):\n",
    "    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int):\n",
    "        super().__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "        # define RNN layers\n",
    "        self.rnn_layers = nn.ModuleList()\n",
    "        for i in range(num_layers):\n",
    "            if i == 0:\n",
    "                self.rnn_layers.append(nn.Linear(input_dim + hidden_dim, hidden_dim))\n",
    "            else:\n",
    "                self.rnn_layers.append(nn.Linear(hidden_dim + hidden_dim, hidden_dim))\n",
    "\n",
    "        # define the output layer\n",
    "        self.output_layer = nn.Linear(hidden_dim, output_dim)\n",
    "\n",
    "    def forward(self, input: torch.Tensor, hidden: torch.Tensor) -> torch.Tensor:\n",
    "        # concatenate input and the first hidden state along the feature dimension\n",
    "        combined_input = torch.cat((input, hidden[0]), dim=1)\n",
    "        new_hidden = []\n",
    "\n",
    "        for i, rnn_layer in enumerate(self.rnn_layers):\n",
    "            hidden_state = torch.tanh(rnn_layer(combined_input))\n",
    "            new_hidden.append(hidden_state)\n",
    "\n",
    "            # concatenate the current hidden state with the previous one\n",
    "            combined_input = torch.cat((hidden_state, hidden[i]), dim=1)\n",
    "\n",
    "        # use the last hidden state for output\n",
    "        final_hidden = new_hidden[-1]\n",
    "        output = self.output_layer(final_hidden)\n",
    "        return output, torch.stack(new_hidden)\n",
    "\n",
    "    def init_hidden(self, batch_size: int) -> torch.Tensor:\n",
    "        # initialize hidden state with zeros for each layer and batch\n",
    "        return torch.zeros(self.num_layers, batch_size, self.hidden_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "input_dim = 10\n",
    "hidden_dim = 20\n",
    "output_dim = 5\n",
    "num_layers = 3\n",
    "num_data = 128\n",
    "sequence_length = 5\n",
    "batch_size = 32\n",
    "\n",
    "# generate synthetic dataset\n",
    "x = torch.randn(num_data, sequence_length, input_dim)\n",
    "y = torch.randn(num_data)\n",
    "\n",
    "# create dataset and dataloader\n",
    "dataset = TensorDataset(x, y)\n",
    "trainsetloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize model\n",
    "deep_rnn = DeepRNN(input_dim, hidden_dim, output_dim, num_layers)\n",
    "deep_rnn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "summary(deep_rnn, input_size=((batch_size, input_dim), (num_layers, batch_size, hidden_dim)), device=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forward pass through the RNN\n",
    "for c, (x, y_true) in enumerate(trainsetloader):\n",
    "    # initialize hidden state for each batch\n",
    "    hidden = deep_rnn.init_hidden(batch_size)\n",
    "    print(hidden.size())\n",
    "\n",
    "    for i in range(sequence_length):\n",
    "        y_pred, hidden = deep_rnn(x[:, i, :], hidden)\n",
    "        print(f\"Batch: {c+1}/{len(trainsetloader)} | Time step: {i+1} | hidden.size(): {hidden.size()} | output.size(): {y_pred.size()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc4_4_'></a>[RNN using PyTorch](#toc0_)\n",
    "\n",
    "📝 **Docs**:\n",
    "\n",
    "- `torch.nn.RNN`: [pytorch.org/docs/stable/generated/torch.nn.RNN.html](https://pytorch.org/docs/stable/generated/torch.nn.RNN.html)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RNN model for sequence-to-one tasks\n",
    "class RNN(nn.Module):\n",
    "    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int = 1):\n",
    "        super().__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "        self.rnn = nn.RNN(input_dim, hidden_dim, num_layers, batch_first=True)\n",
    "        self.fc = nn.Linear(hidden_dim, output_dim)\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        # initialize hidden state with zeros\n",
    "        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(x.device)\n",
    "\n",
    "        # forward propagate RNN\n",
    "        # out : output of the RNN last layer for each sequence in a batch for each time step [(batch_size, seq_length, hidden_dim)].\n",
    "        # _   : the final hidden state (often denoted as hn) of the RNN [(num_layers, batch_size, hidden_dim)].\n",
    "        out, _ = self.rnn(x, h0)\n",
    "\n",
    "        # decode the hidden state of the last time step [seq-to-one modeling]\n",
    "        # :  -> selects all elements along the first dimension (typically batch size).\n",
    "        # -1 -> selects the last element along the second dimension (which represents the sequence length)\n",
    "        # :  -> selects all elements along the third dimension (feature dimension)\n",
    "        out = self.fc(out[:, -1, :])\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "input_dim = 10\n",
    "hidden_dim = 20\n",
    "output_dim = 5\n",
    "num_layers = 1\n",
    "num_data = 128\n",
    "sequence_length = 5\n",
    "batch_size = 32\n",
    "\n",
    "# generate synthetic dataset\n",
    "x = torch.randn(num_data, sequence_length, input_dim)\n",
    "y = torch.randn(num_data, output_dim)\n",
    "\n",
    "# create dataset and dataloader\n",
    "dataset = TensorDataset(x, y)\n",
    "trainsetloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize model\n",
    "rnn_3 = RNN(input_dim, hidden_dim, output_dim, num_layers)\n",
    "rnn_3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "summary(rnn_3, input_size=(batch_size, *x.size()[1:]), device=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forward pass through the RNN\n",
    "for c, (x, y_true) in enumerate(trainsetloader):\n",
    "    y_pred = rnn_3(x)\n",
    "    print(f\"batch: {c+1}/{len(trainsetloader)} | output.size(): {y_pred.size()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc4_5_'></a>[Long Short-Term Memory (LSTM)](#toc0_)\n",
    "\n",
    "- A type of recurrent neural network (RNN) designed to address the **Vanishing Gradient** problem inherent in traditional RNNs.\n",
    "- **Long Short-Term Memory** signifies a system capable of remembering information over both **long** and **short** durations of time.\n",
    "- **Vanilla RNNs** primarily had **short-term** memory due to their design.\n",
    "- It is based on the [**Long Short-term Memory**](https://www.researchgate.net/publication/13853244_Long_Short-term_Memory) paper, Developed in **1997** by [*Sepp Hochreiter*](https://scholar.google.at/citations?user=tvUH3WMAAAAJ&hl=en) and [*Jürgen Schmidhuber*](https://scholar.google.com/citations?user=gLnCTgIAAAAJ&hl=en).\n",
    "\n",
    "🧾 **Notations**:\n",
    "\n",
    "- $\\mathbf{x}_t$: Input vector at time step $t$\n",
    "- $\\mathbf{h}_t$: Hidden state vector at time step $t$\n",
    "- $\\mathbf{c}_t$: Cell state vector at time step $t$\n",
    "- $\\mathbf{W}_f$: Weight matrix for the forget gate (combined input and hidden state)\n",
    "- $\\mathbf{W}_i$: Weight matrix for the input gate (combined input and hidden state)\n",
    "- $\\mathbf{W}_c$: Weight matrix for the candidate cell state (combined input and hidden state)\n",
    "- $\\mathbf{W}_o$: Weight matrix for the output gate (combined input and hidden state)\n",
    "- $\\mathbf{b}_f$: Bias vector for the forget gate\n",
    "- $\\mathbf{b}_i$: Bias vector for the input gate\n",
    "- $\\mathbf{b}_c$: Bias vector for the candidate cell state\n",
    "- $\\mathbf{b}_o$: Bias vector for the output gate\n",
    "\n",
    "🔬 **Formulations**:\n",
    "\n",
    "- **Concatenation of Input and Hidden State**:\n",
    "   $$\\mathbf{z}_t = [\\mathbf{x}_t; \\mathbf{h}_{t-1}]$$\n",
    "- **Forget Gate**\n",
    "   $$\\mathbf{f}_t = \\sigma(\\mathbf{W}_f \\mathbf{z}_t + \\mathbf{b}_f)$$\n",
    "- **Input Gate**\n",
    "   $$\\mathbf{i}_t = \\sigma(\\mathbf{W}_i \\mathbf{z}_t + \\mathbf{b}_i)$$\n",
    "- **Candidate Cell State**\n",
    "   $$\\tilde{\\mathbf{c}}_t = \\tanh(\\mathbf{W}_c \\mathbf{z}_t + \\mathbf{b}_c)$$\n",
    "- **Cell State**\n",
    "   $$\\mathbf{c}_t = \\mathbf{f}_t \\odot \\mathbf{c}_{t-1} + \\mathbf{i}_t \\odot \\tilde{\\mathbf{c}}_t$$\n",
    "- **Output Gate**\n",
    "   $$\\mathbf{o}_t = \\sigma(\\mathbf{W}_o \\mathbf{z}_t + \\mathbf{b}_o)$$\n",
    "- **Hidden State**\n",
    "   $$\\mathbf{h}_t = \\mathbf{o}_t \\odot \\tanh(\\mathbf{c}_t)$$\n",
    "\n",
    "✍️ **Notes**:\n",
    "\n",
    "- The lack of `Weights` in the `Cell State`, allows the long-term memories to flow through a series of unrolled units without causing the gradient to explode or vanish.  \n",
    "- **LSTMs** do not directly solve **Exploding Gradients** but are often **less prone** to it because their structure avoids excessively amplifying gradients during backpropagation.\n",
    "- **LSTMs** can capture **longer sequences** and handle long-term dependencies more effectively than vanilla RNNs due to controlling **Vanishing/Exploding Gradient**.\n",
    "- The **hidden state** $h_t$ (**short-term** memory) at each time step is the typical **output** for that specific time step.\n",
    "- The **cell state** $C_t$ (**long-term** memory) is internal to the LSTM and **is not** used directly as the **output**.\n",
    "\n",
    "📝 **Docs**:\n",
    "\n",
    "- [pytorch.org/docs/stable/generated/torch.nn.LSTM.html](https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html)\n",
    "\n",
    "<figure style=\"text-align: center;\">\n",
    "  <img src=\"../assets/images/original/rnn/lstm.svg\" alt=\"lstm.svg\" style=\"width: 100%;\">\n",
    "  <figcaption style=\"text-align: center;\">Long Short-Term Memory (LSTM)</figcaption>\n",
    "</figure>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LSTM model for sequence-to-one tasks\n",
    "class LSTMModel(nn.Module):\n",
    "    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int):\n",
    "        super().__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True)\n",
    "        self.fc = nn.Linear(hidden_dim, output_dim)\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(x.device)\n",
    "        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(x.device)\n",
    "\n",
    "        out, _ = self.lstm(x, (h0, c0))\n",
    "        out = self.fc(out[:, -1, :])\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "input_dim = 10\n",
    "hidden_dim = 20\n",
    "output_dim = 5\n",
    "num_layers = 2\n",
    "num_data = 128\n",
    "sequence_length = 5\n",
    "batch_size = 32\n",
    "\n",
    "# generate synthetic dataset\n",
    "x = torch.randn(num_data, sequence_length, input_dim)\n",
    "y = torch.randn(num_data, output_dim)\n",
    "\n",
    "# create dataset and dataloader\n",
    "dataset = TensorDataset(x, y)\n",
    "trainsetloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize model\n",
    "lstm = LSTMModel(input_dim, hidden_dim, output_dim, num_layers)\n",
    "lstm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "summary(lstm, input_size=(batch_size, *x.size()[1:]), device=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forward pass through the LSTM\n",
    "for c, (x_batch, y_true) in enumerate(trainsetloader):\n",
    "    y_pred = lstm(x_batch)\n",
    "    print(f\"Batch: {c+1}/{len(trainsetloader)} | Output Size: {y_pred.size()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc4_6_'></a>[Gated Recurrent Units (GRU)](#toc0_)\n",
    "\n",
    "- A gating mechanism in recurrent neural networks, introduced in 2014 by [*Kyunghyun*](https://dblp.uni-trier.de/search/author?author=Kyunghyun%20Cho).\n",
    "- It is based on the [**Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling**](https://arxiv.org/abs/1412.3555) paper.\n",
    "- Similar to LSTM but lacks a `context vector` or `output gate`, resulting in fewer parameters than LSTM.\n",
    "\n",
    "🧾 **Notations**:\n",
    "\n",
    "- $\\mathbf{x}_t$: Input vector at time step $t$\n",
    "- $\\mathbf{h}_t$: Hidden state vector at time step $t$\n",
    "- $\\mathbf{c}_t$: Concatenated input and hidden state vector at time step $t$\n",
    "- $\\mathbf{W}_z$: Weight matrix for the update gate\n",
    "- $\\mathbf{W}_r$: Weight matrix for the reset gate\n",
    "- $\\mathbf{W}_h$: Weight matrix for the candidate hidden state\n",
    "- $\\mathbf{b}_z$: Bias vector for the update gate\n",
    "- $\\mathbf{b}_r$: Bias vector for the reset gate\n",
    "- $\\mathbf{b}_h$: Bias vector for the candidate hidden state\n",
    "- $\\mathbf{z}_t$: Update gate vector at time step $t$\n",
    "- $\\mathbf{r}_t$: Reset gate vector at time step $t$\n",
    "- $\\mathbf{\\tilde{h}}_t$: Candidate hidden state vector at time step $t$\n",
    "\n",
    "🔬 **Formulations**:\n",
    "\n",
    "- **Concatenated Input and Hidden State**:\n",
    "   $$\\mathbf{c}_t = [\\mathbf{x}_t; \\mathbf{h}_{t-1}]$$\n",
    "- **Reset Gate**:\n",
    "   $$\\mathbf{r}_t = \\sigma(\\mathbf{W}_r \\mathbf{c}_t + \\mathbf{b}_r)$$\n",
    "- **Update Gate**:\n",
    "   $$\\mathbf{z}_t = \\sigma(\\mathbf{W}_z \\mathbf{c}_t + \\mathbf{b}_z)$$\n",
    "- **Candidate Hidden State**:\n",
    "   $$\\tilde{\\mathbf{h}}_t = \\tanh(\\mathbf{W}_h [\\mathbf{x}_t; (\\mathbf{r}_t \\odot \\mathbf{h}_{t-1})] + \\mathbf{b}_h)$$\n",
    "- **Hidden State**:\n",
    "   $$\\mathbf{h}_t = (1 - \\mathbf{z}_t) \\odot \\mathbf{h}_{t-1} + \\mathbf{z}_t \\odot \\tilde{\\mathbf{h}}_t$$\n",
    "\n",
    "📝 **Docs**:\n",
    "\n",
    "- [pytorch.org/docs/stable/generated/torch.nn.GRU.html](https://pytorch.org/docs/stable/generated/torch.nn.GRU.html)\n",
    "\n",
    "<figure style=\"text-align: center;\">\n",
    "  <img src=\"../assets/images/original/rnn/gru.svg\" alt=\"gru.svg\" style=\"width: 100%;\">\n",
    "  <figcaption style=\"text-align: center;\">Gated Recurrent Units (GRU)</figcaption>\n",
    "</figure>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GRUModel(nn.Module):\n",
    "    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int):\n",
    "        super().__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "        self.gru = nn.GRU(input_dim, hidden_dim, num_layers, batch_first=True)\n",
    "        self.fc = nn.Linear(hidden_dim, output_dim)\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(x.device)\n",
    "\n",
    "        out, _ = self.gru(x, h0)\n",
    "        out = self.fc(out[:, -1, :])\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "input_dim = 10\n",
    "hidden_dim = 20\n",
    "output_dim = 5\n",
    "num_layers = 2\n",
    "num_data = 128\n",
    "sequence_length = 5\n",
    "batch_size = 32\n",
    "\n",
    "# generate synthetic dataset\n",
    "x = torch.randn(num_data, sequence_length, input_dim)\n",
    "y = torch.randn(num_data, output_dim)\n",
    "\n",
    "# create dataset and dataloader\n",
    "dataset = TensorDataset(x, y)\n",
    "trainsetloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize model\n",
    "gru = GRUModel(input_dim, hidden_dim, output_dim, num_layers)\n",
    "gru"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "summary(gru, input_size=(batch_size, *x.size()[1:]), device=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forward pass through the GRU\n",
    "for c, (x_batch, y_true) in enumerate(trainsetloader):\n",
    "    y_pred = gru(x_batch)\n",
    "    print(f\"Batch: {c+1}/{len(trainsetloader)} | Output Size: {y_pred.size()}\")"
   ]
  }
 ],
 "metadata": {
  "author_email": "AmirhosseinHeydari78@gmail.com",
  "author_github": "https://github.com/mr-pylin",
  "author_name": "Amirhossein Heydari",
  "kernelspec": {
   "display_name": "pytorch-workshop-U_zYfVTd-py3.12",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  },
  "origin_repo": "https://github.com/mr-pylin/pytorch-workshop"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}