{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  },
  "colab": {
   "provenance": []
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cell-0000",
   "metadata": {},
   "source": [
    "# The Training Loop\n",
    "\n",
    "Model, criterion, optimizer — the three ingredients of training — and the forward-loss-backward-step dance that turns gradients into learning.\n",
    "\n",
    "*Part of the free [Deep Learning with PyTorch](https://ramadnsyh.dev/courses/deep-learning) course by [Muhammad Ramadiansyah](https://ramadnsyh.dev). This notebook is generated from the interactive lesson — [read it online](https://ramadnsyh.dev/courses/deep-learning/training-neural-networks).*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0001",
   "metadata": {},
   "source": [
    "You know what a neural network computes, and you know autograd can produce\n",
    "gradients for free. This lesson connects the two: the **training loop**, the\n",
    "handful of lines at the heart of every PyTorch project. Once you can read and\n",
    "write this loop, everything from linear regression to giant language models is\n",
    "a variation on the same theme."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0002",
   "metadata": {},
   "source": [
    "> **Run this lesson in Colab**\n",
    "> \n",
    "> The code here uses PyTorch, which doesn't run in the browser. Download the\n",
    "> notebook and run it in Google Colab — no GPU needed yet, this lesson's models\n",
    "> train in seconds on CPU."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0003",
   "metadata": {},
   "source": [
    "## Three ingredients before you train: MCO\n",
    "\n",
    "Every training script starts by preparing three objects. A handy mnemonic:\n",
    "**MCO — Model, Criterion, Optimizer**."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0004",
   "metadata": {},
   "source": [
    "### M is for Model\n",
    "\n",
    "The thing that makes predictions. PyTorch's `nn` module gives you ready-made\n",
    "layers; `nn.Sequential` chains them into a network:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0005",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "# 4 inputs -> hidden(3, ReLU) -> hidden(4, ReLU) -> 3 outputs\n",
    "model = nn.Sequential(\n",
    "    nn.Linear(4, 3),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(3, 4),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(4, 3),\n",
    ")\n",
    "\n",
    "print(model)\n",
    "print(model.state_dict().keys())   # every learnable weight and bias"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0006",
   "metadata": {},
   "source": [
    "Each `nn.Linear(in, out)` is exactly the weighted-sum-plus-bias you built in\n",
    "NumPy — a weight matrix and a bias vector, initialized randomly. The\n",
    "`state_dict()` is the model's memory: all its learnable parameters, which is\n",
    "what training will change and what you'll later save to disk."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0007",
   "metadata": {},
   "source": [
    "### C is for Criterion\n",
    "\n",
    "The **criterion** (loss function) measures how wrong the predictions are —\n",
    "one number, where lower is better. You pick it based on the task:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0008",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "criterion = nn.MSELoss()              # regression (model ends in plain Linear)\n",
    "criterion = nn.BCEWithLogitsLoss()    # binary classification (1 output, raw logit)\n",
    "criterion = nn.CrossEntropyLoss()     # multiclass (N outputs, raw logits)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0009",
   "metadata": {},
   "source": [
    "Note a modern PyTorch convention: for classification the model outputs **raw\n",
    "logits** — no sigmoid or softmax layer at the end. `BCEWithLogitsLoss` and\n",
    "`CrossEntropyLoss` apply the squashing internally, which is both numerically\n",
    "more stable and less code. Only add a `sigmoid`/`softmax` yourself when you\n",
    "need actual probabilities at prediction time."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0010",
   "metadata": {},
   "source": [
    "### O is for Optimizer\n",
    "\n",
    "The optimizer updates the weights using the gradients. It's gradient descent\n",
    "with an engine attached:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0011",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from torch import optim\n",
    "\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.01)    # classic gradient descent\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)  # adaptive — great default"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0012",
   "metadata": {},
   "source": [
    "You hand it `model.parameters()` — the list of tensors it's allowed to modify\n",
    "— and a **learning rate**. You already know why the learning rate is the most\n",
    "important knob; here's a refresher on the three regimes:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0013",
   "metadata": {},
   "source": [
    "> 🎛️ **Interactive demo** — this section has a hands-on visualization in the web version of this lesson: [open it here](https://ramadnsyh.dev/courses/deep-learning/training-neural-networks)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0014",
   "metadata": {},
   "source": [
    "Too small and training crawls; too large and the loss bounces or explodes.\n",
    "**Adam** adapts a per-parameter step size on the fly, which makes it far more\n",
    "forgiving of the initial learning-rate choice — that's why `Adam` (or its\n",
    "sibling `AdamW`) with `lr=0.001` is the standard starting point, while plain\n",
    "`SGD` typically needs more tuning to shine."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0015",
   "metadata": {},
   "source": [
    "## The loop: forward, loss, backward, step\n",
    "\n",
    "With MCO in place, one training step is four moves:\n",
    "\n",
    "1. **Forward pass** — `output = model(x)`: push data through the network.\n",
    "2. **Compute loss** — `loss = criterion(output, y)`: one number measuring error.\n",
    "3. **Backward pass** — `loss.backward()`: autograd fills `.grad` for every parameter.\n",
    "4. **Update** — `optimizer.step()`: nudge every weight downhill, `w ← w − lr·grad`.\n",
    "\n",
    "And one bookkeeping move: **`optimizer.zero_grad()`**. Remember from the last\n",
    "lesson that gradients *accumulate* — each `backward()` adds to `.grad` instead\n",
    "of overwriting it. Without zeroing, step 2's gradients would stack on top of\n",
    "step 1's, and the updates would be garbage. So every iteration clears the\n",
    "gradients before (or right after) the update.\n",
    "\n",
    "Repeat the whole dance many times. Each full pass through the training data is\n",
    "called an **epoch**."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0016",
   "metadata": {},
   "source": [
    "## A complete minimal example\n",
    "\n",
    "Let's train the smallest possible network — a single `nn.Linear(1, 1)`, i.e.\n",
    "`y = wx + b` — to recover a line from noisy synthetic data. Every real\n",
    "training script you'll ever write has this exact skeleton:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0017",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn, optim\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "torch.manual_seed(42)\n",
    "\n",
    "# Synthetic data: y = 2x - 1 + noise\n",
    "X = torch.rand(100, 1) * 10                 # shape [100, 1]\n",
    "y = 2 * X - 1 + torch.randn(100, 1) * 0.8   # shape [100, 1]\n",
    "\n",
    "# MCO\n",
    "model = nn.Linear(1, 1)\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.1)\n",
    "\n",
    "# Training loop\n",
    "losses = []\n",
    "for epoch in range(200):\n",
    "    output = model(X)                # 1. forward\n",
    "    loss = criterion(output, y)      # 2. loss\n",
    "\n",
    "    optimizer.zero_grad()            #    clear old gradients\n",
    "    loss.backward()                  # 3. backward\n",
    "    optimizer.step()                 # 4. update\n",
    "\n",
    "    losses.append(loss.item())\n",
    "    if (epoch + 1) % 50 == 0:\n",
    "        print(f\"epoch {epoch+1:3d} | loss {loss.item():.4f}\")\n",
    "\n",
    "w, b = model.weight.item(), model.bias.item()\n",
    "print(f\"\\nlearned: y = {w:.2f}x + {b:.2f}   (true: y = 2.00x - 1.00)\")\n",
    "\n",
    "plt.plot(losses)\n",
    "plt.xlabel(\"epoch\"); plt.ylabel(\"MSE loss\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0018",
   "metadata": {},
   "source": [
    "That's it. A 175-million-parameter network trains with the same eight lines —\n",
    "only the model, the data, and the number of epochs change."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0019",
   "metadata": {},
   "source": [
    "> **loss.item(), not loss**\n",
    "> \n",
    "> The loss is a tensor still attached to the autograd graph. Store or print\n",
    "> `loss.item()` (a plain Python float) — appending the raw tensor to a list\n",
    "> keeps the whole computation graph alive and quietly eats your memory."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0020",
   "metadata": {},
   "source": [
    "## Reading the loss curve\n",
    "\n",
    "The loss-per-epoch plot is your training EKG. Learn to read these shapes —\n",
    "they'll tell you what to fix (plotting a *validation* loss alongside the\n",
    "training loss, which we'll set up properly in the DataLoader lesson, makes the\n",
    "diagnosis even sharper):\n",
    "\n",
    "- **Smooth decline that flattens out at a low value** — healthy. Training\n",
    "  converged.\n",
    "- **Still clearly falling when the loop ends** — underfitting by impatience:\n",
    "  train longer, or raise the learning rate a bit.\n",
    "- **Flattens early at a high value** — underfitting by capacity: the model is\n",
    "  too simple for the pattern, or the learning rate is too small to make\n",
    "  progress.\n",
    "- **Training loss keeps dropping while validation loss turns around and\n",
    "  rises** — overfitting: the model has started memorizing the training set.\n",
    "  More data, regularization (dropout, coming soon), or early stopping.\n",
    "- **Spiky, oscillating, or growing** — learning rate too high. The optimizer\n",
    "  is overshooting the valley. Cut `lr` by 10x and try again."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0021",
   "metadata": {},
   "source": [
    "## Sensible defaults, in one place\n",
    "\n",
    "When in doubt, start here and adjust only when the loss curve tells you to:\n",
    "\n",
    "| Task | Model output | Criterion | Optimizer |\n",
    "|---|---|---|---|\n",
    "| Regression | 1 linear value per target | `nn.MSELoss` | `Adam, lr=0.001` |\n",
    "| Binary classification | 1 raw logit | `nn.BCEWithLogitsLoss` | `Adam, lr=0.001` |\n",
    "| Multiclass classification | one raw logit per class | `nn.CrossEntropyLoss` | `Adam, lr=0.001` |\n",
    "\n",
    "Two pairings to burn in: `CrossEntropyLoss` wants raw logits **and** integer\n",
    "class labels (`torch.long`, not one-hot). `MSELoss` wants the target shaped\n",
    "exactly like the output — a `[100]` target against a `[100, 1]` output will\n",
    "\"work\" via broadcasting and silently ruin your training."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0022",
   "metadata": {},
   "source": [
    "### 🏋️ Exercise — Break it, then fix it\n",
    "\n",
    "loss={loss:10.4f}  w={w:6.2f}  b={b:6.2f}\")\n",
    "\n",
    "# sgd lr=0.1  diverges (nan) - steps overshoot on this steep loss surface\n",
    "# sgd lr=0.001 converges slowly, b still far from 5\n",
    "# adam lr=0.1 lands near w=-3, b=5\n",
    "`}\n",
    ">\n",
    "Generate a new dataset `y = -3x + 5` plus Gaussian noise (100 points, x in 0–10)\n",
    "and train `nn.Linear(1, 1)` on it three times: (1) `SGD` with `lr=0.1`,\n",
    "(2) `SGD` with `lr=0.001`, (3) `Adam` with `lr=0.1` — 200 epochs each. Print the\n",
    "final loss and the learned `w`, `b` for each run. Which run diverges, which\n",
    "crawls, and which nails it? Relate each to the loss-curve shapes above."
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0023",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your solution here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0024",
   "metadata": {},
   "source": [
    "<details><summary>✅ Show solution</summary>\n",
    "\n",
    "```python\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "\n",
    "torch.manual_seed(0)\n",
    "X = torch.rand(100, 1) * 10\n",
    "y = -3 * X + 5 + torch.randn(100, 1) * 0.8\n",
    "\n",
    "def train(opt_name, lr, epochs=200):\n",
    "    torch.manual_seed(0)\n",
    "    model = nn.Linear(1, 1)\n",
    "    criterion = nn.MSELoss()\n",
    "    opt_cls = optim.SGD if opt_name == \"sgd\" else optim.Adam\n",
    "    optimizer = opt_cls(model.parameters(), lr=lr)\n",
    "    for _ in range(epochs):\n",
    "        loss = criterion(model(X), y)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    return loss.item(), model.weight.item(), model.bias.item()\n",
    "\n",
    "for name, lr in [(\"sgd\", 0.1), (\"sgd\", 0.001), (\"adam\", 0.1)]:\n",
    "    loss, w, b = train(name, lr)\n",
    "    print(f\"{name:4s} lr={lr:<6} -> loss={loss:10.4f}  w={w:6.2f}  b={b:6.2f}\")\n",
    "\n",
    "# sgd lr=0.1  diverges (nan) - steps overshoot on this steep loss surface\n",
    "# sgd lr=0.001 converges slowly, b still far from 5\n",
    "# adam lr=0.1 lands near w=-3, b=5\n",
    "```\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0025",
   "metadata": {},
   "source": [
    "Next up: `nn.Module` — the class-based way to define networks, plus just enough\n",
    "object-oriented Python to read any PyTorch model ever written."
   ]
  }
 ]
}