{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0631b8ca",
   "metadata": {},
   "source": [
    "# Build a Residual CNN\n",
    "\n",
    "*(Click the button in the top right corner to download the lab.)*\n",
    "\n",
    "Deep convolutional networks are more expressive than shallow networks, but they are also more difficult to optimize. This lab explores several architectural mechanisms that help make deeper convolutional networks easier to train. You will first examine batch normalization, then implement residual blocks and combine them into a complete residual CNN with global average pooling. Finally, you will compare this model with a plain CNN that uses neither batch normalization nor residual shortcuts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "895f3ed6",
   "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 torch import nn\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "from torchvision.datasets import FashionMNIST\n",
    "import torchvision.transforms.v2 as v2\n",
    "\n",
    "def set_seed(seed: int) -> None:\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    if torch.cuda.is_available():\n",
    "        torch.cuda.manual_seed_all(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d13c0a9",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "SEED = 7\n",
    "set_seed(SEED)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "print(\"PyTorch version:\", torch.__version__)\n",
    "print(\"Selected device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e35be69d",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 1. Experimental Setup\n",
    "\n",
    "This lab uses FashionMNIST to study architectural changes in a small image-classification problem. A fixed subset of the official training data is used throughout the lab so that the convolutional models can be trained repeatedly without making the experiments unnecessarily expensive.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "356280d4",
   "metadata": {},
   "source": [
    "### Data partitions\n",
    "\n",
    "The official FashionMNIST training set contains the examples available for model development. From it, 12,000 examples are used for training and 6,000 for validation; the remaining training examples are left unused. The split is generated reproducibly. No fitted preprocessing is required. The official test set is loaded but no test loader is created yet, so it remains untouched until the final evaluation.\n",
    "\n",
    "The class counts printed below provide a quick check that the reduced training subset has not introduced a severe class imbalance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcc3b40a",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "DATA_DIRECTORY = \".data\"\n",
    "TRAIN_SIZE = 12_000\n",
    "VALIDATION_SIZE = 6_000\n",
    "\n",
    "transform = v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])\n",
    "\n",
    "train_source = FashionMNIST(DATA_DIRECTORY, True,  transform, download=True)\n",
    "test_dataset = FashionMNIST(DATA_DIRECTORY, False, transform, download=True)\n",
    "\n",
    "class_names = train_source.classes\n",
    "\n",
    "split_generator = torch.Generator().manual_seed(SEED)\n",
    "LEFTOVER = len(train_source) - TRAIN_SIZE - VALIDATION_SIZE\n",
    "train_dataset, validation_dataset, _ = random_split(train_source, [TRAIN_SIZE, VALIDATION_SIZE, LEFTOVER], split_generator)\n",
    "\n",
    "print(\"Training examples:   \", len(train_dataset))\n",
    "print(\"Validation examples: \", len(validation_dataset))\n",
    "print(\"Sealed test examples:\", len(test_dataset))\n",
    "\n",
    "print(\"\\nClass distribution in training set:\")\n",
    "count = torch.bincount(torch.tensor([train_dataset[i][1] for i in range(len(train_dataset))]))\n",
    "for i, class_name in enumerate(class_names):\n",
    "    print(f\"  {class_name}: {count[i]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8229af50",
   "metadata": {},
   "source": [
    "The data loaders are constructed with a default batch size of 128. For a shuffled loader, the random seed is fixed so that the batch order can be reproduced within the lab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a87b2fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 128\n",
    "\n",
    "def make_loader(dataset, shuffle: bool, seed = SEED, batch_size = BATCH_SIZE):\n",
    "    generator = torch.Generator().manual_seed(seed) if shuffle else None\n",
    "    return DataLoader(dataset, batch_size, shuffle, generator=generator)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcf8eb5e",
   "metadata": {},
   "source": [
    "### Training workflow\n",
    "\n",
    "The functions below provide the training and evaluation workflow used throughout the lab. \n",
    "- `train_epoch` performs one parameter-updating pass through the training loader and accumulates epoch-level loss and accuracy. \n",
    "- `evaluate` measures the model at a fixed parameter state using evaluation mode and inference mode.\n",
    "- `fit_model` combines these operations across epochs. After every training epoch, it evaluates the model on the validation set, records the history, and saves the parameters that achieve the lowest validation loss. At the end of training, that best validation checkpoint is restored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab444fba",
   "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",
    "\n",
    "        if not torch.isfinite(loss):\n",
    "            raise FloatingPointError(\"Training loss became non-finite.\")\n",
    "\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": "eb304ef8",
   "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": "a6636532",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def fit_model(model, train_ds, val_ds, epochs, learning_rate, seed=SEED):\n",
    "    set_seed(seed)\n",
    "\n",
    "    train_loader = make_loader(train_ds, shuffle=True, seed=seed)\n",
    "    valid_loader = make_loader(val_ds, shuffle=False, seed=seed)\n",
    "\n",
    "    model = model.to(device)\n",
    "    loss_fn = nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), learning_rate)\n",
    "\n",
    "    history = {\n",
    "        \"train_loss\": [],\n",
    "        \"train_accuracy\": [],\n",
    "        \"valid_loss\": [],\n",
    "        \"valid_accuracy\": [],\n",
    "    }\n",
    "\n",
    "    best_valid_loss = float(\"inf\")\n",
    "    best_valid_accuracy = None\n",
    "    best_epoch = None\n",
    "    best_state = None\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_valid_accuracy = valid_metrics[\"accuracy\"]\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 {\n",
    "        \"model\": model,\n",
    "        \"history\": history,\n",
    "        \"best_epoch\": best_epoch,\n",
    "        \"valid_loss\": best_valid_loss,\n",
    "        \"valid_accuracy\": best_valid_accuracy,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6b720f0",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def plot_metric(results, metric, title):\n",
    "    plt.figure(figsize=(9, 5))\n",
    "\n",
    "    for label, result in results.items():\n",
    "        values = result[\"history\"][metric]\n",
    "        epochs = np.arange(1, len(values) + 1)\n",
    "        plt.plot(epochs, values, marker=\"o\", label=label)\n",
    "\n",
    "    plt.xlabel(\"Epoch\")\n",
    "    plt.ylabel(metric.replace(\"_\", \" \").title())\n",
    "    plt.title(title)\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b874fc8",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## 2. Batch Normalization\n",
    "\n",
    "Batch normalization is a layer that rescales activations using statistics computed for each channel separately. \n",
    "\n",
    "For an activation tensor with shape $(N,C,H,W)$, `BatchNorm2d` calculates one mean and one variance for each channel using all examples in the batch and all spatial positions. An activation $x$ is first normalized as\n",
    "\n",
    "$$\n",
    "\\hat{x}=\\frac{x-\\mu}{\\sqrt{\\sigma^2+\\epsilon}},\n",
    "$$\n",
    "\n",
    "where $\\mu$ and $\\sigma^2$ are the statistics for its channel. The normalized value is then transformed using two learned parameters,\n",
    "\n",
    "$$\n",
    "y=\\gamma\\hat{x}+\\beta,\n",
    "$$\n",
    "\n",
    "where $\\gamma$ controls the scale and $\\beta$ controls the offset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fee45d9",
   "metadata": {},
   "source": [
    "### Calculate channel statistics\n",
    "\n",
    "We start with a small three-channel tensor whose channels deliberately have different offsets and scales. This makes the effect of channel-wise normalization easy to observe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa38f999",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_seed(SEED)\n",
    "\n",
    "activations = torch.randn(4, 3, 5, 5)\n",
    "activations[:, 0] = activations[:, 0] * 0.5 + 3.0\n",
    "activations[:, 1] = activations[:, 1] * 2.0 - 1.0\n",
    "activations[:, 2] = activations[:, 2] * 5.0 + 10.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "404b0326",
   "metadata": {},
   "source": [
    "### <span style=\"background: #3B82F62E; border-left: 5px solid #3b82f6; padding: 4px 8px; border-radius: 4px;\">Exercise 1</span>\n",
    "\n",
    "Calculate the mean and variance of each channel across the batch and spatial dimensions. When computing the variance, use `unbiased=False` so that the calculation matches the batch-normalization formula.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a5863bd",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "# TODO: reduce over batch, height, and width while preserving channels\n",
    "channel_mean = None # YOUR CODE HERE\n",
    "channel_variance = None # YOUR CODE HERE\n",
    "\n",
    "print(\"Means:\", channel_mean)\n",
    "print(\"Variances:\", channel_variance)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "215d8149",
   "metadata": {},
   "source": [
    "### <span style=\"background: #22C55E24; border-left: 5px solid #22c55e; padding: 4px 8px; border-radius: 4px;\">Checkpoint</span>\n",
    "\n",
    "Verify that your calculations are correct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c93e956",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "# TODO: assert that channel_mean and channel_variance each contain 3 values\n",
    "# TODO: compare channel_mean[0] with a direct mean of activations[:, 0, :, :]\n",
    "# TODO: compare channel_variance[0] with a direct variance of activations[:, 0, :, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bc12ec5",
   "metadata": {},
   "source": [
    "### Check with PyTorch\n",
    "\n",
    "Now compare the manual calculation with PyTorch's `BatchNorm2d`. Configure the layer without learned scale and shift parameters, and use the statistics from the current batch. This isolates the normalization step that you calculated above."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7274b78f",
   "metadata": {},
   "source": [
    "### <span style=\"background: #3B82F62E; border-left: 5px solid #3b82f6; padding: 4px 8px; border-radius: 4px;\">Exercise 2</span>\n",
    "\n",
    "Using the channel means and variances calculated previously, normalize the activation tensor manually. Then compare your result with the output produced by `BatchNorm2d`.\n",
    "\n",
    "*Hint:* You can use broadcasting to subtract the channel means and divide by the channel standard deviations. Use `eps` to avoid division by zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c65b6bb4",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "bn_without_affine = nn.BatchNorm2d(3, affine=False, track_running_stats=False)\n",
    "bn_without_affine.train()\n",
    "\n",
    "# TODO: calculate the manual normalization using channel_mean and channel_variance\n",
    "eps = bn_without_affine.eps\n",
    "manual_normalized = None # YOUR CODE HERE\n",
    "\n",
    "pytorch_normalized = bn_without_affine(activations)\n",
    "\n",
    "print(\"Maximum absolute difference:\", (manual_normalized - pytorch_normalized).abs().max().item())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4482867",
   "metadata": {},
   "source": [
    "Before the learned scale and shift are applied, normalization should produce values with mean close to zero and variance close to one in each channel. An ordinary `BatchNorm2d` layer can then rescale and recenter these normalized activations through its learned `weight` and `bias` parameters.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe7721ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "normalized_means = pytorch_normalized.mean(dim=(0, 2, 3))\n",
    "normalized_variances = pytorch_normalized.var(dim=(0, 2, 3), unbiased=False)\n",
    "\n",
    "print(\"Normalized means:\", torch.round(normalized_means, decimals=2))\n",
    "print(\"Normalized variances:\", torch.round(normalized_variances, decimals=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d62966a",
   "metadata": {},
   "source": [
    "### Learned parameters and running statistics\n",
    "\n",
    "An ordinary `BatchNorm2d` layer contains two different kinds of state. The learned `weight` and `bias` are trainable parameters, one pair for each channel. The layer also stores running estimates of the channel mean and variance. These running statistics are not optimized by gradient descent; they are updated from training batches and later used during evaluation.\n",
    "\n",
    "Inspect the initial state of a three-channel batch-normalization layer.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72228e79",
   "metadata": {},
   "outputs": [],
   "source": [
    "bn = nn.BatchNorm2d(3)\n",
    "\n",
    "print(\"Learned scale shape:\", tuple(bn.weight.shape))\n",
    "print(\"Learned shift shape:\", tuple(bn.bias.shape))\n",
    "print(\"Running mean:\", bn.running_mean)\n",
    "print(\"Running variance:\", bn.running_var)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9083507a",
   "metadata": {},
   "source": [
    "### Training and evaluation behaviour\n",
    "\n",
    "In training mode, `BatchNorm2d` normalizes activations using the statistics of the current mini-batch and updates its running estimates. In evaluation mode, it instead normalizes activations using the stored running estimates and leaves them unchanged.\n",
    "\n",
    "Run the following code and compare the running mean and variance before training, after several training batches, and after an evaluation batch.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "198cbd56",
   "metadata": {},
   "outputs": [],
   "source": [
    "bn.train()\n",
    "\n",
    "running_mean_before = bn.running_mean.clone()\n",
    "running_var_before = bn.running_var.clone()\n",
    "\n",
    "for _ in range(5):\n",
    "    training_batch = torch.randn(8, 3, 5, 5) * 2 + 4\n",
    "    _ = bn(training_batch)\n",
    "\n",
    "running_mean_after_training = bn.running_mean.clone()\n",
    "running_var_after_training = bn.running_var.clone()\n",
    "\n",
    "bn.eval()\n",
    "\n",
    "_ = bn(torch.randn(8, 3, 5, 5) * 10 - 20)\n",
    "\n",
    "running_mean_after_evaluation = bn.running_mean.clone()\n",
    "running_var_after_evaluation = bn.running_var.clone()\n",
    "\n",
    "print(\"Running mean before training:    \", running_mean_before)\n",
    "print(\"Running mean after training:     \", running_mean_after_training)\n",
    "print(\"Running mean after evaluation:   \", running_mean_after_evaluation)\n",
    "print()\n",
    "print(\"Running variance before training: \", running_var_before)\n",
    "print(\"Running variance after training:  \", running_var_after_training)\n",
    "print(\"Running variance after evaluation:\", running_var_after_evaluation)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4605af2",
   "metadata": {},
   "source": [
    "### <span style=\"background: #22C55E24; border-left: 5px solid #22c55e; padding: 4px 8px; border-radius: 4px;\">Checkpoint</span>\n",
    "\n",
    "Verify that the running mean and variance changed during training but remained unchanged during evaluation. Also verify that `bn.weight` and `bn.bias` are trainable parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d961d0b",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "# TODO: assert that running statistics changed during training\n",
    "# TODO: assert that they remained unchanged during evaluation\n",
    "# TODO: assert that bn.weight and bn.bias require gradients"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "785fb959",
   "metadata": {},
   "source": [
    "The distinction between training and evaluation modes has an important consequence. During training, the normalized representation of an example can depend on the other examples in the same mini-batch because the normalization statistics come from that batch. During evaluation, the stored running statistics provide a fixed reference instead.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d902a0b3",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 3. Residual Block\n",
    "\n",
    "A residual block combines a learned transformation with a shortcut path. If the residual branch computes a transformation $F(x)$ and the input already has the correct shape, the block can use an identity shortcut:\n",
    "\n",
    "$$\n",
    "y = x + F(x).\n",
    "$$\n",
    "\n",
    "If the residual branch changes the number of channels or the spatial resolution, the shortcut must transform the input as well:\n",
    "\n",
    "$$\n",
    "y = P(x) + F(x),\n",
    "$$\n",
    "\n",
    "where $P(x)$ is a projection that makes the shortcut compatible with the residual output.\n",
    "\n",
    "### Reusable implementation\n",
    "\n",
    "The residual branch below follows the sequence **Conv → BatchNorm → ReLU → Conv → BatchNorm**. The shortcut uses `nn.Identity` when the shapes already match. Otherwise, a `1×1` convolution with the required stride changes the channel count and spatial dimensions, followed by `BatchNorm2d`. The two branches are then added and passed through a final ReLU.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c8a5913",
   "metadata": {},
   "source": [
    "### <span style=\"background: #3B82F62E; border-left: 5px solid #3b82f6; padding: 4px 8px; border-radius: 4px;\">Exercise 3</span>\n",
    "\n",
    "Complete the `ResidualBlock` class. Use an identity shortcut when the input and residual output shapes match. Otherwise, construct the projection shortcut with a `1×1` convolution using the block stride, followed by `BatchNorm2d`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d5666f4",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "class ResidualBlock(nn.Module):\n",
    "\n",
    "    def __init__(self, in_channels, out_channels, stride=1):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(out_channels)\n",
    "        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False)\n",
    "        self.bn2 = nn.BatchNorm2d(out_channels)\n",
    "\n",
    "        # TODO: use nn.Identity when shapes already match.\n",
    "        #       Otherwise use a 1x1 projection with the required stride,\n",
    "        #       followed by BatchNorm2d.\n",
    "        if in_channels == out_channels and stride == 1:\n",
    "            self.shortcut = None # YOUR CODE HERE\n",
    "        else:\n",
    "            self.shortcut = nn.Sequential(\n",
    "                # YOUR CODE HERE\n",
    "            )\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        residual = self._residual_branch(x)\n",
    "        shortcut = self.shortcut(x)\n",
    "\n",
    "        # TODO: combine the two paths and apply ReLU\n",
    "        output = None # YOUR CODE HERE\n",
    "        return output\n",
    "\n",
    "    def _residual_branch(self, x):\n",
    "        residual = self.conv1(x)\n",
    "        residual = self.bn1(residual)\n",
    "        residual = torch.relu(residual)\n",
    "        residual = self.conv2(residual)\n",
    "        residual = self.bn2(residual)\n",
    "        return residual"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "567de212",
   "metadata": {},
   "source": [
    "### <span style=\"background: #3B82F62E; border-left: 5px solid #3b82f6; padding: 4px 8px; border-radius: 4px;\">Exercise 4</span>\n",
    "\n",
    "Indicate the shapes of the residual output and shortcut output for both blocks in the table below.\n",
    "\n",
    "| Block | Input | Residual output | Shortcut output |\n",
    "|---|---|---|---|\n",
    "| `ResidualBlock(32, 32)` | `(8, 32, 28, 28)` | | |\n",
    "| `ResidualBlock(32, 64, stride=2)` | `(8, 32, 28, 28)` | | |\n",
    "\n",
    "Then create the two blocks and inspect the output of the residual and shortcut branches directly.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e11642f",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "identity_block = ResidualBlock(32, 32)\n",
    "projection_block = ResidualBlock(32, 64, stride=2)\n",
    "\n",
    "x = torch.randn(8, 32, 28, 28)\n",
    "\n",
    "# TODO: compute the outputs of both branches for each block\n",
    "identity_residual = None # YOUR CODE HERE\n",
    "identity_shortcut = None # YOUR CODE HERE\n",
    "projection_residual = None # YOUR CODE HERE\n",
    "projection_shortcut = None # YOUR CODE HERE\n",
    "\n",
    "print(\"Identity block residual: \", tuple(identity_residual.shape))\n",
    "print(\"Identity block shortcut: \", tuple(identity_shortcut.shape))\n",
    "print(\"Projection block residual:\", tuple(projection_residual.shape))\n",
    "print(\"Projection block shortcut:\", tuple(projection_shortcut.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8cbafc3",
   "metadata": {},
   "source": [
    "### <span style=\"background: #22C55E24; border-left: 5px solid #22c55e; padding: 4px 8px; border-radius: 4px;\">Checkpoint</span>\n",
    "\n",
    "Verify that the residual and shortcut branches produce compatible shapes in both cases, and confirm that the projection shortcut has the expected `1×1` convolution.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6199cd7c",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "# TODO: assert that residual and shortcut shapes match for both blocks\n",
    "# TODO: assert that identity_block.shortcut is nn.Identity\n",
    "# TODO: verify that the projection convolution has weight shape (64, 32, 1, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f18225f",
   "metadata": {},
   "source": [
    "\n",
    "---\n",
    "\n",
    "## 4. Residual CNN\n",
    "\n",
    "The residual blocks can now be combined into a complete convolutional network. Starting from a FashionMNIST image with shape `(1, 28, 28)`, the model follows this progression:\n",
    "\n",
    "```\n",
    "input    →  stem      →  block 1   →  block 2   →  block 3  →  GAP  →  classifier\n",
    "-------     --------     --------     --------     -------     ---     ----------\n",
    "1×28×28  →  32×28×28  →  32×28×28  →  64×14×14  →  128×7×7  →  128  →  10\n",
    "```\n",
    "\n",
    "The stem first produces 32 feature channels. An identity residual block keeps the same shape, then two projection blocks reduce the spatial resolution while increasing the channel count. Global average pooling converts the final feature maps into one value per channel before the linear classifier produces ten logits."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf481e7f",
   "metadata": {},
   "source": [
    "### Global average pooling\n",
    "\n",
    "Global average pooling converts each feature map into a single value by averaging over all of its spatial positions. For an activation tensor with shape $(N,C,H,W)$, it produces a tensor with shape $(N,C,1,1)$.\n",
    "\n",
    "Each output value summarizes how strongly one learned feature is present across the entire feature map, without preserving its exact spatial location. The resulting vector contains one value per channel and can be passed directly to a classifier.\n",
    "\n",
    "In this network, global average pooling reduces a representation of shape `(128, 7, 7)` to 128 values. Flattening the same representation would produce $128 \\times 7 \\times 7 = 6272$ values, requiring a much larger final linear layer. Global average pooling therefore provides a compact transition from convolutional features to class logits while deliberately discarding exact spatial position."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c6b3c06",
   "metadata": {},
   "source": [
    "### <span style=\"background: #3B82F62E; border-left: 5px solid #3b82f6; padding: 4px 8px; border-radius: 4px;\">Exercise 5</span>\n",
    "\n",
    "Calculate the number of parameters required by a ten-class linear classifier in two cases: after flattening the `(128, 7, 7)` feature maps, and after global average pooling. Include both weights and biases, but do not instantiate the classifier layers.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b7087ad",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "flattened_features = 128 * 7 * 7\n",
    "global_features = 128\n",
    "n_classes = 10\n",
    "\n",
    "# TODO: count parameters include both weights and biases\n",
    "flatten_classifier_parameters = None # YOUR CODE HERE\n",
    "global_classifier_parameters = None # YOUR CODE HERE\n",
    "\n",
    "print(\"Flatten classifier parameters:\", flatten_classifier_parameters)\n",
    "print(\"Global-pooling classifier parameters:\", global_classifier_parameters)\n",
    "print(\"Reduction factor:\", flatten_classifier_parameters / global_classifier_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "384a97a2",
   "metadata": {},
   "source": [
    "### <span style=\"background: #3B82F62E; border-left: 5px solid #3b82f6; padding: 4px 8px; border-radius: 4px;\">Exercise 6</span>\n",
    "\n",
    "Complete the feature extractor in `ResidualCNN` by assembling the three residual blocks shown in the architecture above.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe3e3742",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "class ResidualCNN(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self.stem = nn.Sequential(\n",
    "            nn.Conv2d(1, 32, kernel_size=3, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "\n",
    "        # TODO: assemble the three residual blocks\n",
    "        self.features = nn.Sequential(\n",
    "            # YOUR CODE HERE\n",
    "        )\n",
    "\n",
    "        self.pool = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        self.classifier = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, images):\n",
    "        x = self.stem(images)\n",
    "        x = self.features(x)\n",
    "        x = self.pool(x)\n",
    "        x = torch.flatten(x, start_dim=1)\n",
    "        return self.classifier(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a95ff515",
   "metadata": {},
   "source": [
    "### <span style=\"background: #22C55E24; border-left: 5px solid #22c55e; padding: 4px 8px; border-radius: 4px;\">Checkpoint</span>\n",
    "\n",
    "Complete the expected shape at each stage in the table below.\n",
    "\n",
    "| Stage | Expected shape for batch size 16 |\n",
    "|---|---|\n",
    "| Input | `(16, 1, 28, 28)` |\n",
    "| Stem | |\n",
    "| Residual block 1 | |\n",
    "| Residual block 2 | |\n",
    "| Residual block 3 | |\n",
    "| Global average pooling | |\n",
    "| Flatten | |\n",
    "| Logits | |\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88005bd7",
   "metadata": {},
   "source": [
    "Then, run the following cell and compare your trace with the actual tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78fa8450",
   "metadata": {
    "tags": [
     "hide-inputs"
    ]
   },
   "outputs": [],
   "source": [
    "set_seed(SEED)\n",
    "residual_cnn = ResidualCNN().to(device)\n",
    "example_images = torch.randn(16, 1, 28, 28, device=device)\n",
    "\n",
    "residual_cnn.eval()\n",
    "with torch.inference_mode():\n",
    "    stem_output = residual_cnn.stem(example_images)\n",
    "    block1_output = residual_cnn.features[0](stem_output)\n",
    "    block2_output = residual_cnn.features[1](block1_output)\n",
    "    block3_output = residual_cnn.features[2](block2_output)\n",
    "    pooled_output = residual_cnn.pool(block3_output)\n",
    "    flat_output = torch.flatten(pooled_output, start_dim=1)\n",
    "    logits = residual_cnn.classifier(flat_output)\n",
    "\n",
    "for name, tensor in [\n",
    "    (\"input\", example_images),\n",
    "    (\"stem\", stem_output),\n",
    "    (\"block 1\", block1_output),\n",
    "    (\"block 2\", block2_output),\n",
    "    (\"block 3\", block3_output),\n",
    "    (\"pooled\", pooled_output),\n",
    "    (\"flattened\", flat_output),\n",
    "    (\"logits\", logits),\n",
    "]:\n",
    "    print(f\"{name:10s}: {tuple(tensor.shape)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de5ccd36",
   "metadata": {},
   "source": [
    "### Train the residual CNN\n",
    "\n",
    "The model now has a valid end-to-end architecture. Train a fresh `ResidualCNN` for 6 epochs with a learning rate of 0.001. This run establishes the residual model that will later be compared with a plain CNN under the same training budget.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7717d6ba",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "set_seed(SEED)\n",
    "EPOCHS = 6\n",
    "LEARNING_RATE = 1e-3\n",
    "\n",
    "# TODO: create and train a fresh ResidualCNN\n",
    "model = ResidualCNN()\n",
    "residual_result = None # YOUR CODE HERE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9aca1dff",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "epochs = np.arange(1, EPOCHS + 1)\n",
    "plt.plot(epochs, residual_result[\"history\"][\"train_loss\"], marker=\"o\", label=\"Train Loss\")\n",
    "plt.plot(epochs, residual_result[\"history\"][\"valid_loss\"], marker=\"o\", label=\"Validation Loss\")\n",
    "plt.xlabel(\"Epoch\")\n",
    "plt.ylabel(\"Loss\")\n",
    "plt.title(\"Residual CNN\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc3ff0bc",
   "metadata": {},
   "source": [
    "### <span style=\"background: #8B5CF624; border-left: 5px solid #8b5cf6; padding: 4px 8px; border-radius: 4px;\">Analysis</span>\n",
    "\n",
    "Inspect the training and validation loss curves. Does the training loss decrease consistently? At which epoch is the validation loss lowest? Do the curves show an obvious optimization failure, or does the model appear to learn the task over this training run?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f54f182f",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 5. Plain CNN\n",
    "\n",
    "The residual CNN combines batch normalization and residual connections. To provide a simpler reference, the model below mirrors the sequence of `3×3` convolutions in the residual branches, together with the same channel progression, downsampling locations, global average pooling, and final classifier. It omits both batch normalization and the shortcut paths, including their `1×1` projection convolutions. Each `3×3` convolution is followed directly by ReLU.\n",
    "\n",
    "This is a **broad architectural comparison** between a plain CNN and the complete residual CNN. Because batch normalization and residual shortcuts change together, the comparison cannot isolate the individual contribution of either mechanism.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "442e24b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PlainCNN(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self.stem = nn.Sequential(\n",
    "            nn.Conv2d(1, 32, kernel_size=3, padding=1),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "\n",
    "        self.features = nn.Sequential(\n",
    "            nn.Conv2d(32, 32, kernel_size=3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(32, 32, kernel_size=3, padding=1),\n",
    "            nn.ReLU(),\n",
    "\n",
    "            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(64, 64, kernel_size=3, padding=1),\n",
    "            nn.ReLU(),\n",
    "\n",
    "            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(128, 128, kernel_size=3, padding=1),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "\n",
    "        self.pool = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        self.classifier = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, images):\n",
    "        x = self.stem(images)\n",
    "        x = self.features(x)\n",
    "        x = self.pool(x)\n",
    "        x = torch.flatten(x, start_dim=1)\n",
    "        return self.classifier(x)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64721ab0",
   "metadata": {},
   "source": [
    "### Train the plain CNN\n",
    "\n",
    "Train the plain CNN for the same 6 epochs and with the same learning rate used for the residual CNN. This makes the resulting learning curves directly comparable under these conditions.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfc44db7",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "set_seed(SEED)\n",
    "EPOCHS = 6\n",
    "LEARNING_RATE = 1e-3\n",
    "\n",
    "plain_model = PlainCNN()\n",
    "\n",
    "# TODO: train the plain CNN\n",
    "plain_result = None # YOUR CODE HERE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b18fbfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "comparison_results = {\n",
    "    \"Plain CNN\": plain_result,\n",
    "    \"Residual CNN\": residual_result,\n",
    "}\n",
    "\n",
    "plot_metric(comparison_results, \"train_loss\", \"Plain CNN versus residual CNN\")\n",
    "plot_metric(comparison_results, \"valid_loss\", \"Plain CNN versus residual CNN\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b729358",
   "metadata": {},
   "source": [
    "### <span style=\"background: #8B5CF624; border-left: 5px solid #8b5cf6; padding: 4px 8px; border-radius: 4px;\">Analysis</span>\n",
    "\n",
    "Compare the training-loss curves of the two models. Which model reduces training loss more quickly, and which reaches the lower training loss by the end of the run? Then compare their best validation losses and the relationship between training and validation loss.\n",
    "\n",
    "Use these observations to describe how the two architectures behaved under the tested training conditions. What can this comparison establish, and what can it **not** establish because batch normalization and residual shortcuts changed together?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8428d728",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 6. Test Evaluation\n",
    "\n",
    "Validation loss has been used throughout model development to compare the two candidates and choose among their checkpoints. The test set serves a different purpose: it estimates the generalization performance of the model selected without test feedback.\n",
    "\n",
    "Select the candidate with the lowest validation loss first. Only after that choice has been made should a test loader be created and the selected model evaluated once on the test set.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5d1c660",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_candidates = {\n",
    "    \"Plain CNN\": plain_result,\n",
    "    \"Residual CNN\": residual_result,\n",
    "}\n",
    "\n",
    "candidate_table = pd.DataFrame(\n",
    "    [\n",
    "        {\n",
    "            \"model\": label,\n",
    "            \"best epoch\": result[\"best_epoch\"],\n",
    "            \"validation loss\": result[\"valid_loss\"],\n",
    "            \"validation accuracy\": result[\"valid_accuracy\"],\n",
    "        }\n",
    "        for label, result in final_candidates.items()\n",
    "    ]\n",
    ").sort_values(\"validation loss\")\n",
    "\n",
    "candidate_table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1232cfe",
   "metadata": {},
   "source": [
    "### <span style=\"background: #3B82F62E; border-left: 5px solid #3b82f6; padding: 4px 8px; border-radius: 4px;\">Exercise 7</span>\n",
    "\n",
    "Select the lowest-validation-loss candidate programmatically, then create the test loader and evaluate the selected checkpoint once.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "490ff303",
   "metadata": {
    "tags": [
     "quiz"
    ]
   },
   "outputs": [],
   "source": [
    "selected_label, selected_result = min(\n",
    "    final_candidates.items(),\n",
    "    key=lambda item: item[1][\"valid_loss\"],\n",
    ")\n",
    "\n",
    "print(\"Selected model:\", selected_label)\n",
    "print(\"Validation loss:\", selected_result[\"valid_loss\"])\n",
    "print(\"Validation accuracy:\", selected_result[\"valid_accuracy\"])\n",
    "\n",
    "# TODO: create a non-shuffled test loader only now\n",
    "# TODO: evaluate selected_result[\"model\"] once on the test set\n",
    "test_loader = None # YOUR CODE HERE\n",
    "final_test_metrics = None # YOUR CODE HERE\n",
    "\n",
    "print(\"Test loss:\", final_test_metrics[\"loss\"])\n",
    "print(\"Test accuracy:\", final_test_metrics[\"accuracy\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfa54aa7",
   "metadata": {},
   "source": [
    "### <span style=\"background: #8B5CF624; border-left: 5px solid #8b5cf6; padding: 4px 8px; border-radius: 4px;\">Analysis</span>\n",
    "\n",
    "Compare the final test loss and accuracy with the selected model's validation performance. Are they broadly consistent? Use the test result to describe how well the selected model generalizes to examples that were not used for training or model selection.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e2b0904",
   "metadata": {},
   "source": [
    "The completed model combines several ideas that support deeper convolutional architectures. Batch normalization controls activation scale during training, residual shortcuts preserve a direct path for existing representations, projection shortcuts reconcile incompatible shapes, and global average pooling converts the final feature maps into a compact representation for classification. The comparison with the plain CNN shows how the complete residual architecture behaves under the same training budget, without attributing any difference to a single mechanism.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6652e361",
   "metadata": {},
   "source": [
    "## <span style=\"background: #8B5CF624; border-left: 5px solid #8b5cf6; padding: 4px 8px; border-radius: 4px;\">Final Reflection</span>\n",
    "\n",
    "1. How do the learned parameters and running statistics of `BatchNorm2d` play different roles during training and evaluation?\n",
    "2. When can a residual block use an identity shortcut, and when is a projection shortcut required?\n",
    "3. What does global average pooling preserve, and what spatial information does it discard?\n",
    "4. What evidence from the learning curves distinguishes the residual CNN from the plain CNN in this experiment, and what causal conclusion can you **not** draw from that comparison?\n",
    "5. How does the selected model's test performance compare with its validation performance, and what does that suggest about its generalization?\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "deep-learning-book (3.14.5.final.0)",
   "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
}