datalab_image_classification/train_image_classifier.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Image Classifier Training on Caltech-256 Subset\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/cleanlab/examples/blob/master/datalab_image_classification/train_image_classifier.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we train a Swin Transformer model for image classification on a subset of the Caltech-256 dataset using the Timm library and Pytorch.\n",
    "\n",
    "We train the model with K-fold cross-validation and use it to produce out-of-sample predicted class probabilities for each image in our dataset, as well as a feature embedding of each image.\n",
    "\n",
    "Please install the dependencies specified in this [requirements-train.txt](https://github.com/cleanlab/examples/blob/master/datalab_image_classification/requirements-train.txt) file before running the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision.transforms as transforms\n",
    "import numpy as np\n",
    "from torchvision.datasets import ImageFolder\n",
    "from tqdm import tqdm\n",
    "import timm\n",
    "from sklearn.model_selection import StratifiedKFold"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset Preparation\n",
    "\n",
    "We'll download the dataset to disk and load it with a Torchvision data loader (applying necessary transformations, which we define below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -nc 'https://cleanlab-public.s3.amazonaws.com/Datalab/caltech256-subset.tar.gz'\n",
    "!mkdir -p data\n",
    "!tar -xf caltech256-subset.tar.gz -C data/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_to_rgb(img):\n",
    "    if img.mode == 'L':\n",
    "        img = img.convert('RGB')\n",
    "    return img\n",
    "\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.Lambda(convert_to_rgb),\n",
    "    transforms.Resize((224, 224)),\n",
    "    transforms.CenterCrop(224),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(  # ImageNet stats for normalization (mean and std) of RGB channels\n",
    "        mean=[0.485, 0.456, 0.406],\n",
    "        std=[0.229, 0.224, 0.225]\n",
    "    )\n",
    "])\n",
    "\n",
    "DATA_DIR = './data' # Save notebook artifacts in this directory\n",
    "\n",
    "# Load data from disk\n",
    "dataset = ImageFolder(\n",
    "    os.path.join(DATA_DIR, \"caltech256-subset\"),\n",
    "    transform=transform\n",
    ")\n",
    "n_classes = len(dataset.classes)\n",
    "labels = np.array(dataset.targets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Training\n",
    "\n",
    "With the dataset ready, we now define the hyperparameters for training and perform k-fold cross-validation to train a Swin Transformer.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define hyperparameters\n",
    "\n",
    "batch_size = 32 # Resnet50: 64, Swin-Transformer-patch-4-window-7-224: 32\n",
    "learning_rate = 0.00001 # Resnet50: 0.001, Swin-Transformer-patch-4-window-7-224: 0.0001\n",
    "num_epochs = 10\n",
    "num_folds = 5  # Use 3 for faster training\n",
    "patience = 2\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model_name = \"swin_base_patch4_window7_224\"\n",
    "model_prefix = \"caltech256_subset_\" + model_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### K-Fold Cross-Validation and Training\n",
    "\n",
    "In this notebook, we use k-fold cross-validation to train the model and extract out-of-sample predicted probabilities for all data points.\n",
    "\n",
    "During the training process, we'll just use the validation accuracy on the held-out fold to allow early stopping for each fold. This approach helps us prevent overfitting and obtain a better estimate of the model's performance.\n",
    "\n",
    "While we're not specifically interested in the model artifacts themselves, we aim to get a general idea of whether the chosen model architecture is accurate enough for our purpose.\n",
    "\n",
    "**Warning**: This cell may take a long time to execute and should be run with a GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kf = StratifiedKFold(n_splits=num_folds, shuffle=True, random_state=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train model\n",
    "\n",
    "# Set seed\n",
    "seed = 42\n",
    "torch.manual_seed(seed)\n",
    "np.random.seed(seed)\n",
    "for fold, (train_idx, test_idx) in enumerate(kf.split(dataset, labels)):\n",
    "    print(f'Fold {fold + 1}/{num_folds}')\n",
    "    print('-' * 10)\n",
    "\n",
    "    # Define data loaders for current fold\n",
    "    train_subset = torch.utils.data.Subset(dataset, train_idx)\n",
    "    val_subset = torch.utils.data.Subset(dataset, test_idx)\n",
    "    # Print train and validation set sizes\n",
    "    print(f'Train set size: {len(train_idx)}')\n",
    "    print(f'Test set size: {len(test_idx)}')\n",
    "\n",
    "    train_loader = torch.utils.data.DataLoader(train_subset, batch_size=batch_size, shuffle=True)\n",
    "    val_loader = torch.utils.data.DataLoader(val_subset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    # Initialize model for current fold\n",
    "    model = timm.create_model(model_name, pretrained=True, num_classes=n_classes)\n",
    "    model = model.to(device)\n",
    "    num_features = model.num_features\n",
    "\n",
    "    # Define loss function and optimizer\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
    "\n",
    "    best_val_accuracy = 0\n",
    "    best_epoch = 0\n",
    "    # Train model for current fold\n",
    "    for epoch in range(num_epochs):\n",
    "        print(f'Epoch {epoch + 1}/{num_epochs}')\n",
    "        model.train()\n",
    "        for inputs, targets in tqdm(train_loader):\n",
    "            inputs = inputs.to(device)\n",
    "            targets = targets.to(device)\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, targets)\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        print(f'Train loss: {loss:.4f}')\n",
    "\n",
    "        # Evaluate model on training set for current fold\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            correct = 0\n",
    "            total = 0\n",
    "            eval_loader = train_loader\n",
    "            for inputs, targets in tqdm(val_loader):\n",
    "                inputs = inputs.to(device)\n",
    "                targets = targets.to(device)\n",
    "                outputs = model(inputs)\n",
    "                _, predicted = torch.max(outputs.data, 1)\n",
    "                total += targets.size(0)\n",
    "                correct += (predicted == targets).sum().item()\n",
    "            val_accuracy = 100 * correct / total\n",
    "            print(f'Validation accuracy: {val_accuracy:.2f}%')\n",
    "\n",
    "        # Save model checkpoint if it is the best so far\n",
    "        if val_accuracy > best_val_accuracy:\n",
    "            print('Saving model...')\n",
    "            path = f'{model_prefix}_fold_{fold + 1}.pt'\n",
    "            torch.save(model.state_dict(), path)\n",
    "            best_val_accuracy = val_accuracy\n",
    "            best_epoch = epoch\n",
    "\n",
    "        # Early stopping\n",
    "        if epoch - best_epoch > patience:\n",
    "            print(f'Early stopping at epoch {epoch + 1}')\n",
    "            break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting Predicted Class Probabilities and Extracting Feature Embeddings \n",
    "\n",
    "After training, we will compute predicted class probabilities for the entire dataset using the trained models from each fold.\n",
    "\n",
    "In addition, to keep things simple, we'll use the model trained on the first fold as a feature extractor to obtain embeddings for every image in the dataset.\n",
    "\n",
    "These artifacts will be used by `Datalab` to inspect the dataset for potential issues, so we save them to files used in the next notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = timm.create_model(model_name, pretrained=True, num_classes=n_classes)\n",
    "path = f'{model_prefix}_fold_1.pt'\n",
    "model.load_state_dict(torch.load(path))\n",
    "model.eval()\n",
    "model.to(device)\n",
    "num_features = model.num_features\n",
    "\n",
    "features = np.zeros((len(dataset),num_features))\n",
    "pred_probs = np.zeros((len(dataset), n_classes))\n",
    "\n",
    "for fold, (_, test_idx) in enumerate(kf.split(dataset, labels)):\n",
    "    # Save out-of-sample predictions and features for current fold\n",
    "    # This is the validation set\n",
    "    # Define data loaders for current fold\n",
    "    test_subset = torch.utils.data.Subset(dataset, test_idx)\n",
    "    test_loader = torch.utils.data.DataLoader(test_subset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "\n",
    "    model = timm.create_model(model_name, pretrained=True, num_classes=n_classes)\n",
    "    path = f'{model_prefix}_fold_{fold + 1}.pt'\n",
    "    model.load_state_dict(torch.load(path))\n",
    "\n",
    "    model.eval()\n",
    "    model.to(device)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        pred_probs_fold = []\n",
    "        for inputs, _ in tqdm(test_loader):\n",
    "            inputs = inputs.to(device)\n",
    "            outputs = model(inputs)\n",
    "            # Predicted probabilities\n",
    "            outputs = nn.functional.softmax(outputs, dim=1)\n",
    "            pred_probs_fold.append(outputs.cpu().numpy())\n",
    "        pred_probs[test_idx] = np.concatenate(pred_probs_fold, axis=0)\n",
    "        \n",
    "    model = timm.create_model(model_name, pretrained=True, num_classes=n_classes)\n",
    "    path = f'{model_prefix}_fold_1.pt'\n",
    "    model.load_state_dict(torch.load(path))\n",
    "    model.eval()\n",
    "    model.to(device)\n",
    "    with torch.no_grad():\n",
    "        features_fold = []\n",
    "        model.reset_classifier(0)\n",
    "        for inputs, _ in tqdm(test_loader):\n",
    "            inputs = inputs.to(device)\n",
    "            features_fold.append(model(inputs).cpu().numpy())\n",
    "        features[test_idx] = np.concatenate(features_fold, axis=0)\n",
    "\n",
    "features_path = os.path.join(DATA_DIR, \"features.npy\")\n",
    "pred_probs_path = os.path.join(DATA_DIR, \"pred_probs.npy\")\n",
    "\n",
    "np.save(features_path, features)\n",
    "np.save(pred_probs_path, pred_probs)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "949777d72b0d2535278d3dc13498b2535136f6dfe0678499012e853ee9abcab1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}