{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cc8e9454",
   "metadata": {},
   "source": [
    "# PyTorch CNN Activity: Experimenting with Model Architecture and Data\n",
    "\n",
    "In this activity, you will experiment with different aspects of the CNN model from the Lesson 31 demo to understand how architectural choices and data preprocessing affect model performance.\n",
    "\n",
    "## Activity Overview\n",
    "\n",
    "You will conduct **three experiments** using the code provided in the demo notebook:\n",
    "\n",
    "1. **Experiment 1**: Add more convolutional blocks and/or modify filters and filter sizes\n",
    "2. **Experiment 2**: Use RGB images instead of grayscale\n",
    "3. **Experiment 3**: Add image augmentation using PyTorch transforms\n",
    "\n",
    "For each experiment, you will:\n",
    "- Modify the relevant code sections\n",
    "- Train the model\n",
    "- Compare results with the baseline model\n",
    "- Document your observations\n",
    "\n",
    "**Note**: You may want to reduce the number of epochs (e.g., to 20-30) to speed up experimentation.\n",
    "\n",
    "## Notebook Setup\n",
    "\n",
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdd302c7",
   "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",
    "\n",
    "from sklearn.metrics import (\n",
    "    confusion_matrix, roc_curve, auc, \n",
    "    precision_recall_curve, average_precision_score\n",
    ")\n",
    "\n",
    "from sklearn.preprocessing import label_binarize\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": "665bb71a",
   "metadata": {},
   "source": [
    "### Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df8b361b",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 1000 # Training images come in 5 batches of 10,000\n",
    "learning_rate = 1e-3\n",
    "epochs = 30\n",
    "print_every = 5 # Print training progress every n epochs\n",
    "\n",
    "# CIFAR-10 class names in class order\n",
    "class_names = [\n",
    "    'airplane', 'automobile', 'bird', 'cat', 'deer',\n",
    "    'dog', 'frog', 'horse', 'ship', 'truck'\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc030588",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Experiment 1: Modify Convolutional Architecture\n",
    "\n",
    "**Objective**: Explore how adding more convolutional blocks or changing filter counts and kernel sizes affects model performance.\n",
    "\n",
    "### 1.1. Load and Preprocess Data (Baseline - Grayscale)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4da77825",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make sure data directory exists\n",
    "data_dir = Path('../data')\n",
    "data_dir.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# Data preprocessing: convert to grayscale, tensor, and normalize\n",
    "transform_exp1 = transforms.Compose([\n",
    "    transforms.Grayscale(num_output_channels=1),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5,), (0.5,))\n",
    "])\n",
    "\n",
    "# Load training and test datasets\n",
    "train_dataset_exp1 = datasets.CIFAR10(\n",
    "    root=data_dir,\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=transform_exp1\n",
    ")\n",
    "\n",
    "test_dataset_exp1 = datasets.CIFAR10(\n",
    "    root=data_dir,\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=transform_exp1\n",
    ")\n",
    "\n",
    "print(f'Training samples: {len(train_dataset_exp1)}')\n",
    "print(f'Test samples: {len(test_dataset_exp1)}')\n",
    "print(f'Image shape: {train_dataset_exp1[0][0].shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1eb6f135",
   "metadata": {},
   "source": [
    "### 1.2. Create Data Loaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65ff48e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create training, validation and testing tensors\n",
    "X_train_full = torch.stack([img for img, _ in train_dataset_exp1]).to(device)\n",
    "y_train_full = torch.tensor([label for _, label in train_dataset_exp1]).to(device)\n",
    "\n",
    "X_test_exp1 = torch.stack([img for img, _ in test_dataset_exp1]).to(device)\n",
    "y_test_exp1 = torch.tensor([label for _, label in test_dataset_exp1]).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_exp1 = X_train_full[indices[:n_train]]\n",
    "y_train_exp1 = y_train_full[indices[:n_train]]\n",
    "X_val_exp1 = X_train_full[indices[n_train:]]\n",
    "y_val_exp1 = y_train_full[indices[n_train:]]\n",
    "\n",
    "# Create TensorDatasets and DataLoaders\n",
    "train_tensor_dataset_exp1 = torch.utils.data.TensorDataset(X_train_exp1, y_train_exp1)\n",
    "val_tensor_dataset_exp1 = torch.utils.data.TensorDataset(X_val_exp1, y_val_exp1)\n",
    "test_tensor_dataset_exp1 = torch.utils.data.TensorDataset(X_test_exp1, y_test_exp1)\n",
    "\n",
    "train_loader_exp1 = DataLoader(train_tensor_dataset_exp1, batch_size=batch_size, shuffle=True)\n",
    "val_loader_exp1 = DataLoader(val_tensor_dataset_exp1, batch_size=batch_size, shuffle=False)\n",
    "test_loader_exp1 = DataLoader(test_tensor_dataset_exp1, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "print(f'Training batches: {len(train_loader_exp1)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e926ebcb",
   "metadata": {},
   "source": [
    "### 1.3. Define Modified CNN Architecture\n",
    "\n",
    "**TODO**: Modify the model architecture below. Try one or more of these changes:\n",
    "\n",
    "- Add another convolutional block (Conv2d -> BatchNorm2d -> ReLU -> Conv2d -> BatchNorm2d -> ReLU -> MaxPool2d -> Dropout)\n",
    "- Increase the number of filters (e.g., change 32 to 64 or 128)\n",
    "- Experiment with different kernel sizes (e.g., 5x5 instead of 3x3)\n",
    "- Try different pooling strategies\n",
    "\n",
    "**Important**: Remember to update the input size to the first Linear layer if you change the architecture! The size depends on:\n",
    "- Number of filters in the last conv layer\n",
    "- Final spatial dimensions after pooling\n",
    "\n",
    "*Hint*: For a 32x32 image, each MaxPool2d(2,2) layer divides dimensions by 2. So:\n",
    "- After 1 pooling: 32 → 16\n",
    "- After 2 poolings: 32 → 16 → 8\n",
    "- After 3 poolings: 32 → 16 → 8 → 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e834bc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_classes = 10\n",
    "\n",
    "model_exp1 = nn.Sequential(\n",
    "\n",
    "    # Conv block: grayscale input\n",
    "    nn.Conv2d(1, 32, kernel_size=3, padding=1),\n",
    "    nn.BatchNorm2d(32),\n",
    "    nn.ReLU(),\n",
    "    nn.Conv2d(32, 32, kernel_size=3, padding=1),\n",
    "    nn.BatchNorm2d(32),\n",
    "    nn.ReLU(),\n",
    "    nn.MaxPool2d(2, 2),\n",
    "    nn.Dropout(0.5),\n",
    "    \n",
    "    # Classifier\n",
    "    nn.Flatten(),\n",
    "    nn.Linear(32 * 16 * 16, 128),\n",
    "    nn.ReLU(),\n",
    "    nn.Dropout(0.5),\n",
    "    nn.Linear(128, num_classes)\n",
    "\n",
    ").to(device)\n",
    "\n",
    "trainable_params = sum(p.numel() for p in model_exp1.parameters() if p.requires_grad)\n",
    "print(model_exp1)\n",
    "print(f'\\nTotal parameters: {trainable_params}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0c435f3",
   "metadata": {},
   "source": [
    "### 1.4. Train Modified Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e4102e5",
   "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",
    "    device: torch.device = None\n",
    ") -> dict[str, list[float]]:\n",
    "    '''Training loop for PyTorch classification model.\n",
    "    \n",
    "    Args:\n",
    "        device: If provided, moves batches to this device on-the-fly.\n",
    "                If None, assumes data is already on the correct device.\n",
    "    '''\n",
    "\n",
    "    history = {'train_loss': [], 'val_loss': [], 'train_accuracy': [], 'val_accuracy': []}\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "\n",
    "        # Training phase\n",
    "        model.train()\n",
    "        running_loss = 0.0\n",
    "        correct = 0\n",
    "        total = 0\n",
    "\n",
    "        for images, labels in train_loader:\n",
    "            \n",
    "            # Move batch to device if specified\n",
    "            if device is not None:\n",
    "                images, labels = images.to(device), labels.to(device)\n",
    "\n",
    "            # Forward pass\n",
    "            optimizer.zero_grad()\n",
    "            outputs = model(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "\n",
    "            # Backward pass\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            # Track metrics\n",
    "            running_loss += loss.item()\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "\n",
    "        # Calculate training metrics\n",
    "        train_loss = running_loss / len(train_loader)\n",
    "        train_accuracy = 100 * correct / total\n",
    "\n",
    "        # Validation phase\n",
    "        model.eval()\n",
    "        val_running_loss = 0.0\n",
    "        val_correct = 0\n",
    "        val_total = 0\n",
    "\n",
    "        with torch.no_grad():\n",
    "\n",
    "            for images, labels in val_loader:\n",
    "                \n",
    "                # Move batch to device if specified\n",
    "                if device is not None:\n",
    "                    images, labels = images.to(device), labels.to(device)\n",
    "\n",
    "                outputs = model(images)\n",
    "                loss = criterion(outputs, labels)\n",
    "\n",
    "                val_running_loss += loss.item()\n",
    "                _, predicted = torch.max(outputs.data, 1)\n",
    "                val_total += labels.size(0)\n",
    "                val_correct += (predicted == labels).sum().item()\n",
    "\n",
    "        val_loss = val_running_loss / len(val_loader)\n",
    "        val_accuracy = 100 * val_correct / val_total\n",
    "\n",
    "        # Record metrics\n",
    "        history['train_loss'].append(train_loss)\n",
    "        history['val_loss'].append(val_loss)\n",
    "        history['train_accuracy'].append(train_accuracy)\n",
    "        history['val_accuracy'].append(val_accuracy)\n",
    "\n",
    "        # Print progress\n",
    "        if (epoch + 1) % print_every == 0 or epoch == 0:\n",
    "\n",
    "            print(\n",
    "                f'Epoch {epoch+1}/{epochs} - ' +\n",
    "                f'loss: {train_loss:.4f} - ' +\n",
    "                f'accuracy: {train_accuracy:.2f}% - ' +\n",
    "                f'val_loss: {val_loss:.4f} - ' +\n",
    "                f'val_accuracy: {val_accuracy:.2f}%'\n",
    "            )\n",
    "\n",
    "    print('\\nTraining complete.')\n",
    "\n",
    "    return history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "865041b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "criterion_exp1 = nn.CrossEntropyLoss()\n",
    "optimizer_exp1 = optim.Adam(model_exp1.parameters(), lr=learning_rate)\n",
    "\n",
    "history_exp1 = train_model(\n",
    "    model=model_exp1,\n",
    "    train_loader=train_loader_exp1,\n",
    "    val_loader=val_loader_exp1,\n",
    "    criterion=criterion_exp1,\n",
    "    optimizer=optimizer_exp1,\n",
    "    epochs=epochs,\n",
    "    print_every=print_every\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73fdf140",
   "metadata": {},
   "source": [
    "### 1.5. Evaluate and Visualize Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08942db9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_model(\n",
    "    model: nn.Module,\n",
    "    test_loader: DataLoader,\n",
    "    device: torch.device = None\n",
    ") -> tuple[float, np.ndarray, np.ndarray]:\n",
    "    '''Evaluate model on test set.\n",
    "    \n",
    "    Args:\n",
    "        device: If provided, moves batches to this device on-the-fly.\n",
    "                If None, assumes data is already on the correct device.\n",
    "    '''\n",
    "\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    all_predictions = []\n",
    "    all_labels = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "\n",
    "        for images, labels in test_loader:\n",
    "            \n",
    "            # Move batch to device if specified\n",
    "            if device is not None:\n",
    "                images, labels = images.to(device), labels.to(device)\n",
    "\n",
    "            outputs = model(images)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "\n",
    "            all_predictions.extend(predicted.cpu().numpy())\n",
    "            all_labels.extend(labels.cpu().numpy())\n",
    "\n",
    "    accuracy = 100 * correct / total\n",
    "    return accuracy, np.array(all_predictions), np.array(all_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c0afcf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Learning curves\n",
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
    "\n",
    "axes[0].set_title('Loss - Experiment 1')\n",
    "axes[0].plot(history_exp1['train_loss'], label='Train')\n",
    "axes[0].plot(history_exp1['val_loss'], label='Validation')\n",
    "axes[0].set_xlabel('Epoch')\n",
    "axes[0].set_ylabel('Loss (cross-entropy)')\n",
    "axes[0].legend(loc='best')\n",
    "\n",
    "axes[1].set_title('Accuracy - Experiment 1')\n",
    "axes[1].plot(history_exp1['train_accuracy'], label='Train')\n",
    "axes[1].plot(history_exp1['val_accuracy'], label='Validation')\n",
    "axes[1].set_xlabel('Epoch')\n",
    "axes[1].set_ylabel('Accuracy (%)')\n",
    "axes[1].legend(loc='best')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Test accuracy\n",
    "test_accuracy_exp1, predictions_exp1, true_labels_exp1 = evaluate_model(model_exp1, test_loader_exp1)\n",
    "print(f'\\nExperiment 1 Test Accuracy: {test_accuracy_exp1:.2f}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62d510de",
   "metadata": {},
   "source": [
    "### 1.6. Experiment 1 Observations\n",
    "\n",
    "**TODO**: Document your findings:\n",
    "- What architectural changes did you make?\n",
    "- How did these changes affect the number of parameters?\n",
    "- Did the model perform better or worse than the baseline?\n",
    "- What did you observe about training time and convergence?\n",
    "- Did you notice any overfitting or underfitting?\n",
    "\n",
    "*Your notes here:*\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbc07230",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Experiment 2: Use RGB Images Instead of Grayscale\n",
    "\n",
    "**Objective**: Compare model performance using full RGB color information versus grayscale.\n",
    "\n",
    "### 2.1. Load and Preprocess RGB Data\n",
    "\n",
    "**TODO**: Modify the transform to keep RGB channels (3 channels) instead of converting to grayscale. Update the normalization to use 3 mean and std values.\n",
    "\n",
    "*Hint*: For RGB, use `transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cf37587",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Modify this transform to use RGB instead of grayscale\n",
    "transform_exp2 = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "])\n",
    "\n",
    "# Load training and test datasets with RGB\n",
    "train_dataset_exp2 = datasets.CIFAR10(\n",
    "    root=data_dir,\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=transform_exp2\n",
    ")\n",
    "\n",
    "test_dataset_exp2 = datasets.CIFAR10(\n",
    "    root=data_dir,\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=transform_exp2\n",
    ")\n",
    "\n",
    "print(f'Image shape: {train_dataset_exp2[0][0].shape}')  # Should be [3, 32, 32]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "742ac7b8",
   "metadata": {},
   "source": [
    "### 2.2. Visualize RGB Sample Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72836e10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot first 10 RGB images from the training dataset\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",
    "    img, label = train_dataset_exp2[i]\n",
    "    \n",
    "    # Unnormalize and transpose for plotting\n",
    "    img = img * 0.5 + 0.5\n",
    "    img = img.numpy().transpose(1, 2, 0)  # Change from CxHxW to HxWxC\n",
    "    ax.set_title(class_names[label])\n",
    "    ax.imshow(img)\n",
    "    ax.axis('off')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da306112",
   "metadata": {},
   "source": [
    "### 2.3. Create Data Loaders for RGB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a11adc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create training, validation and testing tensors\n",
    "X_train_full_exp2 = torch.stack([img for img, _ in train_dataset_exp2]).to(device)\n",
    "y_train_full_exp2 = torch.tensor([label for _, label in train_dataset_exp2]).to(device)\n",
    "\n",
    "X_test_exp2 = torch.stack([img for img, _ in test_dataset_exp2]).to(device)\n",
    "y_test_exp2 = torch.tensor([label for _, label in test_dataset_exp2]).to(device)\n",
    "\n",
    "# Split training data\n",
    "n_train = int(0.8 * len(X_train_full_exp2))\n",
    "indices = torch.randperm(len(X_train_full_exp2))\n",
    "\n",
    "X_train_exp2 = X_train_full_exp2[indices[:n_train]]\n",
    "y_train_exp2 = y_train_full_exp2[indices[:n_train]]\n",
    "X_val_exp2 = X_train_full_exp2[indices[n_train:]]\n",
    "y_val_exp2 = y_train_full_exp2[indices[n_train:]]\n",
    "\n",
    "# Create DataLoaders\n",
    "train_tensor_dataset_exp2 = torch.utils.data.TensorDataset(X_train_exp2, y_train_exp2)\n",
    "val_tensor_dataset_exp2 = torch.utils.data.TensorDataset(X_val_exp2, y_val_exp2)\n",
    "test_tensor_dataset_exp2 = torch.utils.data.TensorDataset(X_test_exp2, y_test_exp2)\n",
    "\n",
    "train_loader_exp2 = DataLoader(train_tensor_dataset_exp2, batch_size=batch_size, shuffle=True)\n",
    "val_loader_exp2 = DataLoader(val_tensor_dataset_exp2, batch_size=batch_size, shuffle=False)\n",
    "test_loader_exp2 = DataLoader(test_tensor_dataset_exp2, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "print(f'X_train shape: {X_train_exp2.shape}')  # Should show 3 channels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee7a5f59",
   "metadata": {},
   "source": [
    "### 2.4. Define CNN for RGB Images\n",
    "\n",
    "**TODO**: Modify the first Conv2d layer to accept 3 input channels instead of 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aadebd5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Update the first conv layer to accept 3 channels\n",
    "model_exp2 = nn.Sequential(\n",
    "\n",
    "    # Conv block: RGB input\n",
    "    nn.Conv2d(1, 32, kernel_size=3, padding=1),\n",
    "    nn.BatchNorm2d(32),\n",
    "    nn.ReLU(),\n",
    "    nn.Conv2d(32, 32, kernel_size=3, padding=1),\n",
    "    nn.BatchNorm2d(32),\n",
    "    nn.ReLU(),\n",
    "    nn.MaxPool2d(2, 2),\n",
    "    nn.Dropout(0.5),\n",
    "    \n",
    "    # Classifier\n",
    "    nn.Flatten(),\n",
    "    nn.Linear(32 * 16 * 16, 128),\n",
    "    nn.ReLU(),\n",
    "    nn.Dropout(0.5),\n",
    "    nn.Linear(128, num_classes)\n",
    "\n",
    ").to(device)\n",
    "\n",
    "trainable_params = sum(p.numel() for p in model_exp2.parameters() if p.requires_grad)\n",
    "print(model_exp2)\n",
    "print(f'\\nTotal parameters: {trainable_params}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c4f01b3",
   "metadata": {},
   "source": [
    "### 2.5. Train RGB Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6104f3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "criterion_exp2 = nn.CrossEntropyLoss()\n",
    "optimizer_exp2 = optim.Adam(model_exp2.parameters(), lr=learning_rate)\n",
    "\n",
    "history_exp2 = train_model(\n",
    "    model=model_exp2,\n",
    "    train_loader=train_loader_exp2,\n",
    "    val_loader=val_loader_exp2,\n",
    "    criterion=criterion_exp2,\n",
    "    optimizer=optimizer_exp2,\n",
    "    epochs=epochs,\n",
    "    print_every=print_every\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8ea58e7",
   "metadata": {},
   "source": [
    "### 2.6. Evaluate RGB Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ed85bed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Learning curves\n",
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
    "\n",
    "axes[0].set_title('Loss - Experiment 2 (RGB)')\n",
    "axes[0].plot(history_exp2['train_loss'], label='Train')\n",
    "axes[0].plot(history_exp2['val_loss'], label='Validation')\n",
    "axes[0].set_xlabel('Epoch')\n",
    "axes[0].set_ylabel('Loss (cross-entropy)')\n",
    "axes[0].legend(loc='best')\n",
    "\n",
    "axes[1].set_title('Accuracy - Experiment 2 (RGB)')\n",
    "axes[1].plot(history_exp2['train_accuracy'], label='Train')\n",
    "axes[1].plot(history_exp2['val_accuracy'], label='Validation')\n",
    "axes[1].set_xlabel('Epoch')\n",
    "axes[1].set_ylabel('Accuracy (%)')\n",
    "axes[1].legend(loc='best')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Test accuracy\n",
    "test_accuracy_exp2, predictions_exp2, true_labels_exp2 = evaluate_model(model_exp2, test_loader_exp2)\n",
    "print(f'\\nExperiment 2 Test Accuracy: {test_accuracy_exp2:.2f}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "629a099e",
   "metadata": {},
   "source": [
    "### 2.7. Experiment 2 Observations\n",
    "\n",
    "**TODO**: Document your findings:\n",
    "- Did RGB improve accuracy compared to grayscale?\n",
    "- How did the number of parameters change?\n",
    "- Which classes benefited most from color information?\n",
    "- Were there any classes that performed similarly with both grayscale and RGB?\n",
    "\n",
    "*Your notes here:*\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed2ee0bd",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Experiment 3: Add Image Augmentation\n",
    "\n",
    "**Objective**: Use PyTorch transforms to augment training data and improve model generalization.\n",
    "\n",
    "### 3.1. Define Augmented Transforms\n",
    "\n",
    "**TODO**: Add image augmentation transforms to the training data. Consider:\n",
    "- `transforms.RandomHorizontalFlip(p=0.5)` - Randomly flip images horizontally\n",
    "- `transforms.RandomCrop(32, padding=4)` - Random crop with padding\n",
    "- `transforms.RandomRotation(degrees=15)` - Small random rotations\n",
    "- `transforms.ColorJitter(brightness=0.2, contrast=0.2)` - Random brightness/contrast adjustments\n",
    "\n",
    "**Note**: Apply augmentation only to training data, not validation or test data!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26c96439",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Add augmentation transforms here\n",
    "transform_train_exp3 = transforms.Compose([\n",
    "    # TODO: Add augmentation transforms here (before ToTensor)    \n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "])\n",
    "\n",
    "# Validation and test transforms (no augmentation)\n",
    "transform_test_exp3 = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "])\n",
    "\n",
    "# Load datasets with augmentation\n",
    "train_dataset_exp3 = datasets.CIFAR10(\n",
    "    root=data_dir,\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=transform_train_exp3  # Augmented transform\n",
    ")\n",
    "\n",
    "test_dataset_exp3 = datasets.CIFAR10(\n",
    "    root=data_dir,\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=transform_test_exp3  # No augmentation\n",
    ")\n",
    "\n",
    "print('Datasets loaded with augmentation')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce43da7e",
   "metadata": {},
   "source": [
    "### 3.2. Visualize Augmented Images\n",
    "\n",
    "Run this cell multiple times to see different augmentations of the same images!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cafa7a78",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize augmented versions of the same 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",
    "    # Get augmented version each time this is run\n",
    "    img, label = train_dataset_exp3[i]\n",
    "    \n",
    "    # Unnormalize and transpose for plotting\n",
    "    img = img * 0.5 + 0.5\n",
    "    img = img.numpy().transpose(1, 2, 0)\n",
    "    ax.set_title(class_names[label])\n",
    "    ax.imshow(img)\n",
    "    ax.axis('off')\n",
    "\n",
    "plt.suptitle('Augmented Training Images (run cell again to see different augmentations)', y=1.02)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba0683e5",
   "metadata": {},
   "source": [
    "### 3.3. Create Data Loaders with Augmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "842ee568",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For data augmentation, we must NOT preload data to GPU as tensors.\n",
    "# Transforms need to be applied on-the-fly during each epoch so each \n",
    "# batch sees different augmented versions of the images.\n",
    "\n",
    "# Split training data into train and validation sets using Subset\n",
    "n_train = int(0.8 * len(train_dataset_exp3))\n",
    "n_val = len(train_dataset_exp3) - n_train\n",
    "indices = torch.randperm(len(train_dataset_exp3)).tolist()\n",
    "\n",
    "train_subset_exp3 = torch.utils.data.Subset(train_dataset_exp3, indices[:n_train])\n",
    "val_subset_exp3 = torch.utils.data.Subset(train_dataset_exp3, indices[n_train:])\n",
    "\n",
    "print(f'Training samples: {len(train_subset_exp3)}')\n",
    "print(f'Validation samples: {len(val_subset_exp3)}')\n",
    "print(f'Test samples: {len(test_dataset_exp3)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70daf63e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create DataLoaders directly from Dataset/Subset objects\n",
    "# Transforms are applied on-the-fly when batches are loaded\n",
    "train_loader_exp3 = DataLoader(\n",
    "    train_subset_exp3,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True\n",
    ")\n",
    "\n",
    "val_loader_exp3 = DataLoader(\n",
    "    val_subset_exp3,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=False\n",
    ")\n",
    "\n",
    "test_loader_exp3 = DataLoader(\n",
    "    test_dataset_exp3,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=False\n",
    ")\n",
    "\n",
    "print(f'Training batches: {len(train_loader_exp3)}')\n",
    "print(f'Validation batches: {len(val_loader_exp3)}')\n",
    "print(f'Test batches: {len(test_loader_exp3)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed1a28d8",
   "metadata": {},
   "source": [
    "### 3.4. Define Model for Augmented Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac5365b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Same architecture as Experiment 2 (RGB)\n",
    "model_exp3 = nn.Sequential(\n",
    "\n",
    "    # Conv block: RGB input\n",
    "    nn.Conv2d(3, 32, kernel_size=3, padding=1),\n",
    "    nn.BatchNorm2d(32),\n",
    "    nn.ReLU(),\n",
    "    nn.Conv2d(32, 32, kernel_size=3, padding=1),\n",
    "    nn.BatchNorm2d(32),\n",
    "    nn.ReLU(),\n",
    "    nn.MaxPool2d(2, 2),\n",
    "    nn.Dropout(0.5),\n",
    "    \n",
    "    # Classifier\n",
    "    nn.Flatten(),\n",
    "    nn.Linear(32 * 16 * 16, 128),\n",
    "    nn.ReLU(),\n",
    "    nn.Dropout(0.5),\n",
    "    nn.Linear(128, num_classes)\n",
    "\n",
    ").to(device)\n",
    "\n",
    "trainable_params = sum(p.numel() for p in model_exp3.parameters() if p.requires_grad)\n",
    "print(f'Total parameters: {trainable_params}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b46d131b",
   "metadata": {},
   "source": [
    "### 3.5. Train Model with Augmented Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00a6fc79",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "criterion_exp3 = nn.CrossEntropyLoss()\n",
    "optimizer_exp3 = optim.Adam(model_exp3.parameters(), lr=learning_rate)\n",
    "\n",
    "# Pass device to move batches on-the-fly (required for on-the-fly augmentation)\n",
    "history_exp3 = train_model(\n",
    "    model=model_exp3,\n",
    "    train_loader=train_loader_exp3,\n",
    "    val_loader=val_loader_exp3,\n",
    "    criterion=criterion_exp3,\n",
    "    optimizer=optimizer_exp3,\n",
    "    epochs=epochs,\n",
    "    print_every=print_every,\n",
    "    device=device\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85ad8659",
   "metadata": {},
   "source": [
    "### 3.6. Evaluate Augmented Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc7e3481",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Learning curves\n",
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
    "\n",
    "axes[0].set_title('Loss - Experiment 3 (Augmented)')\n",
    "axes[0].plot(history_exp3['train_loss'], label='Train')\n",
    "axes[0].plot(history_exp3['val_loss'], label='Validation')\n",
    "axes[0].set_xlabel('Epoch')\n",
    "axes[0].set_ylabel('Loss (cross-entropy)')\n",
    "axes[0].legend(loc='best')\n",
    "\n",
    "axes[1].set_title('Accuracy - Experiment 3 (Augmented)')\n",
    "axes[1].plot(history_exp3['train_accuracy'], label='Train')\n",
    "axes[1].plot(history_exp3['val_accuracy'], label='Validation')\n",
    "axes[1].set_xlabel('Epoch')\n",
    "axes[1].set_ylabel('Accuracy (%)')\n",
    "axes[1].legend(loc='best')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Test accuracy (pass device for on-the-fly batch loading)\n",
    "test_accuracy_exp3, predictions_exp3, true_labels_exp3 = evaluate_model(\n",
    "    model_exp3, test_loader_exp3, device=device\n",
    ")\n",
    "print(f'\\nExperiment 3 Test Accuracy: {test_accuracy_exp3:.2f}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcbd1142",
   "metadata": {},
   "source": [
    "### 3.7. Experiment 3 Observations\n",
    "\n",
    "**TODO**: Document your findings:\n",
    "- Which augmentation techniques did you use?\n",
    "- Did augmentation improve test accuracy?\n",
    "- Did you notice any effect on the gap between training and validation accuracy (overfitting)?\n",
    "- How did training time compare to non-augmented training?\n",
    "- Would you recommend augmentation for this dataset?\n",
    "\n",
    "*Your notes here:*\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52aa6bc8",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Summary: Compare All Experiments\n",
    "\n",
    "### Compare Test Accuracies\n",
    "\n",
    "**TODO**: Fill in your results and compare:\n",
    "\n",
    "| Experiment | Description | Test Accuracy | Notes |\n",
    "|------------|-------------|---------------|-------|\n",
    "| Baseline (demo) | Grayscale, simple architecture | ~60% | From demo notebook |\n",
    "| Experiment 1 | Modified architecture | _% | |\n",
    "| Experiment 2 | RGB images | _% | |\n",
    "| Experiment 3 | Image augmentation | _% | |\n",
    "\n",
    "### Final Reflections\n",
    "\n",
    "**TODO**: Based on your experiments, answer these questions:\n",
    "\n",
    "1. Which experiment produced the best results and why do you think that is?\n",
    "\n",
    "2. What trade-offs did you observe between model complexity, training time, and performance?\n",
    "\n",
    "3. If you were to combine multiple improvements (e.g., deeper architecture + RGB + augmentation), what would you expect?\n",
    "\n",
    "4. What other experiments would you like to try?\n",
    "\n",
    "*Your reflections here:*\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8155ac6",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}