identify_mislabeled_data.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Identify mislabeled data\n",
    "\n",
    "In this notebook, we see how AUM Ranking can be used to identify mislabeled data.\n",
    "\n",
    "First, we artificially corrupt the labels of CIFAR-10 to simulate a dataset with labeling errors. Then, we assign threshold samples, train a model and measure the AUM for each sample as per the AUM paper. We use AUM Ranking to decide which samples are likely to be mislabeled. We then repeat the process with a disjoint assignment of threshold samples so that all training samples are flagged as potentially mislabeled or not. Finally, we measure the precision and accuracy to see how well the process identifies the corrupted data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "import os\n",
    "import random\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from torcheval.metrics import Mean, MulticlassAccuracy\n",
    "import torchvision\n",
    "from tqdm import tqdm\n",
    "\n",
    "from aum_ranking import assign_threshold_samples, ThresholdSamplesDataset, AUM, compute_aum_threshold, flag_mislabeled_samples, combine_mislabeled_sample_flags\n",
    "from models import ResNet32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define how to corrupt the training dataset labels\n",
    "\n",
    "def assign_corrupted_samples(dataset, fraction):\n",
    "    corrupted_flags = np.zeros((len(dataset),), dtype=np.uint8)\n",
    "    num_corrupted_samples = round(fraction * len(dataset))\n",
    "    corrupted_flags[:num_corrupted_samples] = 1\n",
    "    np.random.shuffle(corrupted_flags)\n",
    "\n",
    "    real_labels = np.array([y for x, y in dataset])\n",
    "    noise = np.random.randint(1, len(dataset.classes), size=len(dataset))\n",
    "    fake_labels = (real_labels + noise) % len(dataset.classes)\n",
    "    corrupted_labels = np.where(corrupted_flags, fake_labels, real_labels)\n",
    "    return corrupted_labels, corrupted_flags\n",
    "\n",
    "\n",
    "class CorruptedDataset(Dataset):\n",
    "    \"\"\"\n",
    "    A Dataset wrapper which synthetically mislabels some of the data.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, dataset, labels):\n",
    "        self.dataset = dataset\n",
    "        self.labels = labels\n",
    "        self.classes = dataset.classes\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        x, y = self.dataset[index]\n",
    "        return x, self.labels[index]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.dataset)\n",
    "\n",
    "\n",
    "def corrupt_dataset(dataset, fraction):\n",
    "    corrupted_labels, corrupted_flags = assign_corrupted_samples(dataset, fraction)\n",
    "    return CorruptedDataset(dataset, corrupted_labels), corrupted_flags"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we define datasets based on CIFAR-10 and artificially corrupt a fraction of the training set labels. We keep track of which samples were corrupted so we can measure the precision and recall at the end.\n",
    "\n",
    "To try a different labeling error rate, change `labeling_error_fraction = 0.2` to a different value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define datasets\n",
    "\n",
    "pixel_means = torch.tensor([0.4914, 0.4822, 0.4465])\n",
    "pixel_stds = torch.tensor([0.2470, 0.2435, 0.2616])\n",
    "\n",
    "preprocess = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.ToTensor(),\n",
    "    torchvision.transforms.Normalize(mean=pixel_means, std=pixel_stds)\n",
    "])\n",
    "\n",
    "# This will be used later for visualizing images\n",
    "unprocess = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.Normalize(mean=-pixel_means / pixel_stds, std=1.0 / pixel_stds),\n",
    "    torchvision.transforms.ToPILImage()\n",
    "])\n",
    "\n",
    "train_dataset = torchvision.datasets.CIFAR10(root=\"data\", train=True, download=True, transform=torchvision.transforms.Compose([\n",
    "    preprocess,\n",
    "    torchvision.transforms.Pad(4),\n",
    "    torchvision.transforms.RandomHorizontalFlip(),\n",
    "    torchvision.transforms.RandomCrop(32)\n",
    "]))\n",
    "\n",
    "labeling_error_fraction = 0.2  # 20%\n",
    "train_dataset, corrupted_flags = corrupt_dataset(train_dataset, labeling_error_fraction)\n",
    "\n",
    "test_dataset = torchvision.datasets.CIFAR10(root=\"data\", train=False, download=True, transform=preprocess)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we assign threshold samples and apply AUM Ranking, we do not know whether the threshold samples are mislabeled or not. So, we make a second assignment of threshold samples, disjoint from the first, and apply the process again.\n",
    "\n",
    "We make both threshold sample assignments in one go, because it is easier and more efficient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assign two sets of threshold samples\n",
    "\n",
    "threshold_sample_flags_1, threshold_sample_flags_2 = assign_threshold_samples(num_samples=len(train_dataset), num_classes=len(train_dataset.classes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Number of classes: {len(train_dataset.classes)}\")\n",
    "print(f\"Number of samples: {len(train_dataset)}\")\n",
    "print(f\"Number of threshold samples (first pass): {threshold_sample_flags_1.sum()}\")\n",
    "print(f\"Number of threshold samples (second pass): {threshold_sample_flags_2.sum()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next few cells define functions for training and diagnostics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set device\n",
    "\n",
    "device = (\n",
    "    \"cuda\" if torch.cuda.is_available()\n",
    "    else \"mps\" if torch.backends.mps.is_available()\n",
    "    else \"cpu\"\n",
    ")\n",
    "print(f\"Using {device} device\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper functions for visualizing predictions\n",
    "\n",
    "def predict_with_probs(model, x):\n",
    "    \"\"\"\n",
    "    Generates predictions and corresponding probabilities from a trained\n",
    "    network and a list of images\n",
    "    \"\"\"\n",
    "    x = x.to(device)\n",
    "    logits = model(x)\n",
    "    output = F.softmax(logits, dim=1)\n",
    "    probs, preds = torch.max(output, 1)\n",
    "    probs = probs.cpu().numpy()\n",
    "    preds = preds.cpu().numpy()\n",
    "    return preds, probs\n",
    "\n",
    "\n",
    "def plot_classes_preds(model, x, y, classes):\n",
    "    \"\"\"\n",
    "    Generates matplotlib Figure using a trained network, along with images\n",
    "    and labels from a batch, that shows the network's top prediction along\n",
    "    with its probability, alongside the actual label, coloring this\n",
    "    information based on whether the prediction was correct or not.\n",
    "    Uses the \"images_to_probs\" function.\n",
    "    \"\"\"\n",
    "    with torch.no_grad():\n",
    "        preds, probs = predict_with_probs(model, x)\n",
    "    # Plot the images in the batch, along with predicted and true labels\n",
    "    fig, axs = plt.subplots(1, 4, figsize=(10, 3))\n",
    "    for i, ax in enumerate(axs):\n",
    "        ax.xaxis.set_visible(False)\n",
    "        ax.yaxis.set_visible(False)\n",
    "        ax.imshow(unprocess(x[i]))\n",
    "        ax.set_title(\n",
    "            \"{0}, {1:.1%}\\n(actual: {2})\".format(\n",
    "                classes[preds[i]],\n",
    "                probs[i],\n",
    "                classes[y[i]]),\n",
    "            color=(\"green\" if preds[i]==y[i].item() else \"red\")\n",
    "        )\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the training loop, we accumulate the AUM values with `aum.update`.\n",
    "\n",
    "`indexes` is passed in to `aum.update` to address the fact that the training set gets shuffled in each epoch.\n",
    "\n",
    "`progress` is a `tqdm` object showing a progress bar. We reuse the progress bar between epochs to minimize the amount of cell output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(loader, progress, model, loss_fn, optimizer, epoch, epochs, writer, loss_metric, accuracy_metric, aum):\n",
    "    progress.reset()\n",
    "    progress.desc = f\"Epoch {epoch+1}/{epochs}\"\n",
    "\n",
    "    model.train()\n",
    "    for x, y, indexes in loader:\n",
    "        x = x.to(device)\n",
    "        y = y.to(device)\n",
    "\n",
    "        pred = model(x)\n",
    "        loss = loss_fn(pred, y)\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        loss_metric.update(loss.detach())\n",
    "        accuracy_metric.update(pred, y)\n",
    "\n",
    "        progress.set_postfix_str(\n",
    "            f\"loss={loss_metric.compute().item():.4f}, accuracy={accuracy_metric.compute().item():.2%}\",\n",
    "            refresh=False\n",
    "        )\n",
    "\n",
    "        aum.update(pred, y, indexes)\n",
    "\n",
    "        progress.update()\n",
    "\n",
    "    writer.add_scalar(\"loss/train\", scalar_value=loss_metric.compute(), global_step=epoch)\n",
    "    loss_metric.reset()\n",
    "    writer.add_scalar(\"accuracy/train\", scalar_value=accuracy_metric.compute(), global_step=epoch)\n",
    "    accuracy_metric.reset()\n",
    "\n",
    "    progress.refresh()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(loader, model, loss_fn, epoch, writer, loss_metric, accuracy_metric, classes):\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for x, y in loader:\n",
    "            x = x.to(device)\n",
    "            y = y.to(device)\n",
    "\n",
    "            pred = model(x)\n",
    "            loss = loss_fn(pred, y)\n",
    "\n",
    "            loss_metric.update(loss)\n",
    "            accuracy_metric.update(pred, y)\n",
    "    \n",
    "    writer.add_scalar(\"loss/test\", scalar_value=loss_metric.compute(), global_step=epoch)\n",
    "    loss_metric.reset()\n",
    "    writer.add_scalar(\"accuracy/test\", scalar_value=accuracy_metric.compute(), global_step=epoch)\n",
    "    accuracy_metric.reset()\n",
    "\n",
    "    x, y = zip(*random.choices(loader.dataset, k=4))\n",
    "    x = torch.stack(x)\n",
    "    y = torch.tensor(y)\n",
    "    writer.add_figure(\"predictions\", plot_classes_preds(model, x, y, classes), global_step=epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_on_dataset(model, loader):\n",
    "    preds = np.empty(len(loader.dataset), dtype=np.int32)\n",
    "    probs = np.empty(len(loader.dataset), dtype=np.int32)\n",
    "    \n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for x, y, indexes in loader:\n",
    "            pred, prob = predict_with_probs(model, x)\n",
    "            preds[indexes] = pred\n",
    "            probs[indexes] = prob\n",
    "\n",
    "    return preds, probs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we do a training run on `threshold_dataset`, a specific assignment of threshold samples.\n",
    "\n",
    "We also collect some \"suggested labels\" from the model. This is not part of AUM Ranking, but it makes the visualization at the end nicer. In practice, it could be used to pseudo-label the mislabeled samples instead of removing them from the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def training_run(pass_index, threshold_dataset, test_dataset, epochs, aum):\n",
    "    # Make data loaders\n",
    "    batch_size = 64\n",
    "    threshold_loader = DataLoader(threshold_dataset, batch_size=batch_size, shuffle=True)\n",
    "    test_loader = DataLoader(test_dataset, batch_size=batch_size)\n",
    "\n",
    "    # Build model\n",
    "    model = ResNet32(num_classes=len(threshold_dataset.classes)).to(device)\n",
    "\n",
    "    # Set loss function and optimizer\n",
    "    loss_fn = torch.nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, weight_decay=1e-4, momentum=0.9, nesterov=True)\n",
    "\n",
    "    # Make TensorBoard writer\n",
    "    now = datetime.now().strftime(\"%b%d_%H-%M-%S\")\n",
    "    log_dir = os.path.join(\"runs\", \"identify_mislabeled_data\", f\"{now}_pass_{pass_index}\")\n",
    "    writer = SummaryWriter(log_dir=log_dir)\n",
    "\n",
    "    # Define metrics\n",
    "    loss_metric = Mean(device=device)\n",
    "    accuracy_metric = MulticlassAccuracy(device=device)\n",
    "\n",
    "    # Train the model\n",
    "    print(\"Training model\")\n",
    "    print(\"-------------------------------\")\n",
    "    with tqdm(total=len(threshold_loader)) as train_progress:\n",
    "        for epoch in range(epochs):\n",
    "            train(threshold_loader, train_progress, model, loss_fn, optimizer, epoch, epochs, writer, loss_metric, accuracy_metric, aum)\n",
    "            test(test_loader, model, loss_fn, epoch, writer, loss_metric, accuracy_metric, classes=threshold_loader.dataset.classes)\n",
    "    suggested_labels, suggested_label_probs = predict_on_dataset(model, threshold_loader)\n",
    "    print(\"-------------------------------\")\n",
    "    print(\"Done!\")\n",
    "\n",
    "    return model, suggested_labels, suggested_label_probs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These functions collect together all the steps of AUM Ranking. We assign threshold samples, train a model, compute AUM values, then flag them as potentially mislabeled or not based on the threshold samples. Repeating the process, we flag a second subset of samples as potentially mislabeled or not. Finally, we combine the two subsets of flags by OR-ing them together.\n",
    "\n",
    "As an additional convenience, we compute \"suggested labels\" from both models and combine them.\n",
    "\n",
    "To reduce notebook running time, change `epochs = 150` to a smaller value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def identify_mislabeled_samples(pass_index, threshold_sample_flags, train_dataset, test_dataset):\n",
    "    print(f\"Performing pass {pass_index}\")\n",
    "\n",
    "    # Make threshold samples dataset\n",
    "    threshold_dataset = ThresholdSamplesDataset(train_dataset, threshold_sample_flags)\n",
    "\n",
    "    # Create AUM calculator\n",
    "    aum = AUM(num_samples=len(threshold_dataset))\n",
    "\n",
    "    # Train a model to populate the margin values\n",
    "    epochs = 150\n",
    "    model, suggested_labels, suggested_label_probs = training_run(pass_index, threshold_dataset, test_dataset, epochs, aum)\n",
    "\n",
    "    # Compute AUM values\n",
    "    aum_values = aum.compute(epochs).numpy()\n",
    "    print(f\"AUM values: {aum_values.shape}, {aum_values.dtype}\")\n",
    "    print(f\"mean: {np.mean(aum_values):.4f}, min: {np.min(aum_values):.4f}, max: {np.max(aum_values):.4f}, std: {np.std(aum_values):.4f}\")\n",
    "\n",
    "    # Compute AUM threshold\n",
    "    aum_threshold = compute_aum_threshold(aum_values, threshold_sample_flags)\n",
    "    print(f\"AUM threshold: {aum_threshold}\")\n",
    "\n",
    "    # Flag (potentially) mislabeled samples\n",
    "    mislabeled_sample_flags = flag_mislabeled_samples(aum_values, threshold_sample_flags, aum_threshold)\n",
    "    print(f\"Potentially mislabeled samples: {np.sum(mislabeled_sample_flags)}\")\n",
    "    print(f\"Finished pass {pass_index}\")\n",
    "    print(\"===============================\")\n",
    "    return mislabeled_sample_flags, suggested_labels, suggested_label_probs\n",
    "\n",
    "\n",
    "# We combine suggested labels based on the most confident prediction. Compared to averaging the two sets of probabilities,\n",
    "# this avoids needing to store all class probabilities for all samples - we can just store the top one.\n",
    "def combine_suggested_labels(suggested_labels_1, suggested_label_probs_1, suggested_labels_2, suggested_label_probs_2):\n",
    "    choose_1 = suggested_label_probs_1 >= suggested_label_probs_2\n",
    "    suggested_labels = np.where(choose_1, suggested_labels_1, suggested_labels_2)\n",
    "    return suggested_labels\n",
    "\n",
    "\n",
    "def identify_and_combine_mislabeled_samples(train_dataset, test_dataset, threshold_sample_flags_1, threshold_sample_flags_2):\n",
    "    mislabeled_sample_flags_1, suggested_labels_1, suggested_label_probs_1 = identify_mislabeled_samples(1, threshold_sample_flags_1, train_dataset, test_dataset)\n",
    "    mislabeled_sample_flags_2, suggested_labels_2, suggested_label_probs_2 = identify_mislabeled_samples(2, threshold_sample_flags_2, train_dataset, test_dataset)\n",
    "    mislabeled_sample_flags = combine_mislabeled_sample_flags(mislabeled_sample_flags_1, mislabeled_sample_flags_2)\n",
    "    suggested_labels = combine_suggested_labels(suggested_labels_1, suggested_label_probs_1, suggested_labels_2, suggested_label_probs_2)\n",
    "    return mislabeled_sample_flags, suggested_labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Applying the above functions, we get mislabeled sample flags (and suggested labels) for all training samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mislabeled_sample_flags, suggested_labels = identify_and_combine_mislabeled_samples(train_dataset, test_dataset, threshold_sample_flags_1, threshold_sample_flags_2)\n",
    "print(f\"Mislabeled sample flags: {mislabeled_sample_flags.shape}, {mislabeled_sample_flags.dtype}\")\n",
    "print(f\"Potentially mislabeled samples: {np.sum(mislabeled_sample_flags)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we plot some training samples identified as mislabeled, along with their labels in the corrupted dataset and the labels suggested by the models. Re-run the cell to see more images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_mislabeled_samples(dataset, mislabeled_sample_flags, suggested_labels):\n",
    "    (mislabeled_sample_indexes,) = np.nonzero(mislabeled_sample_flags)\n",
    "    indexes = np.random.choice(mislabeled_sample_indexes, 6)\n",
    "    fig, axs = plt.subplots(1, 6, figsize=(15, 3.6))\n",
    "    fig.suptitle(\"Potentially mislabeled samples\")\n",
    "    for i, ax in zip(indexes, axs):\n",
    "        x, y = dataset[i]\n",
    "        ax.xaxis.set_visible(False)\n",
    "        ax.yaxis.set_visible(False)\n",
    "        ax.imshow(unprocess(x))\n",
    "        ax.set_title(f\"{i}\\n{dataset.classes[y]}\\n(suggested: {dataset.classes[suggested_labels[i]]})\")\n",
    "\n",
    "\n",
    "plot_mislabeled_samples(train_dataset, mislabeled_sample_flags, suggested_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We calculate the precision and recall of identifying mislabeled data using AUM Ranking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mislabeled_true_positives = np.count_nonzero(mislabeled_sample_flags & corrupted_flags)\n",
    "mislabeled_false_positives = np.count_nonzero(mislabeled_sample_flags & ~corrupted_flags)\n",
    "mislabeled_false_negatives = np.count_nonzero(~mislabeled_sample_flags & corrupted_flags)\n",
    "\n",
    "precision = mislabeled_true_positives / (mislabeled_true_positives + mislabeled_false_positives)\n",
    "recall = mislabeled_true_positives / (mislabeled_true_positives + mislabeled_false_negatives)\n",
    "\n",
    "print(f\"Mislabeled sample identification\")\n",
    "print(f\"precision: {precision:.2%}\")\n",
    "print(f\"recall:    {recall:.2%}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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": 2
}