{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8a97d117",
   "metadata": {},
   "source": [
    "# PyTorch CNN model\n",
    "\n",
    "In this notebook, we build a convolutional neural network (CNN) classifier for the CIFAR-10 dataset using PyTorch's `nn.Sequential` module.\n",
    "\n",
    "## Notebook set-up\n",
    "\n",
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "041d3264",
   "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": "35d44b72",
   "metadata": {},
   "source": [
    "### Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26de3062",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 10000 # Training images come in 5 batches of 10,000\n",
    "learning_rate = 1e-3\n",
    "epochs = 100\n",
    "print_every = 10 # 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": "01ea5a8b",
   "metadata": {},
   "source": [
    "## 1. Load and preprocess CIFAR-10 data\n",
    "\n",
    "CIFAR-10 contains 32x32 color images (3 channels) across 10 classes. We convert the images to grayscale for this demonstration.\n",
    "\n",
    "### 1.1. Load datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a4112b9",
   "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 = 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 = 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": "6c3aad08",
   "metadata": {},
   "source": [
    "### 1.2. Visualize sample images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10c3b8ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot first 10 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",
    "    # Get image and label from dataset\n",
    "    img, label = train_dataset[i]\n",
    "    \n",
    "    # Unnormalize and squeeze for plotting\n",
    "    img = img * 0.5 + 0.5\n",
    "    img = img.numpy().squeeze()\n",
    "    ax.set_title(class_names[label])\n",
    "    ax.imshow(img, cmap='gray')\n",
    "    ax.axis('off')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b5a77a6",
   "metadata": {},
   "source": [
    "### 1.3. Create training, validation and testing tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e1cff44",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reshape and preload 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": "markdown",
   "id": "ac80fde1",
   "metadata": {},
   "source": [
    "### 1.4. Create `DataLoader()` objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4760746",
   "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": "dae5aa25",
   "metadata": {},
   "source": [
    "## 2. Build CNN classifier with nn.Sequential\n",
    "\n",
    "We build a convolutional neural network using `nn.Sequential`. The CNN uses convolutional layers to extract spatial features from the 32x32x1 grayscale images before classification.\n",
    "\n",
    "### 2.1. Define model architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "763963f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input: 1 x 32 x 32 (grayscale)\n",
    "num_classes = 10\n",
    "\n",
    "model = nn.Sequential(\n",
    "\n",
    "    # Conv block: 1 -> 32 channels, 32 x 32 -> 16 x 16\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",
    ").to(device)\n",
    "\n",
    "trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "print(model)\n",
    "print(f'\\nTotal parameters: {trainable_params}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96a88a92",
   "metadata": {},
   "source": [
    "### 2.2. Define loss function and optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5eb64318",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8eb281d6",
   "metadata": {},
   "source": [
    "### 2.3. Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5f22f21",
   "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",
    "    Note: 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",
    "            # 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",
    "                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": "873f2d4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\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": "9106ffd4",
   "metadata": {},
   "source": [
    "### 2.5. Learning curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f5d53e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
    "\n",
    "axes[0].set_title('Loss')\n",
    "axes[0].plot(history['train_loss'], label='Train')\n",
    "axes[0].plot(history['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')\n",
    "axes[1].plot(history['train_accuracy'], label='Train')\n",
    "axes[1].plot(history['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()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5181cdc2",
   "metadata": {},
   "source": [
    "## 3. Evaluate model on test set\n",
    "\n",
    "### 3.1. Calculate test accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59bdd8c6",
   "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",
    "    Note: 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",
    "            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": "8ad4007b",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_accuracy, predictions, true_labels = evaluate_model(model, test_loader)\n",
    "print(f'Test accuracy: {test_accuracy:.2f}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e39a7521",
   "metadata": {},
   "source": [
    "### 3.2. Per-class accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83199dee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate per-class accuracy\n",
    "class_correct = {name: 0 for name in class_names}\n",
    "class_total = {name: 0 for name in class_names}\n",
    "\n",
    "for pred, true in zip(predictions, true_labels):\n",
    "\n",
    "    class_name = class_names[true]\n",
    "    class_total[class_name] += 1\n",
    "\n",
    "    if pred == true:\n",
    "        class_correct[class_name] += 1\n",
    "\n",
    "print('Per-class accuracy:')\n",
    "print('-' * 30)\n",
    "\n",
    "for name in class_names:\n",
    "    acc = 100 * class_correct[name] / class_total[name]\n",
    "    print(f'{name:12s}: {acc:.2f}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34c71064",
   "metadata": {},
   "source": [
    "### 3.4. Confusion matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a39efd36",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute confusion matrix\n",
    "cm = confusion_matrix(true_labels, predictions)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 8))\n",
    "\n",
    "ax.set_title('Confusion matrix')\n",
    "im = ax.imshow(cm, cmap='Blues')\n",
    "\n",
    "# Add labels\n",
    "ax.set_xticks(range(len(class_names)))\n",
    "ax.set_yticks(range(len(class_names)))\n",
    "ax.set_xticklabels(class_names, rotation=45, ha='right')\n",
    "ax.set_yticklabels(class_names)\n",
    "ax.set_xlabel('Predicted label')\n",
    "ax.set_ylabel('True label')\n",
    "\n",
    "# Add text annotations\n",
    "for i in range(len(class_names)):\n",
    "\n",
    "    for j in range(len(class_names)):\n",
    "        color = 'white' if cm[i, j] > cm.max() / 2 else 'black'\n",
    "        ax.text(j, i, str(cm[i, j]), ha='center', va='center', color=color)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ffa5ebe",
   "metadata": {},
   "source": [
    "### 3.5. Predicted class probability distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fe7ea38",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get predicted probabilities for all test samples\n",
    "model.eval()\n",
    "all_probs = []\n",
    "\n",
    "with torch.no_grad():\n",
    "    for images, _ in test_loader:\n",
    "        outputs = model(images)\n",
    "        probs = torch.softmax(outputs, dim=1)\n",
    "        all_probs.append(probs.cpu().numpy())\n",
    "\n",
    "all_probs = np.concatenate(all_probs, axis=0)\n",
    "\n",
    "# Plot probability distributions for each class\n",
    "fig, axes = plt.subplots(nrows=2, ncols=5, figsize=(12, 4))\n",
    "\n",
    "fig.suptitle('Predicted probability distributions by class', fontsize=14, y=1.02)\n",
    "fig.supxlabel('Predicted probability', fontsize=12)\n",
    "fig.supylabel('Count', fontsize=12)\n",
    "\n",
    "axes = axes.flatten()\n",
    "\n",
    "for i, (ax, class_name) in enumerate(zip(axes, class_names)):\n",
    "\n",
    "    # Get probabilities for this class across all samples\n",
    "    class_probs = all_probs[:, i]\n",
    "    \n",
    "    # Plot density using histogram with density normalization\n",
    "    ax.hist(class_probs, bins=50, color='black')\n",
    "    ax.set_title(class_name)\n",
    "    ax.set_xlim(0, 1)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73cdc456",
   "metadata": {},
   "source": [
    "### 3.6. Evaluation curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f574d9eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Binarize true labels for one-vs-rest evaluation\n",
    "y_test_bin = label_binarize(true_labels, classes=range(len(class_names)))\n",
    "\n",
    "# Create figure with ROC and PR curves side by side\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))\n",
    "\n",
    "# Plot ROC curves for each class\n",
    "ax1.set_title('ROC curves (one-vs-rest)')\n",
    "\n",
    "for i, class_name in enumerate(class_names):\n",
    "    fpr, tpr, _ = roc_curve(y_test_bin[:, i], all_probs[:, i])\n",
    "    roc_auc = auc(fpr, tpr)\n",
    "    ax1.plot(fpr, tpr, label=class_name)\n",
    "\n",
    "ax1.plot([0, 1], [0, 1], 'k--', label='Random classifier')\n",
    "ax1.set_xlabel('False positive rate')\n",
    "ax1.set_ylabel('True positive rate')\n",
    "ax1.legend(loc='lower right', fontsize=12)\n",
    "ax1.set_xlim([0, 1])\n",
    "ax1.set_ylim([0, 1.05])\n",
    "\n",
    "# Plot Precision-Recall curves for each class\n",
    "ax2.set_title('Precision-recall curves (one-vs-rest)')\n",
    "\n",
    "for i, class_name in enumerate(class_names):\n",
    "    precision, recall, _ = precision_recall_curve(y_test_bin[:, i], all_probs[:, i])\n",
    "    ap = average_precision_score(y_test_bin[:, i], all_probs[:, i])\n",
    "    ax2.plot(recall, precision)\n",
    "\n",
    "# Random classifier baseline (horizontal line at class prevalence = 1/num_classes)\n",
    "baseline = 1 / len(class_names)\n",
    "ax2.axhline(y=baseline, color='k', linestyle='--')\n",
    "\n",
    "ax2.set_xlabel('Recall')\n",
    "ax2.set_ylabel('Precision')\n",
    "ax2.set_xlim([0, 1])\n",
    "ax2.set_ylim([0, 1.05])\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7652023",
   "metadata": {},
   "source": [
    "## 4. Save model for inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6743f566",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define model save path\n",
    "model_dir = Path('../models')\n",
    "model_dir.mkdir(parents=True, exist_ok=True)\n",
    "model_path = model_dir / 'cifar10_cnn_model.pth'\n",
    "\n",
    "# Save the model state dict (recommended for inference)\n",
    "torch.save(model.state_dict(), model_path)\n",
    "print(f'Model saved to: {model_path.resolve()}')"
   ]
  }
 ],
 "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
}