{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "91456bfd",
   "metadata": {},
   "source": [
    "# Exercise: Image Classification with CNNs on Fashion-MNIST\n",
    "\n",
    "## Objective\n",
    "\n",
    "In this exercise, you will train a **Convolutional Neural Network (CNN)** to classify clothing images from the **Fashion-MNIST** dataset.\n",
    "\n",
    "By the end of the exercise, you should be able to:\n",
    "\n",
    "- download and inspect an image dataset using `torchvision`;\n",
    "- understand the shape of image tensors used by CNNs;\n",
    "- create training, validation, and test data loaders;\n",
    "- construct a CNN using convolution, activation, pooling, and fully connected layers;\n",
    "- understand how the spatial dimensions of feature maps change;\n",
    "- train a CNN using forward propagation, loss computation, backpropagation, and an optimizer;\n",
    "- compare training and validation performance;\n",
    "- compute classification accuracy;\n",
    "- identify signs of overfitting and underfitting;\n",
    "- save and reload a trained CNN;\n",
    "- experiment with CNN architecture and hyperparameters.\n",
    "\n",
    "---\n",
    "\n",
    "## Dataset\n",
    "\n",
    "Fashion-MNIST contains **28 × 28 grayscale images** belonging to 10 clothing categories.\n",
    "\n",
    "The classes are:\n",
    "\n",
    "| Label | Class |\n",
    "|---:|---|\n",
    "| 0 | T-shirt/top |\n",
    "| 1 | Trouser |\n",
    "| 2 | Pullover |\n",
    "| 3 | Dress |\n",
    "| 4 | Coat |\n",
    "| 5 | Sandal |\n",
    "| 6 | Shirt |\n",
    "| 7 | Sneaker |\n",
    "| 8 | Bag |\n",
    "| 9 | Ankle boot |\n",
    "\n",
    "Unlike the original MNIST digit dataset, Fashion-MNIST is a slightly more challenging image-classification problem.\n",
    "\n",
    "---\n",
    "\n",
    "## Rules for the exercise\n",
    "\n",
    "Complete every cell marked:\n",
    "\n",
    "```python\n",
    "# TODO\n",
    "```\n",
    "\n",
    "Run the notebook sequentially from top to bottom.\n",
    "\n",
    "Do not immediately copy a complete CNN implementation from the internet. The goal is to understand **what each layer does and why it is needed**.\n",
    "Also jusitfy **the choice of each hyperparameter, choice of loss function and optimizer**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89fb01cc",
   "metadata": {},
   "source": [
    "## 1. Imports and reproducibility"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a52a3ce9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "SEED = 42\n",
    "\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed_all(SEED)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(\"Using device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06b9ad10",
   "metadata": {},
   "source": [
    "## 2. Download Fashion-MNIST\n",
    "\n",
    "`torchvision.datasets.FashionMNIST` can automatically download the dataset. We first convert every image into a PyTorch tensor using:\n",
    "\n",
    "```python \n",
    "transforms.ToTensor()\n",
    "```\n",
    "For a grayscale image, the tensor shape is: $(C,H,W)$\n",
    "\n",
    "where: \n",
    "- \\(C\\) = number of channels,\n",
    "- \\(H\\) = height,\n",
    "- \\(W\\) = width.\n",
    "\n",
    "For Fashion-MNIST: $(1,28,28)$\n",
    "\n",
    "Create the training and test datasets using:\n",
    "```python \n",
    "root=\"./data\"\n",
    "download=True\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d058bda",
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = transforms.ToTensor()\n",
    "train_full = datasets.FashionMNIST(\n",
    "    root=\"./data\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=transform,\n",
    ")\n",
    "test_dataset = datasets.FashionMNIST(\n",
    "    root=\"./data\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=transform,\n",
    ")\n",
    "\n",
    "print(\"Training samples:\", len(train_full))\n",
    "print(\"Test samples:\", len(test_dataset))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "383e65d5",
   "metadata": {},
   "source": [
    "Before designing a neural network, always inspect your data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cb92b04",
   "metadata": {},
   "outputs": [],
   "source": [
    "class_names = [\n",
    "    \"T-shirt/top\",\n",
    "    \"Trouser\",\n",
    "    \"Pullover\",\n",
    "    \"Dress\",\n",
    "    \"Coat\",\n",
    "    \"Sandal\",\n",
    "    \"Shirt\",\n",
    "    \"Sneaker\",\n",
    "    \"Bag\",\n",
    "    \"Ankle boot\",\n",
    "]\n",
    "\n",
    "# TODO: Extract one image and its label.\n",
    "image, label = train_full[0]\n",
    "\n",
    "print(\"Image shape:\", image.shape)\n",
    "print(\"Label:\", label)\n",
    "print(\"Class:\", class_names[label])\n",
    "print(\"Minimum pixel value:\", image.min().item())\n",
    "print(\"Maximum pixel value:\", image.max().item())\n",
    "\n",
    "plt.imshow(image.squeeze(), cmap=\"gray\")\n",
    "plt.title(class_names[label])\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57f3d9f4",
   "metadata": {},
   "source": [
    "## 4. Visualize several images\n",
    "\n",
    "Looking at several examples is useful for understanding the classification problem.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe6564f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(3, 4, figsize=(8, 7))\n",
    "\n",
    "for i, ax in enumerate(axes.flat):\n",
    "    image, label = train_full[i]\n",
    "    ax.imshow(image.squeeze(), cmap=\"gray\")\n",
    "    ax.set_title(class_names[label])\n",
    "    ax.axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a79e4cb",
   "metadata": {},
   "source": [
    "## 5. Create training and validation sets\n",
    "\n",
    "The original Fashion-MNIST training set will be divided into:\n",
    "\n",
    "- **training set**: used to update network weights;\n",
    "- **validation set**: used to monitor generalization during training.\n",
    "\n",
    "The test set must remain untouched until the final evaluation.\n",
    "\n",
    "We will use:\n",
    "\n",
    "- 50,000 training images;\n",
    "- 10,000 validation images;\n",
    "- 10,000 test images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db25ce46",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_size = 50_000\n",
    "val_size = 10_000\n",
    "\n",
    "generator = torch.Generator().manual_seed(SEED)\n",
    "# TODO: Split train_full into training and validation sets.\n",
    "train_dataset, val_dataset = random_split( #TODO))\n",
    "\n",
    "print(\"Training:\", len(train_dataset))\n",
    "print(\"Validation:\", len(val_dataset))\n",
    "print(\"Test:\", len(test_dataset))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d38384a6",
   "metadata": {},
   "source": [
    "## 6. Create DataLoaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "245e35d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64\n",
    "train_loader = DataLoader(\n",
    "    train_dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    ")\n",
    "val_loader = DataLoader(\n",
    "    val_dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=False,\n",
    ")\n",
    "test_loader = DataLoader(\n",
    "    test_dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=False,\n",
    ")\n",
    "images, labels = next(iter(train_loader))\n",
    "\n",
    "print(\"Batch image shape:\", images.shape)\n",
    "print(\"Batch label shape:\", labels.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fda280cd",
   "metadata": {},
   "source": [
    "# Part II — Understanding the CNN\n",
    "\n",
    "## 7. Convolution layer\n",
    "\n",
    "For a convolutional layer:\n",
    "\n",
    "```python\n",
    "nn.Conv2d(\n",
    "    in_channels,\n",
    "    out_channels,\n",
    "    kernel_size,\n",
    "    stride,\n",
    "    padding\n",
    ")\n",
    "```\n",
    "\n",
    "The output spatial dimension is:\n",
    "\n",
    "$H_{\\text{out}}= \\left\\lfloor\\frac{H_{\\text{in}} + 2P - K}{S} \\right\\rfloor + 1$\n",
    "\n",
    "and similarly for the width.\n",
    "\n",
    "where:\n",
    "\n",
    "- \\(K\\) = kernel size;\n",
    "- \\(S\\) = stride;\n",
    "- \\(P\\) = padding.\n",
    "\n",
    "---\n",
    "Consider:\n",
    "\n",
    "```python\n",
    "nn.Conv2d(\n",
    "    in_channels=1,\n",
    "    out_channels=16,\n",
    "    kernel_size=3,\n",
    "    stride=1,\n",
    "    padding=1\n",
    ")\n",
    "```\n",
    "\n",
    "Given an input image of:\n",
    "\n",
    "```text\n",
    "1 × 28 × 28\n",
    "```\n",
    "\n",
    "answer the following **before running any code**:\n",
    "\n",
    "1. How many kernels are learned?\n",
    "2. What is the size of each kernel?\n",
    "3. What will the output tensor shape be for one image?\n",
    "4. Why does the spatial size remain 28 × 28?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4ef15fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_test = nn.Conv2d(\n",
    "    in_channels=1,\n",
    "    out_channels=16,\n",
    "    kernel_size=3,\n",
    "    stride=1,\n",
    "    padding=1,\n",
    ")\n",
    "\n",
    "x = torch.randn(1, 1, 28, 28)\n",
    "y = conv_test(x)\n",
    "\n",
    "print(\"Input shape :\", x.shape)\n",
    "print(\"Output shape:\", y.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00b6583c",
   "metadata": {},
   "source": [
    "## 8. Pooling\n",
    "\n",
    "A max-pooling operation such as:\n",
    "\n",
    "```python\n",
    "nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "```\n",
    "\n",
    "reduces the spatial dimensions.\n",
    "\n",
    "For example:\n",
    "\n",
    "```text\n",
    "16 × 28 × 28\n",
    "```\n",
    "\n",
    "becomes:\n",
    "\n",
    "```text\n",
    "16 × 14 × 14\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdab7317",
   "metadata": {},
   "outputs": [],
   "source": [
    "pool = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "\n",
    "x = torch.randn(1, 16, 28, 28)\n",
    "y = pool(x)\n",
    "\n",
    "print(\"Before pooling:\", x.shape)\n",
    "print(\"After pooling :\", y.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b43c0cd8",
   "metadata": {},
   "source": [
    "## 9. Activation functions\n",
    "\n",
    "A convolution is a linear operation. Without nonlinear activation functions, stacking several linear layers would still behave like one large linear transformation.\n",
    "\n",
    "We will use:\n",
    "\n",
    "```python\n",
    "nn.ReLU()\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97aca08c",
   "metadata": {},
   "source": [
    "## 10. Design the CNN\n",
    "\n",
    "Create your own network with the final layer contains **10 outputs**, one for each Fashion-MNIST class.\n",
    "\n",
    "These outputs are called **logits**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "451ae31e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FashionCNN(nn.Module):\n",
    "    def __init__(self):\n",
    "    ## TO DO create your network\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.features(x)\n",
    "        x = self.classifier(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "model = FashionCNN().to(device)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76371934",
   "metadata": {},
   "source": [
    "## 11. Verify the output shape\n",
    "\n",
    "A very useful debugging technique is to send a dummy batch through the network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8ce04b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "dummy_batch = torch.randn(64, 1, 28, 28).to(device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    dummy_output = model(dummy_batch)\n",
    "\n",
    "print(\"Input shape :\", dummy_batch.shape)\n",
    "print(\"Output shape:\", dummy_output.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a66f0683",
   "metadata": {},
   "source": [
    "# Part III — Training the CNN\n",
    "\n",
    "## 13. Loss function\n",
    "\n",
    "This is a 10-class classification problem.\n",
    "Justify the choice of loss function and try different loss functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fb13367",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = #TODO"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b8c77de",
   "metadata": {},
   "source": [
    "## 14. Optimizer\n",
    "\n",
    "The optimizer updates the network parameters using gradients computed by backpropagation.\n",
    "Justify the choice of optimizer and try different optimizer functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62ddc67e",
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = \n",
    "\n",
    "# TODO: Create the optimizer.\n",
    "optimizer = #TODO\n",
    "\n",
    "print(optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fbea878",
   "metadata": {},
   "source": [
    "## 15. Training loop\n",
    "Train your CNN and then compute the accuracy using:\n",
    "\n",
    "$\\text{Accuracy}=\\frac{\\text{number of correct predictions}}{\\text{number of samples}}\\times 100$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02129d22",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_one_epoch(model, loader, criterion, optimizer, device):\n",
    "    model.train()\n",
    "\n",
    "    running_loss = 0.0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "\n",
    "    for images, labels in loader:\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "\n",
    "        # TODO 1: Reset gradients.\n",
    "        \n",
    "\n",
    "        # TODO 2: Forward pass.\n",
    "        \n",
    "\n",
    "        # TODO 3: Compute loss.\n",
    "        \n",
    "\n",
    "        # TODO 4: Backpropagation.\n",
    "        \n",
    "\n",
    "        # TODO 5: Update network parameters.\n",
    "        \n",
    "\n",
    "        running_loss += loss.item() * images.size(0)\n",
    "\n",
    "        predictions = outputs.argmax(dim=1)\n",
    "        correct += (predictions == labels).sum().item()\n",
    "        total += labels.size(0)\n",
    "\n",
    "    epoch_loss = running_loss / total\n",
    "    epoch_accuracy = 100.0 * correct / total\n",
    "\n",
    "    return epoch_loss, epoch_accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3921d97a",
   "metadata": {},
   "source": [
    "## 16. Validation function\n",
    "\n",
    "During validation:\n",
    "\n",
    "- we do **not** update model parameters;\n",
    "- gradients are unnecessary;\n",
    "- dropout and batch-normalization layers, if present, should behave in evaluation mode.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "567dcff8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model, loader, criterion, device):\n",
    "    model.eval()\n",
    "\n",
    "    running_loss = 0.0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for images, labels in loader:\n",
    "            images = images.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            # TODO: Forward pass.\n",
    "            \n",
    "\n",
    "            # TODO: Compute loss.\n",
    "        \n",
    "\n",
    "            running_loss += loss.item() * images.size(0)\n",
    "\n",
    "            predictions = outputs.argmax(dim=1)\n",
    "            correct += (predictions == labels).sum().item()\n",
    "            total += labels.size(0)\n",
    "\n",
    "    epoch_loss = running_loss / total\n",
    "    epoch_accuracy = 100.0 * correct / total\n",
    "\n",
    "    return epoch_loss, epoch_accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5a18440",
   "metadata": {},
   "source": [
    "## 17. Train the network\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82944b3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = #TODO\n",
    "\n",
    "history = {\"train_loss\": [],\n",
    "    \"val_loss\": [],\n",
    "    \"train_acc\": [],\n",
    "    \"val_acc\": [],}\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    train_loss, train_acc = train_one_epoch(\n",
    "        model,\n",
    "        train_loader,\n",
    "        criterion,\n",
    "        optimizer,\n",
    "        device,\n",
    "    )\n",
    "\n",
    "    val_loss, val_acc = evaluate(\n",
    "        model,\n",
    "        val_loader,\n",
    "        criterion,\n",
    "        device,\n",
    "    )\n",
    "\n",
    "    history[\"train_loss\"].append(train_loss)\n",
    "    history[\"val_loss\"].append(val_loss)\n",
    "    history[\"train_acc\"].append(train_acc)\n",
    "    history[\"val_acc\"].append(val_acc)\n",
    "\n",
    "    print(\n",
    "        f\"Epoch {epoch + 1:2d}/{epochs} | \"\n",
    "        f\"Train Loss: {train_loss:.4f} | \"\n",
    "        f\"Val Loss: {val_loss:.4f} | \"\n",
    "        f\"Train Acc: {train_acc:.2f}% | \"\n",
    "        f\"Val Acc: {val_acc:.2f}%\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab29de4f",
   "metadata": {},
   "source": [
    "## 18. Plot training and validation loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b17a7794",
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs_axis = range(1, epochs + 1)\n",
    "\n",
    "plt.figure(figsize=(7, 5))\n",
    "plt.plot(epochs_axis, history[\"train_loss\"], marker=\"o\", label=\"Training Loss\")\n",
    "plt.plot(epochs_axis, history[\"val_loss\"], marker=\"o\", label=\"Validation Loss\")\n",
    "plt.xlabel(\"Epoch\")\n",
    "plt.ylabel(\"Cross-Entropy Loss\")\n",
    "plt.title(\"Training and Validation Loss\")\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92fa389d",
   "metadata": {},
   "source": [
    "## 19. Plot training and validation accuracy\n",
    "\n",
    "Accuracy is easier to interpret than loss because it directly tells us the percentage of correctly classified images.\n",
    "\n",
    "However, accuracy does **not** replace the loss function during training.\n",
    "\n",
    "### Question:\n",
    "\n",
    "Plot:\n",
    "\n",
    "- training accuracy;\n",
    "- validation accuracy.\n",
    "\n",
    "Then explain why we usually inspect **both loss and accuracy**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66785620",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(7, 5))\n",
    "plt.plot(epochs_axis, history[\"train_acc\"], marker=\"o\", label=\"Training Accuracy\")\n",
    "plt.plot(epochs_axis, history[\"val_acc\"], marker=\"o\", label=\"Validation Accuracy\")\n",
    "plt.xlabel(\"Epoch\")\n",
    "plt.ylabel(\"Accuracy (%)\")\n",
    "plt.title(\"Training and Validation Accuracy\")\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4b4695b",
   "metadata": {},
   "source": [
    "# Part IV — Final evaluation\n",
    "\n",
    "## 20. Test-set performance\n",
    "\n",
    "The test set should be used only after model development is complete.\n",
    "\n",
    "### Question\n",
    "\n",
    "Why would repeatedly checking the test accuracy during model development be a bad practice?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4420e4e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loss, test_accuracy = evaluate(\n",
    "    model,\n",
    "    test_loader,\n",
    "    criterion,\n",
    "    device,\n",
    ")\n",
    "\n",
    "print(f\"Test Loss    : {test_loss:.4f}\")\n",
    "print(f\"Test Accuracy: {test_accuracy:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a44416b3",
   "metadata": {},
   "source": [
    "## 21. Inspect individual predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ef5382c",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "\n",
    "images, labels = next(iter(test_loader))\n",
    "images_device = images.to(device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    outputs = model(images_device)\n",
    "    predictions = outputs.argmax(dim=1).cpu()\n",
    "\n",
    "fig, axes = plt.subplots(3, 4, figsize=(10, 8))\n",
    "\n",
    "for i, ax in enumerate(axes.flat):\n",
    "    ax.imshow(images[i].squeeze(), cmap=\"gray\")\n",
    "\n",
    "    pred_name = class_names[predictions[i].item()]\n",
    "    true_name = class_names[labels[i].item()]\n",
    "\n",
    "    ax.set_title(\n",
    "        f\"Pred: {pred_name}\\nTrue: {true_name}\",\n",
    "        fontsize=9,\n",
    "    )\n",
    "    ax.axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5adbc329",
   "metadata": {},
   "source": [
    "## 22. Per-class accuracy\n",
    "\n",
    "A model may perform well overall but poorly on one particular class.\n",
    "Calculate the classification accuracy separately for each Fashion-MNIST class.\n",
    "### Questions\n",
    "1. Which class is easiest for the CNN?\n",
    "2. Which class is hardest?\n",
    "3. Why might visually similar classes be confused?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94b8a9de",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "\n",
    "class_correct = [0 for _ in range(10)]\n",
    "class_total = [0 for _ in range(10)]\n",
    "\n",
    "with torch.no_grad():\n",
    "    for images, labels in test_loader:\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "\n",
    "        outputs = model(images)\n",
    "        predictions = outputs.argmax(dim=1)\n",
    "\n",
    "        for label, prediction in zip(labels, predictions):\n",
    "            label_index = label.item()\n",
    "            class_total[label_index] += 1\n",
    "\n",
    "            if prediction.item() == label_index:\n",
    "                class_correct[label_index] += 1\n",
    "\n",
    "for i, name in enumerate(class_names):\n",
    "    accuracy = 100.0 * class_correct[i] / class_total[i]\n",
    "    print(f\"{name:15s}: {accuracy:6.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0bdc2e8",
   "metadata": {},
   "source": [
    "# Part V — Understanding what the CNN learned\n",
    "\n",
    "## 23. Feature maps\n",
    "What is a feature map and visualize it\n",
    "Run one test image through the **first convolutional layer** and visualize several feature maps.\n",
    "Do they look like the original image?\n",
    "Check if you can see evidence of:\n",
    "- edges;\n",
    "- textures;\n",
    "- local shapes;\n",
    "- intensity transitions?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c62f217",
   "metadata": {},
   "outputs": [],
   "source": [
    "first_conv = model.features[0]\n",
    "\n",
    "sample_image, sample_label = test_dataset[0]\n",
    "sample_batch = sample_image.unsqueeze(0).to(device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    feature_maps = first_conv(sample_batch).cpu()\n",
    "\n",
    "print(\"Feature-map tensor shape:\", feature_maps.shape)\n",
    "\n",
    "fig, axes = plt.subplots(2, 4, figsize=(10, 5))\n",
    "\n",
    "for i, ax in enumerate(axes.flat):\n",
    "    ax.imshow(feature_maps[0, i], cmap=\"gray\")\n",
    "    ax.set_title(f\"Feature map {i + 1}\")\n",
    "    ax.axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cf92c5d",
   "metadata": {},
   "source": [
    "## 24. What exactly is learned?\n",
    "\n",
    "### Questions\n",
    "\n",
    "Answer in your own words.\n",
    "\n",
    "1. What is a convolutional kernel?\n",
    "2. Are convolutional kernels manually chosen or learned?\n",
    "3. Why can one convolutional layer contain many kernels?\n",
    "4. What does `out_channels=32` mean?\n",
    "5. Why are feature maps different from the original image?\n",
    "6. What happens to feature abstraction as we move deeper into a CNN?\n",
    "7. Why do CNNs usually work better than a fully connected network directly applied to raw images?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9df1eb13",
   "metadata": {},
   "source": [
    "# Part VI — Model diagnosis\n",
    "\n",
    "## 25. Overfitting and underfitting\n",
    "\n",
    "### Case A\n",
    "\n",
    "Suppose after many epochs:\n",
    "\n",
    "```text\n",
    "Training accuracy   = 99%\n",
    "Validation accuracy = 86%\n",
    "```\n",
    "\n",
    "and validation loss starts increasing.\n",
    "\n",
    "**Question:** What is happening?\n",
    "\n",
    "---\n",
    "\n",
    "### Case B\n",
    "\n",
    "Suppose:\n",
    "\n",
    "```text\n",
    "Training accuracy   = 72%\n",
    "Validation accuracy = 70%\n",
    "```\n",
    "\n",
    "and both losses remain relatively high.\n",
    "\n",
    "**Question:** What is happening?\n",
    "\n",
    "---\n",
    "\n",
    "### Case C\n",
    "\n",
    "Suppose:\n",
    "\n",
    "```text\n",
    "Training accuracy   = 92%\n",
    "Validation accuracy = 91%\n",
    "Test accuracy       = 90%\n",
    "```\n",
    "\n",
    "**Question:** What can you conclude about generalization?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e30c5eb",
   "metadata": {},
   "source": [
    "## 26. Regularization experiment\n",
    "\n",
    "One method for reducing overfitting is **dropout**.\n",
    "\n",
    "For example:\n",
    "\n",
    "```python\n",
    "nn.Dropout(p=0.5)\n",
    "```\n",
    "\n",
    "During training, dropout randomly sets some activations to zero.\n",
    "\n",
    "Create a second model where you add:\n",
    "\n",
    "```python\n",
    "nn.Dropout(p=0.5)\n",
    "```\n",
    "\n",
    "between the hidden fully connected layer and the final output layer.\n",
    "\n",
    "Train it and compare:\n",
    "\n",
    "- training accuracy;\n",
    "- validation accuracy;\n",
    "- validation loss.\n",
    "\n",
    "### Question\n",
    "\n",
    "Does dropout always improve performance?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b10c96f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Create a second CNN architecture with dropout.\n",
    "\n",
    "class FashionCNNWithDropout(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self.features = nn.Sequential(\n",
    "            nn.Conv2d(1, 16, kernel_size=3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2),\n",
    "\n",
    "            nn.Conv2d(16, 32, kernel_size=3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2),\n",
    "        )\n",
    "\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(32 * 7 * 7, 128),\n",
    "            nn.ReLU(),\n",
    "\n",
    "            # TODO: Add dropout here.\n",
    "\n",
    "            nn.Linear(128, 10),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.features(x)\n",
    "        x = self.classifier(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "340beddc",
   "metadata": {},
   "source": [
    "# Part VII — Saving and loading the network\n",
    "\n",
    "## 27. Save the learned weights\n",
    "\n",
    "The learned parameters can be stored using:\n",
    "\n",
    "```python\n",
    "torch.save(model.state_dict(), filename)\n",
    "```\n",
    "Save the model as:\n",
    "\n",
    "```text\n",
    "fashion_cnn.pth\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07e11e02",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Save model weights.\n",
    "\n",
    "torch.save(#TODO)\n",
    "\n",
    "print(\"Model weights saved.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ceef1cf",
   "metadata": {},
   "source": [
    "## 28. Reload the model\n",
    "\n",
    "To reload the model:\n",
    "\n",
    "1. create a new instance of the same architecture;\n",
    "2. load the saved `state_dict`;\n",
    "3. move the model to the required device;\n",
    "4. switch to evaluation mode.\n",
    "Reload the saved model and verify its test accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9324b8ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "loaded_model = FashionCNN().to(device)\n",
    "\n",
    "loaded_model.load_state_dict(\n",
    "    torch.load(\n",
    "        \"fashion_cnn.pth\",\n",
    "        map_location=device,\n",
    "    )\n",
    ")\n",
    "\n",
    "loaded_model.eval()\n",
    "\n",
    "loaded_test_loss, loaded_test_accuracy = evaluate(\n",
    "    loaded_model,\n",
    "    test_loader,\n",
    "    criterion,\n",
    "    device,\n",
    ")\n",
    "\n",
    "print(f\"Reloaded model test accuracy: {loaded_test_accuracy:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63d330bd",
   "metadata": {},
   "source": [
    "# Part VIII — CNN design challenge\n",
    "\n",
    "## 29. Architecture experiment\n",
    "\n",
    "Modify **one factor at a time**.\n",
    "\n",
    "Possible experiments:\n",
    "\n",
    "- increase the first convolution layer from 16 to 32 filters;\n",
    "- increase the second convolution layer from 32 to 64 filters;\n",
    "- add a third convolutional layer;\n",
    "- remove one pooling layer;\n",
    "- change kernel size from 3 × 3 to 5 × 5;\n",
    "- replace ReLU with GELU;\n",
    "- add dropout;\n",
    "- change the learning rate;\n",
    "- change the batch size;\n",
    "- increase the number of epochs.\n",
    "\n",
    "For every experiment, record:\n",
    "\n",
    "| Experiment | Parameters changed | Train Acc. | Val. Acc. | Test Acc. |\n",
    "|---|---|---:|---:|---:|\n",
    "| Baseline | — | | | |\n",
    "| Experiment 1 | | | | |\n",
    "| Experiment 2 | | | | |\n",
    "| Experiment 3 | | | | |\n",
    "\n",
    "### Important\n",
    "\n",
    "Change **one main factor at a time**. Otherwise, you will not know which modification caused the performance difference."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a57ac5e",
   "metadata": {},
   "source": [
    "## 30. Final questions\n",
    "\n",
    "Answer the following in your own words.\n",
    "\n",
    "1. Why are convolutional layers particularly suitable for images?\n",
    "2. What is weight sharing in a CNN?\n",
    "3. What is a receptive field?\n",
    "4. What is the role of pooling?\n",
    "5. What is the purpose of flattening?\n",
    "6. What does the final linear layer represent?\n",
    "7. What are logits?\n",
    "8. Why is Softmax not required before `CrossEntropyLoss` in PyTorch?\n",
    "9. What does `loss.backward()` calculate?\n",
    "10. What does `optimizer.step()` do?\n",
    "11. Why do we call `optimizer.zero_grad()`?\n",
    "12. What is the difference between `model.train()` and `model.eval()`?\n",
    "13. Why do we use `torch.no_grad()` during evaluation?\n",
    "14. Why is validation loss important even when training loss is decreasing?\n",
    "15. Why should the test set be used only at the end?\n",
    "16. How would you recognize overfitting from the training curves?\n",
    "17. How would you recognize underfitting?\n",
    "18. Which Fashion-MNIST classes were the most difficult for your model?\n",
    "19. What architectural change gave your best validation accuracy?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5279198",
   "metadata": {},
   "source": [
    "# Optional advanced challenge\n",
    "\n",
    "If you finish early, try one or more of the following.\n",
    "\n",
    "### Challenge A — Normalization\n",
    "\n",
    "Add image normalization using:\n",
    "\n",
    "```python\n",
    "transforms.Normalize((0.5,), (0.5,))\n",
    "```\n",
    "\n",
    "Retrain the model and compare the results.\n",
    "\n",
    "### Challenge B — Batch normalization\n",
    "\n",
    "Add:\n",
    "\n",
    "```python\n",
    "nn.BatchNorm2d(...)\n",
    "```\n",
    "after convolutional layers.\n",
    "\n",
    "### Challenge E — Your own CNN\n",
    "\n",
    "Design a completely new CNN subject to the following rule:\n",
    "\n",
    "> The test set must not be used to choose your architecture.\n",
    "\n",
    "Choose your architecture using the **validation set**, then report the test accuracy only once your design is finalized."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "723322eb",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Submission - import the CIFAR 10 data set and answer all the questions from this script\n",
    "\n",
    "Submit:\n",
    "\n",
    "1. your completed `.ipynb` notebook;\n",
    "2. all requested plots;\n",
    "3. answers to the written questions;\n",
    "4. the final test accuracy;\n",
    "5. a short comparison of at least two CNN architectures.\n",
    "\n",
    "Your notebook should run sequentially from the first cell to the last cell without errors.\n",
    "\n",
    "Note:Downloding the data set can take some time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4479345f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#TO import the CIFAR - 10 data set, use the following code: \n",
    "import torch\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# Convert images to PyTorch tensors\n",
    "transform = transforms.ToTensor()\n",
    "\n",
    "# Download CIFAR-10 training dataset\n",
    "train_dataset = datasets.CIFAR10(\n",
    "    root=\"./data\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=transform\n",
    ")\n",
    "\n",
    "# Download CIFAR-10 test dataset\n",
    "test_dataset = datasets.CIFAR10(\n",
    "    root=\"./data\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=transform\n",
    ")\n",
    "\n",
    "print(\"Training samples:\", len(train_dataset))\n",
    "print(\"Test samples:\", len(test_dataset))\n",
    "\n",
    "# Inspect one sample\n",
    "image, label = train_dataset[0]\n",
    "\n",
    "print(\"Image shape:\", image.shape)\n",
    "print(\"Label:\", label)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (PT-Control)",
   "language": "python",
   "name": "pt-control"
  },
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}