{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7355917d",
   "metadata": {},
   "source": [
    "# Improving deep networks: initialization, regularization, optimizers, schedules and batch norm\n",
    "\n",
    "Real experiments, small enough for a laptop CPU: how initialization changes training, what weight decay and dropout do to the train/dev gap, how SGD, momentum, RMSprop, Adam and AdamW compare, learning-rate schedules, batch normalization, mixed precision and torch.compile.\n",
    "\n",
    "[Open in Google Colab](https://colab.research.google.com/github/mahmoudyosrimahmoud13/inpace-library/blob/main/courses/dl-2-improving-networks/files/dl2_improving_networks_pytorch.ipynb) · part of the inPace course library. Slides © DeepLearning.AI (CC BY-SA 2.0); this notebook: inPace (CC BY-SA 4.0).\n",
    "\n",
    "**Run it top to bottom** (Runtime → Run all in Colab). Each part says what to look for."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "950c92d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Works in Google Colab as it is (PyTorch is already installed there).\n",
    "# On your own computer:  pip install torch torchvision scikit-learn scipy matplotlib pandas\n",
    "# One framework per notebook: PyTorch and TensorFlow can clash when loaded into the same Python process.\n",
    "import numpy as np, matplotlib.pyplot as plt, torch\n",
    "from torch import nn\n",
    "torch.manual_seed(0); np.random.seed(0)\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
    "print(\"PyTorch\", torch.__version__, \"| device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f7221b6",
   "metadata": {},
   "source": [
    "## 0. Data and a helper\n",
    "Handwritten digits (8×8), split into train / dev / test. `run()` trains a model and records the losses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9403398",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.datasets import load_digits\n",
    "from sklearn.model_selection import train_test_split\n",
    "d = load_digits()\n",
    "X, Xt_, y, yt_ = train_test_split(d.data / 16.0, d.target, test_size=0.4, random_state=0)\n",
    "Xd_, Xte, yd_, yte = train_test_split(Xt_, yt_, test_size=0.5, random_state=0)\n",
    "T = lambda a, f=torch.float32: torch.tensor(a, dtype=f)\n",
    "Xtr, ytr, Xdv, ydv, Xts, yts = T(X), T(y, torch.long), T(Xd_), T(yd_, torch.long), T(Xte), T(yte, torch.long)\n",
    "print(\"train\", tuple(Xtr.shape), \"dev\", tuple(Xdv.shape), \"test\", tuple(Xts.shape))\n",
    "loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(Xtr, ytr), batch_size=64, shuffle=True)\n",
    "\n",
    "def accuracy(m, X, y):\n",
    "    m.eval()\n",
    "    with torch.no_grad(): return (m(X).argmax(1) == y).float().mean().item()\n",
    "\n",
    "def run(model, opt, epochs=40, sched=None):\n",
    "    hist = {\"train\": [], \"dev\": []}\n",
    "    for _ in range(epochs):\n",
    "        model.train()\n",
    "        for xb, yb in loader:                      # mini-batch gradient descent\n",
    "            opt.zero_grad(); nn.functional.cross_entropy(model(xb), yb).backward(); opt.step()\n",
    "            if sched: sched.step()\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            hist[\"train\"].append(nn.functional.cross_entropy(model(Xtr), ytr).item())\n",
    "            hist[\"dev\"].append(nn.functional.cross_entropy(model(Xdv), ydv).item())\n",
    "    return hist\n",
    "\n",
    "def mlp(width=256, depth=4, dropout=0.0, bn=False):\n",
    "    layers, n = [], 64\n",
    "    for _ in range(depth):\n",
    "        layers += [nn.Linear(n, width)] + ([nn.BatchNorm1d(width)] if bn else []) + [nn.ReLU(), nn.Dropout(dropout)]\n",
    "        n = width\n",
    "    return nn.Sequential(*layers, nn.Linear(n, 10))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d5e2689",
   "metadata": {},
   "source": [
    "## 1. Initialization\n",
    "PyTorch's `nn.Linear` already uses a sensible default (Kaiming-uniform). Here we force three choices on a 10-layer network and watch the first epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05201a0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_with(model, how):\n",
    "    for m in model.modules():\n",
    "        if isinstance(m, nn.Linear):\n",
    "            if how == \"zeros\": nn.init.zeros_(m.weight)\n",
    "            elif how == \"tiny\": nn.init.normal_(m.weight, std=0.01)\n",
    "            elif how == \"he\": nn.init.kaiming_normal_(m.weight, nonlinearity=\"relu\")\n",
    "            nn.init.zeros_(m.bias)\n",
    "    return model\n",
    "for how in (\"zeros\", \"tiny\", \"he\"):\n",
    "    torch.manual_seed(0)\n",
    "    m = init_with(mlp(width=128, depth=10), how)\n",
    "    h = run(m, torch.optim.SGD(m.parameters(), lr=0.05), epochs=15)\n",
    "    plt.plot(h[\"train\"], label=f\"{how}: dev acc {accuracy(m, Xdv, ydv):.2f}\")\n",
    "plt.legend(); plt.title(\"10 ReLU layers, plain SGD: training loss\"); plt.xlabel(\"epoch\"); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84bce4e4",
   "metadata": {},
   "source": [
    "## 2. Regularization: weight decay and dropout\n",
    "We train on only **150 examples** so the network can memorize them. Look at the gap between train and dev."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae778c5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "small = torch.utils.data.TensorDataset(Xtr[:150], ytr[:150])\n",
    "loader_full, loader = loader, torch.utils.data.DataLoader(small, batch_size=50, shuffle=True)\n",
    "rows = []\n",
    "for name, kw, wd in ((\"nothing\", {}, 0.0), (\"weight decay 1.0\", {}, 1.0), (\"dropout 0.5\", {\"dropout\": 0.5}, 0.0)):\n",
    "    torch.manual_seed(0)\n",
    "    m = mlp(512, 3, **kw)\n",
    "    h = run(m, torch.optim.AdamW(m.parameters(), lr=1e-3, weight_decay=wd), epochs=200)\n",
    "    rows.append((name, accuracy(m, Xtr[:150], ytr[:150]), accuracy(m, Xdv, ydv)))\n",
    "    plt.plot(h[\"dev\"], label=f\"{name} (dev loss)\")\n",
    "loader = loader_full\n",
    "plt.legend(); plt.ylim(0, 1.5); plt.title(\"dev loss with 150 training examples\"); plt.xlabel(\"epoch\"); plt.show()\n",
    "for r in rows: print(f\"{r[0]:<18} train acc {r[1]:.3f} | dev acc {r[2]:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52b76fae",
   "metadata": {},
   "source": [
    "## 3. Optimizers\n",
    "Same network, same mini-batches, five optimizers. `AdamW` is Adam with *decoupled* weight decay: the modern default for most networks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "009740d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "opts = {\"SGD\": lambda p: torch.optim.SGD(p, lr=0.05),\n",
    "        \"SGD + momentum 0.9\": lambda p: torch.optim.SGD(p, lr=0.05, momentum=0.9),\n",
    "        \"RMSprop\": lambda p: torch.optim.RMSprop(p, lr=1e-3),\n",
    "        \"Adam\": lambda p: torch.optim.Adam(p, lr=1e-3),\n",
    "        \"AdamW\": lambda p: torch.optim.AdamW(p, lr=1e-3, weight_decay=0.01)}\n",
    "for name, make in opts.items():\n",
    "    torch.manual_seed(0); m = mlp(256, 4)\n",
    "    h = run(m, make(m.parameters()), epochs=25)\n",
    "    plt.plot(h[\"train\"], label=f\"{name}: test acc {accuracy(m, Xts, yts):.3f}\")\n",
    "plt.yscale(\"log\"); plt.legend(); plt.title(\"training loss\"); plt.xlabel(\"epoch\"); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92af87a1",
   "metadata": {},
   "source": [
    "## 4. Learning-rate schedules\n",
    "A constant learning rate vs **one-cycle** (warm up, then anneal). The scheduler steps after every mini-batch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c10d098",
   "metadata": {},
   "outputs": [],
   "source": [
    "for name in (\"constant 3e-3\", \"one-cycle, max 1e-2\"):\n",
    "    torch.manual_seed(0); m = mlp(256, 4)\n",
    "    opt = torch.optim.AdamW(m.parameters(), lr=3e-3)\n",
    "    sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=1e-2, total_steps=25 * len(loader)) if \"one\" in name else None\n",
    "    h = run(m, opt, epochs=25, sched=sched)\n",
    "    plt.plot(h[\"dev\"], label=f\"{name}: test acc {accuracy(m, Xts, yts):.3f}\")\n",
    "plt.legend(); plt.title(\"dev loss\"); plt.xlabel(\"epoch\"); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bdf812b",
   "metadata": {},
   "source": [
    "## 5. Batch normalization\n",
    "With a learning rate that's too high for plain layers, batch norm keeps each layer's inputs in a steady range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0acb368a",
   "metadata": {},
   "outputs": [],
   "source": [
    "for bn in (False, True):\n",
    "    torch.manual_seed(0); m = mlp(256, 8, bn=bn)\n",
    "    h = run(m, torch.optim.SGD(m.parameters(), lr=0.5, momentum=0.9), epochs=20)\n",
    "    plt.plot(h[\"train\"], label=f\"batch norm {bn}: test acc {accuracy(m, Xts, yts):.3f}\")\n",
    "plt.yscale(\"log\"); plt.legend(); plt.title(\"8 layers, learning rate 0.5\"); plt.xlabel(\"epoch\"); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82072a5a",
   "metadata": {},
   "source": [
    "## 6. Check gradients the PyTorch way\n",
    "`torch.autograd.gradcheck` compares autograd with numerical differences (use float64)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d98d555b",
   "metadata": {},
   "outputs": [],
   "source": [
    "layer = nn.Sequential(nn.Linear(4, 3), nn.Tanh()).double()\n",
    "x = torch.randn(5, 4, dtype=torch.float64, requires_grad=True)\n",
    "print(\"gradcheck passed:\", torch.autograd.gradcheck(lambda t: layer(t).sum(), (x,)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c06b91f",
   "metadata": {},
   "source": [
    "## 7. Faster training on a GPU: mixed precision and torch.compile\n",
    "On a GPU (Colab: Runtime → Change runtime type → T4 GPU), `autocast` runs most maths in 16-bit floats and `torch.compile` fuses operations. On a CPU this cell just runs normally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2100506",
   "metadata": {},
   "outputs": [],
   "source": [
    "m = mlp(256, 4).to(device)\n",
    "opt = torch.optim.AdamW(m.parameters(), lr=1e-3)\n",
    "use_amp = device == \"cuda\"\n",
    "fast = torch.compile(m) if device == \"cuda\" else m      # compiling pays off on GPUs and long runs\n",
    "for xb, yb in loader:\n",
    "    xb, yb = xb.to(device), yb.to(device)\n",
    "    with torch.autocast(device_type=\"cuda\" if use_amp else \"cpu\", dtype=torch.bfloat16, enabled=use_amp):\n",
    "        loss = nn.functional.cross_entropy(fast(xb), yb)\n",
    "    opt.zero_grad(); loss.backward(); opt.step()\n",
    "print(\"one epoch done on\", device, \"| last loss\", round(loss.item(), 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efd23c13",
   "metadata": {},
   "source": [
    "## Try this\n",
    "1. In section 2, try weight decay 0.01, 0.1, 1.0 and 3.0. Plot dev accuracy against weight decay.\n",
    "2. In section 3, give plain SGD a 10× bigger learning rate. Does it catch up, or diverge?\n",
    "3. In section 5, lower the learning rate to 0.05. Does batch norm still matter as much?\n",
    "4. Random search: pick 10 random learning rates on a log scale between 1e-4 and 1e-1 (`10 ** np.random.uniform(-4, -1)`) and keep the best by dev loss."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}