{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a911a1e0",
   "metadata": {},
   "source": [
    "# Build a CNN\n",
    "\n",
    "*(Click the button in the top right corner to download the lab.)*\n",
    "\n",
    "A multilayer perceptron can classify images after flattening them into vectors, but it has no built-in notion that nearby pixels form local patterns or that the same pattern may appear in many positions. A convolutional network introduces exactly those assumptions. Before training a CNN, we will first show the effects of convolution on an ordinary photograph. The objective is to connect the operation of `Conv2d` to a visible image effect, then use the same operation as a learned component of a classifier.\n",
    "\n",
    "\n",
    "*NOTE: The following tutorials are useful when you need API details.*\n",
    "\n",
    "| Tutorial | Use in this lab |\n",
    "|---|---|\n",
    "| CNN — Network Architecture | Review of the overall CNN idea |\n",
    "| MLP — Training a Neural Network | Review of the training cycle |\n",
    "| MLP — Evaluating a Neural Network | Review of evaluation loop |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d41bbcd",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import copy\n",
    "import random\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "from PIL import Image\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "from torchvision import datasets\n",
    "import torchvision.transforms.v2 as v2 \n",
    "\n",
    "SEED = 7\n",
    "\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(\"Selected device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46412df0",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## 1. Convolution\n",
    "\n",
    "A convolutional layer is easiest to understand when its output can be inspected visually. We will begin with a real photograph bundled with Matplotlib. The photograph is not part of dataset used later for classification. It is simply an image with enough structure (object boundaries, clothing, facial features, and background changes) to make local edge detection visible."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "618feea6",
   "metadata": {},
   "source": [
    "### Load an image\n",
    "\n",
    "We load a grayscale image with PIL, then apply a transformation to convert it to a float tensor with values in the range [0, 1]. After this transformation, the tensor has shape `(channels, height, width)`. We then add a batch dimension to the tensor so that it can be passed to a convolutional layer. Run the cell below and inspect the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b53fd279",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import matplotlib.cbook as cbook\n",
    "\n",
    "transform = v2.Compose([\n",
    "    v2.ToImage(),\n",
    "    v2.ToDtype(torch.float32, scale=True)\n",
    "])\n",
    "\n",
    "photo_path = cbook.get_sample_data(\"grace_hopper.jpg\", asfileobj=False)\n",
    "photo = Image.open(photo_path).convert(\"L\")\n",
    "\n",
    "photo_tensor = transform(photo).unsqueeze(0) # Add batch dimension\n",
    "\n",
    "plt.figure(figsize=(6, 6))\n",
    "plt.imshow(photo_tensor[0, 0], cmap=\"gray\")\n",
    "plt.axis(\"off\")\n",
    "plt.colorbar()\n",
    "plt.title(f\"Photo tensor shape: {tuple(photo_tensor.shape)}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e85d75e",
   "metadata": {},
   "source": [
    "### <span style=\"background: #22C55E24; border-left: 5px solid #22c55e; padding: 4px 8px; border-radius: 4px;\">Checkpoint</span>\n",
    "\n",
    "The following properties are assumptions used by every convolution in this section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00400c8e",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "# TODO: assert that photo_tensor has shape (1, 1, 600, 512)\n",
    "# TODO: assert that photo_tensor is floating point\n",
    "# TODO: assert that all pixel values lie between 0 and 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "322a0690",
   "metadata": {},
   "source": [
    "### Edge detectors\n",
    "\n",
    "A learned convolutional layer eventually discovers its own kernels. For the moment, we will set the weights ourselves so that we know what the layer is looking for. The two kernels below are Sobel-style edge detectors. One responds strongly to changes from left to right; the other responds strongly to changes from top to bottom."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06b01cb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "vertical_kernel = torch.tensor(\n",
    "    [[-1.0, 0.0, 1.0],\n",
    "     [-2.0, 0.0, 2.0],\n",
    "     [-1.0, 0.0, 1.0]]\n",
    ")\n",
    "\n",
    "horizontal_kernel = torch.tensor(\n",
    "    [[-1.0, -2.0, -1.0],\n",
    "     [ 0.0,  0.0,  0.0],\n",
    "     [ 1.0,  2.0,  1.0]]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9396fe15",
   "metadata": {},
   "source": [
    "### <span style=\"background: #3B82F62E; border-left: 5px solid #3b82f6; padding: 4px 8px; border-radius: 4px;\">Exercise 1</span>\n",
    "\n",
    "A `Conv2d` layer with one input channel and two output channels contains two complete kernels. You will copy the vertical detector into the first kernel and the horizontal detector into the second. Remember that the `weight` attribute of a convolutional layer is a tensor of shape: `(out_channels, in_channels, kernel_height, kernel_width)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "080115fd",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "edge_detector = nn.Conv2d(\n",
    "    in_channels=1,\n",
    "    out_channels=2,\n",
    "    kernel_size=3,\n",
    "    bias=False,\n",
    ")\n",
    "\n",
    "with torch.no_grad():\n",
    "    # TODO: copy vertical_kernel into the first output kernel\n",
    "    # TODO: copy horizontal_kernel into the second output kernel\n",
    "    pass\n",
    "\n",
    "with torch.inference_mode():\n",
    "    edge_maps = edge_detector(photo_tensor)\n",
    "\n",
    "print(\"Edge-map shape:\", edge_maps.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f12429a2",
   "metadata": {},
   "source": [
    "The convolution produces **signed** responses. Reversing an intensity transition reverses the sign of the response. For visualization, the magnitude is often more convenient because it shows where the detector responds strongly regardless of orientation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fed8665",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "vertical_edges = edge_maps[0, 0].abs()\n",
    "horizontal_edges = edge_maps[0, 1].abs()\n",
    "\n",
    "figure, axes = plt.subplots(1, 3, figsize=(10, 5))\n",
    "\n",
    "axes[0].imshow(photo_tensor[0, 0], cmap=\"gray\")\n",
    "axes[0].set_title(\"Original\")\n",
    "\n",
    "axes[1].imshow(vertical_edges, cmap=\"gray\")\n",
    "axes[1].set_title(\"Vertical-edge response\")\n",
    "\n",
    "axes[2].imshow(horizontal_edges, cmap=\"gray\")\n",
    "axes[2].set_title(\"Horizontal-edge response\")\n",
    "\n",
    "for axis in axes:\n",
    "    axis.axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f202b5f6",
   "metadata": {},
   "source": [
    "### <span style=\"background: #8B5CF624; border-left: 5px solid #8b5cf6; padding: 4px 8px; border-radius: 4px;\">Analysis</span>\n",
    "\n",
    "Inspect several visible boundaries in the photograph. \n",
    "\n",
    "- Which structures produce strong responses in the vertical map?\n",
    "- Which structures appear more strongly in the horizontal map?\n",
    "- Explain why the output is still an image-like grid.\n",
    "- Explain weight sharing and local connectivity."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbc24415",
   "metadata": {},
   "source": [
    "### Convolution settings\n",
    "\n",
    "The kernel determines *what local pattern is measured*. Padding and stride determine *where and how densely the measurement is made*. \n",
    "\n",
    "We will keep the vertical edge detector fixed and change only these settings. This lets us attribute visible differences to the convolution geometry rather than to a different filter."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "814aed7b",
   "metadata": {},
   "source": [
    "### <span style=\"background: #3B82F62E; border-left: 5px solid #3b82f6; padding: 4px 8px; border-radius: 4px;\">Exercise 2 </span>\n",
    "\n",
    "For a `600 × 512` image and a `3 × 3` kernel, predict the spatial output shape for each configuration before running the code.\n",
    "\n",
    "| Padding | Stride | Predicted height × width |\n",
    "|---:|---:|---|\n",
    "| 0 | 1 | |\n",
    "| 1 | 1 | |\n",
    "| 1 | 2 | |\n",
    "\n",
    "Then, complete the function below so that it constructs a convolution, copies the fixed kernel into the layer, and returns the resulting feature map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4868d7ee",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "def apply_fixed_kernel(image, kernel, padding=0, stride=1):\n",
    "\n",
    "    layer = nn.Conv2d(1, 1, 3, stride, padding, bias=False)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        # TODO: copy kernel into layer.weight[0, 0]\n",
    "        layer.weight[0, 0] = None # YOUR CODE HERE\n",
    "        pass\n",
    "\n",
    "    with torch.inference_mode():\n",
    "        # TODO: apply the layer to image\n",
    "        response = None # YOUR CODE HERE\n",
    "\n",
    "    return response\n",
    "\n",
    "\n",
    "#--- Test ---#\n",
    "\n",
    "print(\"Original photo shape:\", tuple(photo_tensor.shape))\n",
    "\n",
    "setting_results = []\n",
    "for padding, stride in [(0, 1), (1, 1), (1, 2)]:\n",
    "    response = apply_fixed_kernel(photo_tensor, vertical_kernel, padding, stride)\n",
    "    setting_results.append((padding, stride, response))\n",
    "    print(f\"padding={padding}, stride={stride} → {tuple(response.shape)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cf6da8d",
   "metadata": {},
   "source": [
    "Run the cell below to see the actual outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f565293",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "figure, axes = plt.subplots(1, 4, figsize=(14, 4))\n",
    "\n",
    "axes[0].imshow(photo_tensor[0, 0], cmap=\"gray\")\n",
    "axes[0].set_title(\"Original\")\n",
    "axes[0].axis(\"off\")\n",
    "\n",
    "for axis, (padding, stride, response) in zip(axes[1:], setting_results):\n",
    "    feature_map = response[0, 0].abs()\n",
    "    height, width = feature_map.shape\n",
    "\n",
    "    axis.imshow(feature_map, cmap=\"gray\", interpolation=\"nearest\")\n",
    "    axis.set_title(f\"p={padding}, s={stride} → {height} × {width}\")\n",
    "    #axis.axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c364e01f",
   "metadata": {},
   "source": [
    "### <span style=\"background: #8B5CF624; border-left: 5px solid #8b5cf6; padding: 4px 8px; border-radius: 4px;\">Analysis</span>\n",
    "\n",
    "Compare the three feature maps. Explain what information is lost at the boundary without padding and why stride two produces a lower-resolution representation. The kernel weights did not change, so what exactly changed about the computation?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3d01d9e",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## 2. Convolutional layer\n",
    "\n",
    "In a CNN, the convolutional kernel values are ordinary model parameters and backpropagation adjusts them to reduce the classification loss. A convolutional layer usually learns several kernels at once. Each kernel produces one output feature map, so the number of kernels is the number of output channels."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c90f6ded",
   "metadata": {},
   "source": [
    "### Several learned filters\n",
    "\n",
    "Consider a layer with one grayscale input channel and eight output channels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae709ee6",
   "metadata": {},
   "outputs": [],
   "source": [
    "learned_conv = nn.Conv2d(\n",
    "    in_channels=1,\n",
    "    out_channels=8,\n",
    "    kernel_size=3,\n",
    "    padding=1,\n",
    ")\n",
    "\n",
    "print(\"Weight shape:\", tuple(learned_conv.weight.shape))\n",
    "print(\"Bias shape:\", tuple(learned_conv.bias.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f2366ab",
   "metadata": {},
   "source": [
    "### Convolution blocks\n",
    "\n",
    "Convolution remains an affine operation, so it is normally followed by a nonlinear activation. Pooling can then reduce the spatial resolution while preserving the number of channels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b3d7404",
   "metadata": {},
   "outputs": [],
   "source": [
    "block = nn.Sequential(\n",
    "    nn.Conv2d(1, 8, kernel_size=3, padding=1),\n",
    "    nn.ReLU(),\n",
    "    nn.MaxPool2d(kernel_size=2),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7ff9706",
   "metadata": {},
   "source": [
    "### <span style=\"background: #3B82F62E; border-left: 5px solid #3b82f6; padding: 4px 8px; border-radius: 4px;\">Exercise 3</span>\n",
    "\n",
    "Starting from `(1, 1, 600, 512)`, predict the shape after the convolution, after ReLU, and after max pooling. Also calculate the number of trainable parameters in each layer. Fill in the table below.\n",
    "\n",
    "| Operation | Predicted shape | Trainable parameters |\n",
    "|---|---|---|\n",
    "| Conv2d | | |\n",
    "| ReLU | | |\n",
    "| MaxPool2d | | |\n",
    "\n",
    "Then, pass the photograph through each operation and record the shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c2beca7",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "activation = photo_tensor\n",
    "block_shapes = []\n",
    "\n",
    "for layer in block:\n",
    "    # TODO: apply layer to activation\n",
    "    activation = None # YOUR CODE HERE\n",
    "    block_shapes.append((layer.__class__.__name__, tuple(activation.shape)))\n",
    "\n",
    "block_shapes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22a4c4f6",
   "metadata": {},
   "source": [
    "### <span style=\"background: #22C55E24; border-left: 5px solid #22c55e; padding: 4px 8px; border-radius: 4px;\">Checkpoint</span>\n",
    "\n",
    "Verify that the parameter counts match your predictions. If they do not, check your calculations and the convolution settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8da0a97",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "# TODO: assert the total number of trainable parameters in Conv2d\n",
    "# TODO: assert MaxPool2d contributes no trainable parameters\n",
    "conv_parameters = sum(parameter.numel() for parameter in block[0].parameters())\n",
    "pool_parameters = sum(parameter.numel() for parameter in block[2].parameters())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "454a4402",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## 3. Build a CNN\n",
    "\n",
    "We have now seen convolution act directly on a photograph. The next step is to let the filters be learned for an actual classification task. FashionMNIST contains `28 × 28` grayscale images from ten clothing categories. The input already has the spatial form expected by a CNN, so we will keep the channel, height, and width dimensions intact until the final classifier."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd289de9",
   "metadata": {},
   "source": [
    "### Prepare the data\n",
    "\n",
    "The original training partition contains 60,000 examples. As in the earlier labs, we will reserve 6,000 for validation and use the official FashionMNIST test set only after model selection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c1c3da6",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_DIRECTORY = \".data\"\n",
    "\n",
    "transform = v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])\n",
    "\n",
    "train_source = datasets.FashionMNIST(root=DATA_DIRECTORY, train=True,  transform=transform, download=True)\n",
    "test_dataset = datasets.FashionMNIST(root=DATA_DIRECTORY, train=False, transform=transform, download=True)\n",
    "\n",
    "split_generator = torch.Generator().manual_seed(SEED)\n",
    "training_ds, validation_ds = random_split(train_source, [54_000, 6_000], split_generator)\n",
    "\n",
    "class_names = train_source.classes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "023bdd3f",
   "metadata": {},
   "source": [
    "The dataloaders are constructed with a batch size of 128. The training dataloader shuffles the data, while the validation and test dataloaders do not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91911490",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 128\n",
    "\n",
    "def make_training_loader(seed=SEED):\n",
    "    return DataLoader(training_ds, BATCH_SIZE, shuffle=True, generator=torch.Generator().manual_seed(seed))\n",
    "\n",
    "def make_unshuffled_loader(ds, seed=SEED):\n",
    "    return DataLoader(ds, BATCH_SIZE, shuffle=False, generator=torch.Generator().manual_seed(seed))\n",
    "\n",
    "\n",
    "training_loader = make_training_loader()\n",
    "validation_loader = make_unshuffled_loader(validation_ds)\n",
    "test_loader = make_unshuffled_loader(test_dataset)\n",
    "\n",
    "images, targets = next(iter(training_loader))\n",
    "print(\"Image batch:\", tuple(images.shape))\n",
    "print(\"Target batch:\", tuple(targets.shape))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8902fd4",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "figure, axes = plt.subplots(2, 5, figsize=(10, 5))\n",
    "\n",
    "for index, axis in enumerate(axes.flat):\n",
    "    axis.imshow(images[index, 0], cmap=\"gray\")\n",
    "    axis.set_title(class_names[targets[index].item()])\n",
    "    axis.axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fd5a481",
   "metadata": {},
   "source": [
    "### Assemble the CNN\n",
    "\n",
    "The model follows the architecture developed in the lesson. Two convolutional blocks construct spatial feature maps, after which the final `7 × 7` maps are flattened and converted to ten logits."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b27b61ab",
   "metadata": {},
   "source": [
    "### <span style=\"background: #3B82F62E; border-left: 5px solid #3b82f6; padding: 4px 8px; border-radius: 4px;\">Exercise 4</span>\n",
    "\n",
    "Complete the `SmallCNN` class below. It should contain two convolutional blocks, each with a convolutional layer, a ReLU activation, and a max pooling layer. The final output should be flattened and passed through a linear layer to produce ten logits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93da1d43",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "class SmallCNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self.features = nn.Sequential(\n",
    "            # TODO: 1 → 16 channels, 3×3 kernel, preserve spatial size\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2),\n",
    "\n",
    "            # TODO: 16 → 32 channels, 3×3 kernel, preserve spatial size\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2),\n",
    "        )\n",
    "\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Flatten(),\n",
    "            # TODO: map the final 32×7×7 representation to 10 logits\n",
    "        )\n",
    "\n",
    "    def forward(self, images):\n",
    "        features = self.features(images)\n",
    "        logits = self.classifier(features)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e1fe260",
   "metadata": {},
   "source": [
    "Before training, trace one batch through the model. Predict each shape first, then run the code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "635d3224",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "torch.manual_seed(SEED)\n",
    "cnn = SmallCNN().to(device)\n",
    "\n",
    "trace = []\n",
    "activation = images[:8].to(device)\n",
    "trace.append((\"input\", tuple(activation.shape)))\n",
    "\n",
    "for layer in cnn.features:\n",
    "    activation = layer(activation)\n",
    "    trace.append((layer.__class__.__name__, tuple(activation.shape)))\n",
    "\n",
    "for layer in cnn.classifier:\n",
    "    activation = layer(activation)\n",
    "    trace.append((layer.__class__.__name__, tuple(activation.shape)))\n",
    "\n",
    "trace"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2577b7e0",
   "metadata": {},
   "source": [
    "### <span style=\"background: #22C55E24; border-left: 5px solid #22c55e; padding: 4px 8px; border-radius: 4px;\">Checkpoint</span>\n",
    "\n",
    "Verify the shape of parameter weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20cf092f",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "# TODO: assert that the first convolution has weight shape (16, 1, 3, 3)\n",
    "# TODO: assert that the second convolution has weight shape (32, 16, 3, 3)\n",
    "# TODO: assert that the final layer has weight shape (10, 1568)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf7d4fa9",
   "metadata": {},
   "source": [
    "### <span style=\"background: #8B5CF624; border-left: 5px solid #8b5cf6; padding: 4px 8px; border-radius: 4px;\">Analysis</span>\n",
    "\n",
    "Identify the point at which the representation stops being spatial. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e680410a",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## 4. Train and Inspect the CNN\n",
    "\n",
    "Nothing special is required to optimize convolutional parameters. They participate in the same differentiable computation as the weights of a linear layer. The training workflow therefore remains the one developed earlier: batches produce losses, backpropagation calculates gradients, the optimizer updates parameters, validation tracks generalization, and the best validation checkpoint is restored."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e86b3260",
   "metadata": {},
   "source": [
    "### Training workflow\n",
    "\n",
    "The functions below are supplied infrastructure. Their purpose here is to keep attention on the CNN rather than on rebuilding the training loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9318e96",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def train_epoch(model, loader, loss_fn, optimizer, device):\n",
    "    model.train()\n",
    "\n",
    "    total_loss = 0.0\n",
    "    total_correct = 0\n",
    "    total_examples = 0\n",
    "\n",
    "    for images, targets in loader:\n",
    "        images = images.to(device)\n",
    "        targets = targets.to(device)\n",
    "\n",
    "        optimizer.zero_grad(set_to_none=True)\n",
    "        logits = model(images)\n",
    "        loss = loss_fn(logits, targets)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        batch_size = targets.shape[0]\n",
    "        total_loss += loss.item() * batch_size\n",
    "        total_correct += (logits.argmax(dim=1) == targets).sum().item()\n",
    "        total_examples += batch_size\n",
    "\n",
    "    return {\n",
    "        \"loss\": total_loss / total_examples,\n",
    "        \"accuracy\": total_correct / total_examples,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53b0a4ef",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def evaluate(model, loader, loss_fn, device):\n",
    "    model.eval()\n",
    "\n",
    "    total_loss = 0.0\n",
    "    total_correct = 0\n",
    "    total_examples = 0\n",
    "\n",
    "    with torch.inference_mode():\n",
    "        for images, targets in loader:\n",
    "            images = images.to(device)\n",
    "            targets = targets.to(device)\n",
    "\n",
    "            logits = model(images)\n",
    "            loss = loss_fn(logits, targets)\n",
    "\n",
    "            batch_size = targets.shape[0]\n",
    "            total_loss += loss.item() * batch_size\n",
    "            total_correct += (logits.argmax(dim=1) == targets).sum().item()\n",
    "            total_examples += batch_size\n",
    "\n",
    "    return {\n",
    "        \"loss\": total_loss / total_examples,\n",
    "        \"accuracy\": total_correct / total_examples,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6b56ed3",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def fit(model, train_loader, valid_loader, loss_fn, optimizer, epochs):\n",
    "\n",
    "    device = next(model.parameters()).device\n",
    "\n",
    "    history = {\n",
    "        \"train_loss\": [],\n",
    "        \"train_accuracy\": [],\n",
    "        \"valid_loss\": [],\n",
    "        \"valid_accuracy\": [],\n",
    "    }\n",
    "\n",
    "    best_state = None\n",
    "    best_epoch = None\n",
    "    best_valid_loss = float(\"inf\")\n",
    "\n",
    "    progress_bar = tqdm(range(1, epochs + 1), unit=\"epoch\", desc=\"Training\")\n",
    "\n",
    "    for epoch in progress_bar:\n",
    "        train_metrics = train_epoch(model, train_loader, loss_fn, optimizer, device)\n",
    "        valid_metrics = evaluate(model, valid_loader, loss_fn, device)\n",
    "\n",
    "        history[\"train_loss\"].append(train_metrics[\"loss\"])\n",
    "        history[\"train_accuracy\"].append(train_metrics[\"accuracy\"])\n",
    "        history[\"valid_loss\"].append(valid_metrics[\"loss\"])\n",
    "        history[\"valid_accuracy\"].append(valid_metrics[\"accuracy\"])\n",
    "\n",
    "        if valid_metrics[\"loss\"] < best_valid_loss:\n",
    "            best_valid_loss = valid_metrics[\"loss\"]\n",
    "            best_epoch = epoch\n",
    "            best_state = copy.deepcopy(model.state_dict())\n",
    "\n",
    "        progress_bar.set_postfix(train_loss=f\"{train_metrics['loss']:.4f}\", valid_loss=f\"{valid_metrics['loss']:.4f}\")\n",
    "\n",
    "    model.load_state_dict(best_state)\n",
    "\n",
    "    return history, {\"epoch\": best_epoch, \"valid_loss\": best_valid_loss}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "939437fa",
   "metadata": {},
   "source": [
    "### <span style=\"background: #3B82F62E; border-left: 5px solid #3b82f6; padding: 4px 8px; border-radius: 4px;\">Exercise 5</span>\n",
    "\n",
    "Train the CNN with cross-entropy loss and Adam with learning rate `1e-3`. Five epochs are enough for this lab: the purpose is to observe the architecture in a complete workflow, not to exhaustively tune FashionMNIST."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02dd9b30",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "torch.manual_seed(SEED)\n",
    "cnn = SmallCNN().to(device)\n",
    "\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(cnn.parameters(), lr=1e-3)\n",
    "\n",
    "EPOCHS = 5\n",
    "\n",
    "cnn_training_loader = make_training_loader(SEED)\n",
    "\n",
    "# TODO: train cnn with fit(...), using cnn_training_loader\n",
    "cnn_history, cnn_checkpoint = None # YOUR CODE HERE\n",
    "\n",
    "print(f\"Selected CNN checkpoint\\n{'\\n'.join(f' - {k}: {v}' for k, v in cnn_checkpoint.items())}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a66aa6e",
   "metadata": {},
   "source": [
    "Plot the resulting histories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5b0257c",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "epochs = np.arange(1, EPOCHS + 1)\n",
    "\n",
    "plt.figure(figsize=(8, 5))\n",
    "plt.plot(epochs, cnn_history[\"train_loss\"], marker=\"o\", label=\"training\")\n",
    "plt.plot(epochs, cnn_history[\"valid_loss\"], marker=\"o\", label=\"validation\")\n",
    "plt.xlabel(\"Epoch\")\n",
    "plt.ylabel(\"Cross-entropy loss\")\n",
    "plt.title(\"CNN loss history\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58c1eeed",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8, 5))\n",
    "plt.plot(epochs, cnn_history[\"train_accuracy\"], marker=\"o\", label=\"training\")\n",
    "plt.plot(epochs, cnn_history[\"valid_accuracy\"], marker=\"o\", label=\"validation\")\n",
    "plt.xlabel(\"Epoch\")\n",
    "plt.ylabel(\"Accuracy\")\n",
    "plt.title(\"CNN accuracy history\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59e8b17c",
   "metadata": {},
   "source": [
    "### <span style=\"background: #8B5CF624; border-left: 5px solid #8b5cf6; padding: 4px 8px; border-radius: 4px;\">Analysis</span>\n",
    "\n",
    "- Did optimization progress? \n",
    "- How strong was the training fit? \n",
    "- Did validation performance improve with training? \n",
    "- Which epoch was retained by the validation-loss checkpoint? "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc627b40",
   "metadata": {},
   "source": [
    "### First convolution kernels\n",
    "\n",
    "At the beginning of the lab, we inserted edge detectors by hand. The first CNN layer has instead learned sixteen `3 × 3` kernels from the classification loss. Let's visualize them after training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b91fa812",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "first_conv = cnn.features[0]\n",
    "learned_kernels = first_conv.weight.detach().cpu()\n",
    "\n",
    "figure, axes = plt.subplots(4, 4, figsize=(7, 7))\n",
    "\n",
    "for index, axis in enumerate(axes.flat):\n",
    "    axis.imshow(learned_kernels[index, 0], cmap=\"gray\")\n",
    "    axis.set_title(f\"kernel {index}\")\n",
    "    axis.axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d1ac03a",
   "metadata": {},
   "source": [
    "Now apply the first convolution and ReLU to one FashionMNIST image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f42ae998",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "example_image = images[0:1].to(device)\n",
    "\n",
    "with torch.inference_mode():\n",
    "    first_feature_maps = torch.relu(first_conv(example_image)).cpu()\n",
    "\n",
    "figure, axes = plt.subplots(4, 4, figsize=(7, 7))\n",
    "for index, axis in enumerate(axes.flat):\n",
    "    axis.imshow(first_feature_maps[0, index], cmap=\"gray\")\n",
    "    axis.set_title(f\"channel {index}\")\n",
    "    axis.axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cfbde7e",
   "metadata": {},
   "source": [
    "### <span style=\"background: #8B5CF624; border-left: 5px solid #8b5cf6; padding: 4px 8px; border-radius: 4px;\">Analysis</span>\n",
    "\n",
    "Explain the difference between the sixteen learned kernels and the sixteen feature maps produced for this particular image. Which one is a learned parameter of the model, and which one depends on the current input? Avoid assigning a precise human meaning to every channel unless the visual evidence supports it."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6eca5c1e",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## 5. Compare with an MLP\n",
    "\n",
    "Let's compare the CNN with a multilayer perceptron (MLP). The CNN and an MLP can use the same loss function, optimizer, batches, validation procedure, and checkpoint rule. The architectural difference is how they represent the image."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c67de16d",
   "metadata": {},
   "source": [
    "### Train the MLP\n",
    "\n",
    "The MLP flattens the image immediately. Its hidden units receive position-specific weights for all 784 pixels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1213d447",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FashionMLP(nn.Module):\n",
    "    def __init__(self, hidden_width=128):\n",
    "        super().__init__()\n",
    "        self.network = nn.Sequential(\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(28 * 28, hidden_width),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_width, 10),\n",
    "        )\n",
    "\n",
    "    def forward(self, images):\n",
    "        return self.network(images)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d96be15",
   "metadata": {},
   "source": [
    "### <span style=\"background: #3B82F62E; border-left: 5px solid #3b82f6; padding: 4px 8px; border-radius: 4px;\">Exercise 6</span>\n",
    "\n",
    "Train the MLP using the same training parameters used for the CNN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a68e6a82",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "torch.manual_seed(SEED)\n",
    "mlp = FashionMLP(hidden_width=128).to(device)\n",
    "\n",
    "mlp_optimizer = torch.optim.Adam(mlp.parameters(), lr=1e-3)\n",
    "\n",
    "mlp_training_loader = make_training_loader(SEED)\n",
    "\n",
    "# TODO: train mlp with the same fit(...) workflow, using mlp_training_loader\n",
    "mlp_history, mlp_checkpoint = None # YOUR CODE HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48244ca3",
   "metadata": {},
   "source": [
    "Evaluate the restored CNN and MLP checkpoints on the validation set and compare their parameter counts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f4dfe15",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "cnn_valid = evaluate(cnn, validation_loader, loss_fn, device)\n",
    "mlp_valid = evaluate(mlp, validation_loader, loss_fn, device)\n",
    "\n",
    "comparison = pd.DataFrame(\n",
    "    [\n",
    "        {\n",
    "            \"model\": \"MLP\",\n",
    "            \"parameters\": sum(p.numel() for p in mlp.parameters()),\n",
    "            \"best epoch\": mlp_checkpoint[\"epoch\"],\n",
    "            \"validation loss\": mlp_valid[\"loss\"],\n",
    "            \"validation accuracy\": mlp_valid[\"accuracy\"],\n",
    "        },\n",
    "        {\n",
    "            \"model\": \"CNN\",\n",
    "            \"parameters\": sum(p.numel() for p in cnn.parameters()),\n",
    "            \"best epoch\": cnn_checkpoint[\"epoch\"],\n",
    "            \"validation loss\": cnn_valid[\"loss\"],\n",
    "            \"validation accuracy\": cnn_valid[\"accuracy\"],\n",
    "        },\n",
    "    ]\n",
    ")\n",
    "\n",
    "comparison"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "befcfa4a",
   "metadata": {},
   "source": [
    "### <span style=\"background: #8B5CF624; border-left: 5px solid #8b5cf6; padding: 4px 8px; border-radius: 4px;\">Analysis</span>\n",
    "\n",
    "Compare the validation results and parameter counts. \n",
    "- Did the CNN provide stronger evidence of generalization under this training protocol? \n",
    "- How does the opening edge-detection experiment help explain why convolution is an appropriate architectural choice for images?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71952e4c",
   "metadata": {},
   "source": [
    "### Evaluate once on test\n",
    "\n",
    "Now, we'll select the best model based on validation loss and evaluate it once on the untouched test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd33a407",
   "metadata": {},
   "outputs": [],
   "source": [
    "candidates = [\n",
    "    (\"MLP\", mlp, mlp_valid),\n",
    "    (\"CNN\", cnn, cnn_valid),\n",
    "]\n",
    "\n",
    "selected_name, selected_model, selected_validation = min(candidates, key=lambda item: item[2][\"loss\"])\n",
    "\n",
    "selected_test = evaluate(selected_model, test_loader, loss_fn, device)\n",
    "\n",
    "print(\"Selected model:\", selected_name)\n",
    "print(f\"Validation accuracy: {selected_validation['accuracy']:.1%}\")\n",
    "print(f\"Test accuracy: {selected_test['accuracy']:.1%}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "deep-learning-book (3.14.5)",
   "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.14.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
