notes/gradient_computation_from_scratch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model definition\n",
    "\n",
    "Suppose we want to train a model to predict the price of a house (`y`) based on two variables: the number of rooms in the house `x1`, and the number of bathrooms (`x2`). We can define the dataset as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data format:\n",
    "# [x1, x2] --> [y]\n",
    "# [NUM_ROOMS, NUM_BATHS] --> [PRICE]\n",
    "\n",
    "training_data = [\n",
    "    [torch.tensor([6, 2], dtype=torch.float), torch.tensor([15], dtype=torch.float)],\n",
    "    [torch.tensor([5, 2], dtype=torch.float), torch.tensor([12], dtype=torch.float)],\n",
    "    [torch.tensor([5, 1], dtype=torch.float), torch.tensor([10], dtype=torch.float)],\n",
    "    [torch.tensor([3, 1], dtype=torch.float), torch.tensor([7], dtype=torch.float)],\n",
    "    [torch.tensor([2, 1], dtype=torch.float), torch.tensor([4.5], dtype=torch.float)],\n",
    "    [torch.tensor([2, 0], dtype=torch.float), torch.tensor([4], dtype=torch.float)],\n",
    "    [torch.tensor([1, 0], dtype=torch.float), torch.tensor([2], dtype=torch.float)],\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppose we know that the relation between the two variables (`x1` and `x2`) and the target variable (`y`) is linear, i.e.,\n",
    "\n",
    "$$y = w_1 \\cdot x_1 + w_2 \\cdot x_2 + b$$\n",
    "\n",
    "where `w1` and `w2` are the weights of the model, and `b` is the bias term.\n",
    "\n",
    "We want to train the model using gradient descent to find the optimal values of `w1`, `w2`, and `b` that minimize the mean squared error (MSE) between the predicted values and the actual values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the model parameters\n",
    "class ModelParameters:\n",
    "    \n",
    "    def __init__(self):\n",
    "        self.w1 = torch.tensor(0.773, dtype=torch.float, requires_grad=True)\n",
    "        self.w2 = torch.tensor(0.321, dtype=torch.float, requires_grad=True)\n",
    "        self.b = torch.tensor(0.067, dtype=torch.float, requires_grad=True)\n",
    "\n",
    "# We will use two training loops: the first one without gradient accumulation, and the second one with gradient accumulation.\n",
    "params_no_accumulate = ModelParameters()\n",
    "params_accumulate = ModelParameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training loop (without gradient accumulation)\n",
    "\n",
    "We run gradient descent using one data item at a time, we calculate the gradient of the loss function w.r.t the parameters, and update the parameters using the gradients at each iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial parameters: w1: 0.773, w2: 0.321, b: 0.067\n",
      "Epoch   1 - Loss:     0.7001\n",
      "Epoch   2 - Loss:     0.3374\n",
      "Epoch   3 - Loss:     0.1454\n",
      "Epoch   4 - Loss:      0.051\n",
      "Epoch   5 - Loss:      0.011\n",
      "Epoch   6 - Loss:     0.0001\n",
      "Epoch   7 - Loss:     0.0039\n",
      "Epoch   8 - Loss:     0.0142\n",
      "Epoch   9 - Loss:     0.0265\n",
      "Epoch  10 - Loss:     0.0387\n",
      "Final parameters: w1: 1.897, w2: 0.692, b: 0.299\n"
     ]
    }
   ],
   "source": [
    "def train_no_accumulate(params: ModelParameters, num_epochs: int = 10, learning_rate: float = 1e-3):\n",
    "    print(f'Initial parameters: w1: {params.w1.item():.3f}, w2: {params.w2.item():.3f}, b: {params.b.item():.3f}')\n",
    "    for epoch in range(1, num_epochs+1):\n",
    "        for (x1, x2), y_target in training_data:\n",
    "            # Calculate the output of the model\n",
    "            z1 = x1 * params.w1\n",
    "            z1.retain_grad()\n",
    "            z2 = x2 * params.w2\n",
    "            z2.retain_grad()\n",
    "            y_pred = z1 + z2 + params.b\n",
    "            y_pred.retain_grad()\n",
    "            loss = (y_pred - y_target) ** 2\n",
    "\n",
    "            # Calculate the gradients of the loss w.r.t. the parameters\n",
    "            loss.backward()\n",
    "\n",
    "            # Update the parameters (at each iteration)\n",
    "            with torch.no_grad():\n",
    "                # Equivalent to calling optimizer.step()\n",
    "                params.w1 -= learning_rate * params.w1.grad\n",
    "                params.w2 -= learning_rate * params.w2.grad\n",
    "                params.b -= learning_rate * params.b.grad\n",
    "\n",
    "                # Reset the gradients to zero\n",
    "                # Equivalent to calling optimizer.zero_grad()\n",
    "                params.w1.grad.zero_()\n",
    "                params.w2.grad.zero_()\n",
    "                params.b.grad.zero_()\n",
    "        print(f\"Epoch {epoch:>3} - Loss: {np.round(loss.item(),4):>10}\")\n",
    "    print(f'Final parameters: w1: {params.w1.item():.3f}, w2: {params.w2.item():.3f}, b: {params.b.item():.3f}')\n",
    "        \n",
    "train_no_accumulate(params_no_accumulate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training loop (with gradient accumulation)\n",
    "\n",
    "We run gradient descent using one data item at a time, but we accumulate the gradients over a fixed number of iterations (batch size) before updating the parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial parameters: w1: 0.773, w2: 0.321, b: 0.067\n",
      "Epoch   1 - Loss:     0.6857\n",
      "Epoch   2 - Loss:     0.3218\n",
      "Epoch   3 - Loss:     0.1335\n",
      "Epoch   4 - Loss:     0.0438\n",
      "Epoch   5 - Loss:     0.0078\n",
      "Epoch   6 - Loss:        0.0\n",
      "Epoch   7 - Loss:     0.0059\n",
      "Epoch   8 - Loss:     0.0174\n",
      "Epoch   9 - Loss:     0.0303\n",
      "Epoch  10 - Loss:     0.0427\n",
      "Final parameters: w1: 1.905, w2: 0.698, b: 0.300\n"
     ]
    }
   ],
   "source": [
    "def train_accumulate(params: ModelParameters, num_epochs: int = 10, learning_rate: float = 1e-3, batch_size: int = 2):\n",
    "    print(f'Initial parameters: w1: {params.w1.item():.3f}, w2: {params.w2.item():.3f}, b: {params.b.item():.3f}')\n",
    "    for epoch in range(1, num_epochs+1):\n",
    "        for index, ((x1, x2), y_target) in enumerate(training_data):\n",
    "            # Calculate the output of the model\n",
    "            z1 = x1 * params.w1\n",
    "            z1.retain_grad()\n",
    "            z2 = x2 * params.w2\n",
    "            z2.retain_grad()\n",
    "            y_pred = z1 + z2 + params.b\n",
    "            y_pred.retain_grad()\n",
    "            loss = (y_pred - y_target) ** 2\n",
    "\n",
    "            # We can also divide the loss by the batch size (equivalent to using nn.MSE loss with the paraemter reduction='mean')\n",
    "            # If we don't divide by the batch size, then it is equivalent to using nn.MSE loss with the parameter reduction='sum'\n",
    "\n",
    "            # Calculate the gradients of the loss w.r.t. the parameters\n",
    "            # If we didn't call zero_() on the gradients on the previous iteration, then the gradients will accumulate (add up) over each iteration\n",
    "            loss.backward()\n",
    "\n",
    "            # Everytime we reach the batch size or the end of the dataset, update the parameters\n",
    "            if (index + 1) % batch_size == 0 or index == len(training_data) - 1:\n",
    "                with torch.no_grad():\n",
    "                    # Equivalent to calling optimizer.step()\n",
    "                    params.w1 -= learning_rate * params.w1.grad\n",
    "                    params.w2 -= learning_rate * params.w2.grad\n",
    "                    params.b -= learning_rate * params.b.grad\n",
    "\n",
    "                    # Reset the gradients to zero\n",
    "                    # Equivalent to calling optimizer.zero_grad()\n",
    "                    params.w1.grad.zero_()\n",
    "                    params.w2.grad.zero_()\n",
    "                    params.b.grad.zero_()\n",
    "\n",
    "        print(f\"Epoch {epoch:>3} - Loss: {np.round(loss.item(),4):>10}\")\n",
    "    print(f'Final parameters: w1: {params.w1.item():.3f}, w2: {params.w2.item():.3f}, b: {params.b.item():.3f}')\n",
    "\n",
    "train_accumulate(params_accumulate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lazy-ml",
   "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}