{ "nbformat": 4, "nbformat_minor": 0, "metadata": { "colab": { "provenance": [], "gpuType": "T4", "authorship_tag": "ABX9TyMscbUSeB2qw1Lue5xsmWqV", "include_colab_link": true }, "kernelspec": { "name": "python3", "display_name": "Python 3" }, "language_info": { "name": "python" }, "accelerator": "GPU" }, "cells": [ { "cell_type": "markdown", "metadata": { "id": "view-in-github", "colab_type": "text" }, "source": [ "\"Open" ] }, { "cell_type": "markdown", "source": [ "# **Model Training and Fine-Tuning**\n", "\n", "**Objective:** Ingest the Kaggle Chest X-Ray dataset, apply strategic data augmentation, handle severe class imbalance via loss weighting, and fine-tune a pre-trained DenseNet121 architecture to detect pneumonia.\n", "\n", "## **Environment Setup**\n", "First, we import the necessary PyTorch libraries and configure our hardware accelerator. Because image processing is computationally expensive, we must ensure PyTorch detects and utilizes the Colab GPU (CUDA)." ], "metadata": { "id": "24IiOoFPN2oN" } }, { "cell_type": "code", "source": [ "%%capture install_output\n", "!pip install optuna" ], "metadata": { "id": "OwBert5NpQwF" }, "execution_count": 5, "outputs": [] }, { "cell_type": "code", "execution_count": 6, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "IGGkAlCRNlro", "outputId": "778acab3-e0b5-4c40-a995-ded91103b0a4" }, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Compute Device: cuda\n", "GPU Model: Tesla T4\n" ] } ], "source": [ "import torch\n", "import torch.nn as nn\n", "import torch.optim as optim\n", "from torchvision import datasets, transforms, models\n", "from torch.utils.data import DataLoader\n", "import os\n", "import optuna\n", "import time\n", "import copy\n", "import numpy as np\n", "\n", "# Configure the device\n", "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n", "print(f'Compute Device: {device}')\n", "if device.type == 'cuda':\n", " print(f'GPU Model: {torch.cuda.get_device_name(0)}')" ] }, { "cell_type": "markdown", "source": [ "## **Data Preprocessing and Augmentation Pipeline**\n", "As determined in our EDA, we must standardize all input resolutions to 224x224.\n", "\n", "To prevent overfitting without introducing artificial artifacts, we apply light data augmentation **only to the training set** (rotation and color jitter). The validation and test sets remain strictly unaltered to provide an honest evaluation metric." ], "metadata": { "id": "2SGyO1PxOmDX" } }, { "cell_type": "code", "source": [ "# Define the dataset paths\n", "data_dir = '../dataset/chest_xray'\n", "train_dir = os.path.join(data_dir, 'train')\n", "val_dir = os.path.join(data_dir, 'val')\n", "test_dir = os.path.join(data_dir, 'test')\n", "\n", "# 1. Training Transforms (with Augmentation)\n", "train_transforms = transforms.Compose([\n", " transforms.Resize((224, 224)),\n", " transforms.RandomRotation(10), # Slight rotation to simulate patient misalignment\n", " transforms.ColorJitter(brightness=0.2, contrast=0.2), # Account for different X-ray exposures\n", " transforms.ToTensor(),\n", " transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # Standard ImageNet normalization\n", "])\n", "\n", "# 2. Validation/Testing Transforms (Strictly unaltered)\n", "eval_transforms = transforms.Compose([\n", " transforms.Resize((224, 224)),\n", " transforms.ToTensor(),\n", " transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n", "])\n", "\n", "print(\"Transforms configured successfully.\")" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "ikceEotXOU_1", "outputId": "077395cc-6cf0-42f5-e0b3-65a1e322b5a8" }, "execution_count": 2, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Transforms configured successfully.\n" ] } ] }, { "cell_type": "markdown", "source": [ "## **Initializing Datasets and DataLoaders**\n", "\n", "With our transforms defined, we will use PyTorch's `ImageFolder` to automatically read our directory structure and assign the correct labels (0 for Normal, 1 for Pneumonia).\n", "\n", "We will then wrap these datasets in `DataLoaders`. DataLoaders handle the heavy lifting of batching the images, shuffling them, and pushing them to the GPU. We are setting the `batch_size` to 64 to fully utilize the Google Colab GPU memory." ], "metadata": { "id": "odyLrQARmmn-" } }, { "cell_type": "code", "source": [ "# 1. Create the Dataset objects\n", "train_dataset = datasets.ImageFolder(train_dir, transform=train_transforms)\n", "val_dataset = datasets.ImageFolder(val_dir, transform=eval_transforms)\n", "test_dataset = datasets.ImageFolder(test_dir, transform=eval_transforms)\n", "\n", "print(f\"Training images: {len(train_dataset)}\")\n", "print(f\"Validation images: {len(val_dataset)}\")\n", "print(f\"Testing images: {len(test_dataset)}\")\n", "print(f\"Detected Classes: {train_dataset.classes}\")\n", "\n", "# 2. Create the DataLoaders\n", "BATCH_SIZE = 64\n", "\n", "# We shuffle the training data so the model doesn't learn sequence patterns\n", "train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True)\n", "\n", "# We do not need to shuffle validation/testing data\n", "val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2, pin_memory=True)\n", "test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2, pin_memory=True)\n", "\n", "print(\"DataLoaders are ready!\")" ], "metadata": { "id": "I6obJOa1QYMy", "outputId": "3a3b35dc-01e5-4980-bdfc-0958467a5f24", "colab": { "base_uri": "https://localhost:8080/" } }, "execution_count": 3, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Training images: 5216\n", "Validation images: 16\n", "Testing images: 624\n", "Detected Classes: ['NORMAL', 'PNEUMONIA']\n", "DataLoaders are ready!\n" ] } ] }, { "cell_type": "markdown", "source": [ "## **Model Architecture and Class Weighting**\n", "\n", "Instead of training a model from scratch, we use **Transfer Learning**. We will download DenseNet121 (which already knows how to detect edges, textures, and shapes from millions of standard images) and replace its final classification layer with a custom one designed specifically for our two classes (Normal vs. Pneumonia).\n", "\n", "Simultaneously, we will calculate the inverse frequency of our classes to generate **Class Weights**, injecting them directly into the PyTorch loss function to penalize the model more heavily if it misclassifies a minority \"Normal\" scan." ], "metadata": { "id": "MPPAHhh2nlVz" } }, { "cell_type": "code", "source": [ "# --- 1. Handling the Imbalance (Class Weights) ---\n", "# ImageFolder automatically stores the label for every image in the .targets attribute\n", "labels = train_dataset.targets\n", "class_counts = np.bincount(labels)\n", "total_samples = len(labels)\n", "num_classes = len(train_dataset.classes)\n", "\n", "# Inverse frequency formula: Weight = Total / (NumClasses * ClassCount)\n", "weights = total_samples / (num_classes * class_counts)\n", "class_weights = torch.FloatTensor(weights).to(device)\n", "\n", "print(f\"Class Counts: NORMAL={class_counts[0]}, PNEUMONIA={class_counts[1]}\")\n", "print(f\"Applied Weights: NORMAL={weights[0]:.4f}, PNEUMONIA={weights[1]:.4f}\\n\")\n", "\n", "# --- 2. Model Architecture (DenseNet121) ---\n", "# Download the pre-trained neural network\n", "model = models.densenet121(weights=models.DenseNet121_Weights.IMAGENET1K_V1)\n", "\n", "# Freeze the core convolutional layers so we don't destroy the pre-trained visual features\n", "for param in model.parameters():\n", " param.requires_grad = False\n", "\n", "# Replace the final classification layer (originally 1000 classes) with our custom 2-class head\n", "num_ftrs = model.classifier.in_features\n", "\n", "# Push the entire model architecture to the GPU\n", "model = model.to(device)\n", "\n", "# --- 3. Loss Function and Optimizer ---\n", "# Inject our calculated class weights into the loss function\n", "criterion = nn.CrossEntropyLoss(weight=class_weights)\n", "\n", "# We only want the optimizer to update our newly added classifier layers\n", "optimizer = optim.Adam(model.classifier.parameters(), lr=0.001)\n", "\n", "print(\"Model architecture, optimizer, and weighted loss function configured successfully!\")" ], "metadata": { "id": "3cTqFqcvnVG2", "outputId": "19533754-f990-4485-b848-d9c2be3f782b", "colab": { "base_uri": "https://localhost:8080/" } }, "execution_count": 7, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Class Counts: NORMAL=1341, PNEUMONIA=3875\n", "Applied Weights: NORMAL=1.9448, PNEUMONIA=0.6730\n", "\n", "Model architecture, optimizer, and weighted loss function configured successfully!\n" ] } ] }, { "cell_type": "code", "source": [], "metadata": { "id": "HrV6nUsjodCM" }, "execution_count": null, "outputs": [] } ] }