{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9d123d30",
   "metadata": {},
   "source": [
    "# Lesson 29: PyTorch CIFAR-10 classifier activity\n",
    "\n",
    "In this activity, you will build a deep neural network classifier for the CIFAR-10 dataset using PyTorch. The data loading and preparation code is provided. Your task is to:\n",
    "\n",
    "1. **Define the model** - Build a DNN using `nn.Sequential`\n",
    "2. **Train the model** - Write a training loop with validation tracking\n",
    "3. **Evaluate the model** - Assess performance on the test set\n",
    "\n",
    "## Notebook set-up\n",
    "\n",
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc5fca69",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Standard library imports\n",
    "from pathlib import Path\n",
    "\n",
    "# Third party imports\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# Set random seeds for reproducibility\n",
    "torch.manual_seed(315)\n",
    "np.random.seed(315)\n",
    "\n",
    "# Check for GPU availability\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(f'Using device: {device}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bac7a0f",
   "metadata": {},
   "source": [
    "### Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e8a0b2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 10000\n",
    "learning_rate = 1e-2\n",
    "epochs = 100\n",
    "print_every = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1ef760e",
   "metadata": {},
   "source": [
    "## 1. Load and preprocess CIFAR-10 data\n",
    "\n",
    "CIFAR-10 contains 32x32 color images across 10 classes. We convert the images to grayscale for this exercise.\n",
    "\n",
    "### 1.1. Define transformations and class names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb419f4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define class names\n",
    "class_names = [\n",
    "    'airplane', 'automobile', 'bird', 'cat', 'deer',\n",
    "    'dog', 'frog', 'horse', 'ship', 'truck'\n",
    "]\n",
    "\n",
    "# Data preprocessing: convert to grayscale, tensor, and normalize\n",
    "transform = transforms.Compose([\n",
    "    transforms.Grayscale(num_output_channels=1),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5,), (0.5,))\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e923908",
   "metadata": {},
   "source": [
    "### 1.2. Load datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5eae2637",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make sure data directory exists\n",
    "data_dir = Path('./data')\n",
    "data_dir.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# Load training and test datasets\n",
    "train_dataset = datasets.CIFAR10(\n",
    "    root=data_dir,\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=transform\n",
    ")\n",
    "\n",
    "test_dataset = datasets.CIFAR10(\n",
    "    root=data_dir,\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=transform\n",
    ")\n",
    "\n",
    "print(f'Training samples: {len(train_dataset)}')\n",
    "print(f'Test samples: {len(test_dataset)}')\n",
    "print(f'Image shape: {train_dataset[0][0].shape}')\n",
    "print(f'Number of classes: {len(class_names)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca4049d1",
   "metadata": {},
   "source": [
    "### 1.3. Pre-load data and create data loaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f7ab15e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pre-load entire dataset to device for faster training\n",
    "X_train_full = torch.stack([img for img, _ in train_dataset]).to(device)\n",
    "y_train_full = torch.tensor([label for _, label in train_dataset]).to(device)\n",
    "\n",
    "X_test = torch.stack([img for img, _ in test_dataset]).to(device)\n",
    "y_test = torch.tensor([label for _, label in test_dataset]).to(device)\n",
    "\n",
    "# Split training data into train and validation sets (80/20 split)\n",
    "n_train = int(0.8 * len(X_train_full))\n",
    "indices = torch.randperm(len(X_train_full))\n",
    "\n",
    "X_train = X_train_full[indices[:n_train]]\n",
    "y_train = y_train_full[indices[:n_train]]\n",
    "X_val = X_train_full[indices[n_train:]]\n",
    "y_val = y_train_full[indices[n_train:]]\n",
    "\n",
    "print(f'X_train shape: {X_train.shape}, device: {X_train.device}')\n",
    "print(f'y_train shape: {y_train.shape}, device: {y_train.device}')\n",
    "print(f'X_val shape: {X_val.shape}, device: {X_val.device}')\n",
    "print(f'y_val shape: {y_val.shape}, device: {y_val.device}')\n",
    "print(f'X_test shape: {X_test.shape}, device: {X_test.device}')\n",
    "print(f'y_test shape: {y_test.shape}, device: {y_test.device}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e563bd2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create TensorDatasets\n",
    "train_tensor_dataset = torch.utils.data.TensorDataset(X_train, y_train)\n",
    "val_tensor_dataset = torch.utils.data.TensorDataset(X_val, y_val)\n",
    "test_tensor_dataset = torch.utils.data.TensorDataset(X_test, y_test)\n",
    "\n",
    "# Create DataLoaders\n",
    "train_loader = DataLoader(\n",
    "    train_tensor_dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True\n",
    ")\n",
    "\n",
    "val_loader = DataLoader(\n",
    "    val_tensor_dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=False\n",
    ")\n",
    "\n",
    "test_loader = DataLoader(\n",
    "    test_tensor_dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=False\n",
    ")\n",
    "\n",
    "print(f'Training batches: {len(train_loader)}')\n",
    "print(f'Validation batches: {len(val_loader)}')\n",
    "print(f'Test batches: {len(test_loader)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6eb1ae0a",
   "metadata": {},
   "source": [
    "### 1.4. Visualize sample images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6b9e3eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get a batch of training images\n",
    "images, labels = next(iter(train_loader))\n",
    "\n",
    "# Plot first 10 images\n",
    "ncols = 5\n",
    "nrows = 2\n",
    "\n",
    "fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(ncols*1.5, nrows*1.5))\n",
    "axes = axes.flatten()\n",
    "\n",
    "for i, ax in enumerate(axes):\n",
    "\n",
    "    img = images[i].cpu() * 0.5 + 0.5\n",
    "    img = img.numpy().squeeze()\n",
    "    ax.set_title(class_names[labels[i]])\n",
    "    ax.imshow(img, cmap='gray')\n",
    "    ax.axis('off')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7991bb09",
   "metadata": {},
   "source": [
    "## 2. Build DNN classifier\n",
    "\n",
    "### Task 1: Define model architecture\n",
    "\n",
    "Build a fully connected neural network using `nn.Sequential` to classify CIFAR-10 images.\n",
    "\n",
    "**Requirements:**\n",
    "- Flatten the input images (32x32x1 = 1024 features)\n",
    "- Use at least 2 hidden layers with ReLU activation\n",
    "- Add dropout for regularization\n",
    "- Output layer should have 10 units (one per class)\n",
    "\n",
    "**Hints:**\n",
    "- Use `nn.Flatten()` as the first layer to convert images to vectors\n",
    "- Use `nn.Linear(in_features, out_features)` for fully connected layers\n",
    "- Use `nn.ReLU()` for activation functions\n",
    "- Use `nn.Dropout(p)` for regularization (e.g., p=0.2)\n",
    "- Don't forget to move the model to the device with `.to(device)`\n",
    "- `nn.CrossEntropyLoss` applies softmax internally, so no activation needed on output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d3c5002",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Define your model architecture\n",
    "input_size = 32 * 32 * 1  # Grayscale image flattened\n",
    "num_classes = 10\n",
    "\n",
    "model = nn.Sequential(\n",
    "    # Add your layers here\n",
    ").to(device)\n",
    "\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17d7b17b",
   "metadata": {},
   "source": [
    "### Task 2: Define loss function and optimizer\n",
    "\n",
    "**Hints:**\n",
    "- Use `nn.CrossEntropyLoss()` for multi-class classification\n",
    "- Use `optim.Adam(model.parameters(), lr=learning_rate)` as the optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5168d727",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Define loss function and optimizer\n",
    "criterion = None  # Replace with loss function\n",
    "optimizer = None  # Replace with optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3385ef14",
   "metadata": {},
   "source": [
    "## 3. Training\n",
    "\n",
    "### Task 3: Write training function\n",
    "\n",
    "Write a training loop that trains the model and tracks both training and validation metrics.\n",
    "\n",
    "**Requirements:**\n",
    "- Iterate over epochs\n",
    "- For each epoch, iterate over batches in the training loader\n",
    "- Track training loss and accuracy\n",
    "- After training, evaluate on validation set (without gradient computation)\n",
    "- Track validation loss and accuracy\n",
    "- Return a history dictionary with all metrics\n",
    "\n",
    "**Hints:**\n",
    "- Use `model.train()` before training and `model.eval()` before validation\n",
    "- Use `optimizer.zero_grad()` to clear gradients before each batch\n",
    "- Use `loss.backward()` for backpropagation\n",
    "- Use `optimizer.step()` to update weights\n",
    "- Use `torch.no_grad()` context manager for validation\n",
    "- Use `torch.max(outputs, 1)` to get predictions from logits\n",
    "- Accuracy = 100 * correct / total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2546213",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(\n",
    "    model: nn.Module,\n",
    "    train_loader: DataLoader,\n",
    "    val_loader: DataLoader,\n",
    "    criterion: nn.Module,\n",
    "    optimizer: optim.Optimizer,\n",
    "    epochs: int = 10,\n",
    "    print_every: int = 1\n",
    ") -> dict[str, list[float]]:\n",
    "    '''Training loop for PyTorch classification model.\n",
    "    \n",
    "    TODO: Implement the training loop with:\n",
    "    1. Training phase - iterate over train_loader batches\n",
    "    2. Validation phase - evaluate on val_loader after each epoch\n",
    "    3. Track and return history of train_loss, val_loss, train_accuracy, val_accuracy\n",
    "    '''\n",
    "\n",
    "    history = {'train_loss': [], 'val_loss': [], 'train_accuracy': [], 'val_accuracy': []}\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "\n",
    "        # TODO: Training phase\n",
    "        # - Set model to training mode\n",
    "        # - Loop over batches in train_loader\n",
    "        # - For each batch: zero gradients, forward pass, compute loss, backward pass, update weights\n",
    "        # - Track running loss and accuracy\n",
    "        pass\n",
    "\n",
    "        # TODO: Validation phase\n",
    "        # - Set model to evaluation mode\n",
    "        # - Use torch.no_grad() context\n",
    "        # - Loop over batches in val_loader\n",
    "        # - Compute loss and accuracy (no gradient computation needed)\n",
    "        pass\n",
    "\n",
    "        # TODO: Record metrics in history dict\n",
    "\n",
    "        # TODO: Print progress\n",
    "\n",
    "    print('\\nTraining complete.')\n",
    "\n",
    "    return history"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70394432",
   "metadata": {},
   "source": [
    "### Task 4: Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f228e200",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Call your training function\n",
    "history = train_model(\n",
    "    model=model,\n",
    "    train_loader=train_loader,\n",
    "    val_loader=val_loader,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    epochs=epochs,\n",
    "    print_every=print_every\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85a3dd9e",
   "metadata": {},
   "source": [
    "### Task 5: Plot learning curves\n",
    "\n",
    "**Hints:**\n",
    "- Create a 1x2 subplot for loss and accuracy\n",
    "- Plot both training and validation metrics on each subplot\n",
    "- Add legends to distinguish the curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9d2b248",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Plot learning curves\n",
    "# - Left plot: training and validation loss over epochs\n",
    "# - Right plot: training and validation accuracy over epochs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59e55813",
   "metadata": {},
   "source": [
    "## 4. Evaluate model on test set\n",
    "\n",
    "### Task 6: Write evaluation function\n",
    "\n",
    "Write a function to evaluate the model on the test set and return accuracy and predictions.\n",
    "\n",
    "**Hints:**\n",
    "- Set model to evaluation mode with `model.eval()`\n",
    "- Use `torch.no_grad()` context\n",
    "- Iterate over test_loader and accumulate predictions\n",
    "- Return accuracy, predictions array, and true labels array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea0c0a63",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_model(\n",
    "    model: nn.Module,\n",
    "    test_loader: DataLoader\n",
    ") -> tuple[float, np.ndarray, np.ndarray]:\n",
    "    '''Evaluate model on test set.\n",
    "    \n",
    "    TODO: Implement evaluation logic\n",
    "    Returns: (accuracy, predictions, true_labels)\n",
    "    '''\n",
    "\n",
    "    # TODO: Implement evaluation\n",
    "    pass\n",
    "\n",
    "\n",
    "# TODO: Call your evaluation function and print test accuracy\n",
    "# test_accuracy, predictions, true_labels = evaluate_model(model, test_loader)\n",
    "# print(f'Test accuracy: {test_accuracy:.2f}%')"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}