{
 "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": [
    "# Building a CNN in PyTorch\n",
    "\n",
    "Train a complete convolutional network on FashionMNIST — transforms, conv-relu-pool blocks with shape bookkeeping, data augmentation, and reading the mistakes.\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/cnn-pytorch).*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0001",
   "metadata": {},
   "source": [
    "You know what a convolution computes; now let's assemble the real thing. In\n",
    "this lesson you'll build the full image-classification pipeline: load a\n",
    "dataset with `torchvision`, stack conv-relu-pool blocks into a CNN, train it\n",
    "with the recipe from the last module, measure whether data augmentation\n",
    "helps, and — most instructive of all — look at the images it gets wrong."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0002",
   "metadata": {},
   "source": [
    "> **This lesson needs a GPU**\n",
    "> \n",
    "> Run the notebook in Google Colab and enable the free GPU first:\n",
    "> **Runtime → Change runtime type → T4 GPU**. Training takes a couple of minutes\n",
    "> on GPU versus ~20x longer on CPU. Verify with `torch.cuda.is_available()` —\n",
    "> it should print `True`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0003",
   "metadata": {},
   "source": [
    "## Loading images with torchvision\n",
    "\n",
    "`torchvision` bundles standard datasets and image transforms. We'll use\n",
    "**FashionMNIST**: 70,000 grayscale 28×28 images of clothing in 10 classes —\n",
    "big enough to be interesting, small enough to train in minutes:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0004",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn, optim\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(device)\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),                      # PIL image -> [1, 28, 28] float in [0, 1]\n",
    "    transforms.Normalize((0.286,), (0.353,)),   # (x - mean) / std, per channel\n",
    "])\n",
    "\n",
    "train_set = datasets.FashionMNIST(\"data\", train=True, download=True, transform=transform)\n",
    "test_set = datasets.FashionMNIST(\"data\", train=False, download=True, transform=transform)\n",
    "\n",
    "trainloader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=2)\n",
    "testloader = DataLoader(test_set, batch_size=256, shuffle=False)\n",
    "\n",
    "classes = train_set.classes\n",
    "print(classes)          # ['T-shirt/top', 'Trouser', 'Pullover', ...]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0005",
   "metadata": {},
   "source": [
    "Two transforms, two jobs: `ToTensor` converts the image to a channels-first\n",
    "float tensor scaled to 0–1, and `Normalize` standardizes it with the dataset's\n",
    "mean and standard deviation — the same \"scale your features\" habit from\n",
    "tabular data, applied per pixel. Always look at your data before training:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0006",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "images, labels = next(iter(trainloader))\n",
    "print(images.shape)     # [128, 1, 28, 28] -> [batch, channels, height, width]\n",
    "\n",
    "fig, axes = plt.subplots(3, 6, figsize=(12, 6))\n",
    "for img, label, ax in zip(images, labels, axes.flatten()):\n",
    "    ax.imshow(img.squeeze(), cmap=\"gray\")\n",
    "    ax.set_title(classes[label])\n",
    "    ax.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0007",
   "metadata": {},
   "source": [
    "## The model: conv blocks plus a linear head\n",
    "\n",
    "A classic small CNN is two parts. The **convolutional body** extracts features\n",
    "while shrinking the spatial size; the **linear head** flattens whatever's left\n",
    "and classifies it. The comments track the tensor shape at every stage —\n",
    "do this in every CNN you write, using the size formula from last lesson:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0008",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "class CNN(nn.Module):\n",
    "    def __init__(self, n_classes=10):\n",
    "        super().__init__()\n",
    "        self.conv = nn.Sequential(\n",
    "            # in: [1, 28, 28]\n",
    "            nn.Conv2d(1, 32, kernel_size=3, padding=1),   # -> [32, 28, 28]\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2),                              # -> [32, 14, 14]\n",
    "\n",
    "            nn.Conv2d(32, 64, kernel_size=3, padding=1),  # -> [64, 14, 14]\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2),                              # -> [64, 7, 7]\n",
    "\n",
    "            nn.Flatten(),                                 # -> [64 * 7 * 7] = [3136]\n",
    "        )\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(64 * 7 * 7, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.2),\n",
    "            nn.Linear(128, n_classes),                    # raw logits\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.fc(self.conv(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0009",
   "metadata": {},
   "source": [
    "The bookkeeping matters because `nn.Flatten` feeds `nn.Linear(64 * 7 * 7, ...)`\n",
    "— get the arithmetic wrong and you'll meet PyTorch's most famous error, a\n",
    "matrix-shape mismatch on the first forward pass. The pattern to remember:\n",
    "**channels grow (1 → 32 → 64) while spatial size shrinks (28 → 14 → 7)**. The\n",
    "network trades *where* for *what*."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0010",
   "metadata": {},
   "source": [
    "## Training: same recipe, new data\n",
    "\n",
    "Nothing about the training loop changes — that's the payoff of the reusable\n",
    "pattern from the DataLoader lesson. Here it is, compact:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0011",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def run_epoch(model, loader, criterion, optimizer=None):\n",
    "    training = optimizer is not None\n",
    "    model.train() if training else model.eval()\n",
    "    total_loss, correct, n = 0.0, 0, 0\n",
    "    with torch.enable_grad() if training else torch.inference_mode():\n",
    "        for xb, yb in loader:\n",
    "            xb, yb = xb.to(device), yb.to(device)\n",
    "            output = model(xb)\n",
    "            loss = criterion(output, yb)\n",
    "            if training:\n",
    "                optimizer.zero_grad()\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "            total_loss += loss.item() * len(xb)\n",
    "            correct += (output.argmax(1) == yb).sum().item()\n",
    "            n += len(xb)\n",
    "    return total_loss / n, correct / n\n",
    "\n",
    "model = CNN().to(device)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "for epoch in range(1, 6):\n",
    "    train_loss, train_acc = run_epoch(model, trainloader, criterion, optimizer)\n",
    "    test_loss, test_acc = run_epoch(model, testloader, criterion)\n",
    "    print(f\"epoch {epoch} | train acc {train_acc:.3f} | test acc {test_acc:.3f}\")\n",
    "\n",
    "torch.save(model.state_dict(), \"cnn_baseline.pth\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0012",
   "metadata": {},
   "source": [
    "Five epochs should land around **91–92% test accuracy** — already far beyond\n",
    "what an MLP on flattened pixels manages with similar effort. Notice the gap\n",
    "between train and test accuracy creeping open by the last epoch: mild\n",
    "overfitting, our cue for the next section."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0013",
   "metadata": {},
   "source": [
    "## Data augmentation: free training data\n",
    "\n",
    "**Data augmentation** applies random, label-preserving distortions to each\n",
    "training image, every time it's loaded — flips, small rotations, crops. The\n",
    "model never sees the exact same pixels twice, so memorizing becomes much\n",
    "harder; effectively you've multiplied your dataset. The transforms go in the\n",
    "*training* transform only — never distort the test set, it's the measuring\n",
    "stick:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0014",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "train_transform = transforms.Compose([\n",
    "    transforms.RandomHorizontalFlip(),       # a mirrored sneaker is still a sneaker\n",
    "    transforms.RandomRotation(10),           # up to +/-10 degrees\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.286,), (0.353,)),\n",
    "])\n",
    "\n",
    "aug_train_set = datasets.FashionMNIST(\"data\", train=True, transform=train_transform)\n",
    "aug_trainloader = DataLoader(aug_train_set, batch_size=128, shuffle=True, num_workers=2)\n",
    "\n",
    "torch.manual_seed(0)\n",
    "model_aug = CNN().to(device)\n",
    "optimizer = optim.Adam(model_aug.parameters(), lr=0.001)\n",
    "\n",
    "for epoch in range(1, 6):\n",
    "    train_loss, train_acc = run_epoch(model_aug, aug_trainloader, criterion, optimizer)\n",
    "    test_loss, test_acc = run_epoch(model_aug, testloader, criterion)\n",
    "    print(f\"epoch {epoch} | train acc {train_acc:.3f} | test acc {test_acc:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0015",
   "metadata": {},
   "source": [
    "Compare the two runs. The augmented model's **training** accuracy is *lower*\n",
    "(the task got harder — every image is warped), but the train/test gap nearly\n",
    "closes, and with more epochs the augmented model overtakes the baseline on\n",
    "test accuracy. That's the signature of regularization working. Augmentation\n",
    "pays off most when data is scarce or training is long; pick transforms that\n",
    "match reality — horizontal flips make sense for clothing, but would be a\n",
    "terrible idea for digit recognition, where a flipped \"3\" is no longer a 3."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0016",
   "metadata": {},
   "source": [
    "## Read the mistakes\n",
    "\n",
    "A single accuracy number hides *what* the model gets wrong. Plotting\n",
    "misclassified images is the fastest way to build intuition:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0017",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "model_aug.eval()\n",
    "wrong_imgs, wrong_true, wrong_pred = [], [], []\n",
    "\n",
    "with torch.inference_mode():\n",
    "    for xb, yb in testloader:\n",
    "        xb, yb = xb.to(device), yb.to(device)\n",
    "        preds = model_aug(xb).argmax(1)\n",
    "        mask = preds != yb\n",
    "        wrong_imgs.append(xb[mask].cpu())\n",
    "        wrong_true.append(yb[mask].cpu())\n",
    "        wrong_pred.append(preds[mask].cpu())\n",
    "\n",
    "wrong_imgs = torch.cat(wrong_imgs)\n",
    "wrong_true = torch.cat(wrong_true)\n",
    "wrong_pred = torch.cat(wrong_pred)\n",
    "print(f\"{len(wrong_imgs)} mistakes out of {len(test_set)}\")\n",
    "\n",
    "fig, axes = plt.subplots(3, 6, figsize=(14, 7))\n",
    "for img, t, p, ax in zip(wrong_imgs, wrong_true, wrong_pred, axes.flatten()):\n",
    "    ax.imshow(img.squeeze(), cmap=\"gray\")\n",
    "    ax.set_title(f\"true: {classes[t]}\\npred: {classes[p]}\", color=\"red\", fontsize=9)\n",
    "    ax.axis(\"off\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0018",
   "metadata": {},
   "source": [
    "You'll find the errors are *systematic*, not random: shirts confused with\n",
    "coats, pullovers with shirts — pairs even humans squint at in 28×28\n",
    "grayscale. That tells you the remaining errors need better inputs or bigger\n",
    "models, not more epochs. And \"bigger models\" has a shortcut: instead of\n",
    "training a deeper CNN from scratch, borrow one already trained on millions of\n",
    "images."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0019",
   "metadata": {},
   "source": [
    "### 🏋️ Exercise — Go deeper: a three-block CNN\n",
    "\n",
    "Add a **third** conv block (`64 → 128` channels) to the CNN and retrain with\n",
    "the augmented loader for 5 epochs. Before running, compute by hand what\n",
    "spatial size reaches the flatten layer and fix the head's input size to match.\n",
    "Then compare against the two-block model: test accuracy *and* total parameter\n",
    "count. Which model is bigger — and is the answer what you expected?"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0020",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your solution here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0021",
   "metadata": {},
   "source": [
    "<details><summary>✅ Show solution</summary>\n",
    "\n",
    "```python\n",
    "class CNN3(nn.Module):\n",
    "    def __init__(self, n_classes=10):\n",
    "        super().__init__()\n",
    "        self.conv = nn.Sequential(\n",
    "            nn.Conv2d(1, 32, 3, padding=1),    # [32, 28, 28]\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2),                   # [32, 14, 14]\n",
    "\n",
    "            nn.Conv2d(32, 64, 3, padding=1),   # [64, 14, 14]\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2),                   # [64, 7, 7]\n",
    "\n",
    "            nn.Conv2d(64, 128, 3, padding=1),  # [128, 7, 7]\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2),                   # [128, 3, 3]  (floor(7/2)=3)\n",
    "\n",
    "            nn.Flatten(),                      # [1152]\n",
    "        )\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(128 * 3 * 3, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.2),\n",
    "            nn.Linear(128, n_classes),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.fc(self.conv(x))\n",
    "\n",
    "torch.manual_seed(0)\n",
    "model3 = CNN3().to(device)\n",
    "optimizer = optim.Adam(model3.parameters(), lr=0.001)\n",
    "\n",
    "for epoch in range(1, 6):\n",
    "    tr_loss, tr_acc = run_epoch(model3, aug_trainloader, criterion, optimizer)\n",
    "    te_loss, te_acc = run_epoch(model3, testloader, criterion)\n",
    "    print(f\"epoch {epoch} | train acc {tr_acc:.3f} | test acc {te_acc:.3f}\")\n",
    "\n",
    "n2 = sum(p.numel() for p in CNN().parameters())\n",
    "n3 = sum(p.numel() for p in model3.parameters())\n",
    "print(f\"2-block params: {n2:,} | 3-block params: {n3:,}\")\n",
    "# The 3-block model usually edges ahead on test accuracy — with FEWER\n",
    "# parameters, because the flatten vector shrank from 3136 to 1152.\n",
    "```\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0022",
   "metadata": {},
   "source": [
    "Next up: the Transfer Learning module — where we stop training CNNs from\n",
    "scratch and fine-tune networks pretrained on millions of images instead."
   ]
  }
 ]
}