examples/07_Optimizing_Multi_Label_AUROC_Loss_with_DenseNet121_on_CheXpert.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9Qkb6bYy3rOx"
      },
      "source": [
        "# **Optimizing Multi-label AUROC loss on Chest X-Ray Dataset (CheXpert)**\n",
        "\n",
        "**Author**: Zhuoning Yuan\n",
        "\n",
        "**Introduction**\n",
        "\n",
        "In this tutorial, you will learn how to quickly train a DenseNet121 model by optimizing AUROC using our novel AUCMLoss and PESG optimizer on Chest X-Ray dataset, e.g.,[CheXpert](https://stanfordmlgroup.github.io/competitions/chexpert/). After completion of this tutorial, you should be able to use LibAUC to train your own models on your own datasets.\n",
        "\n",
        "\n",
        "\n",
        "**Useful Resources**:\n",
        "* Website: https://libauc.org\n",
        "* Github: https://github.com/Optimization-AI/LibAUC\n",
        "\n",
        "**Reference**:  \n",
        "\n",
        "If you find this tutorial helpful in your work,  please acknowledge our library and cite the following paper:\n",
        "\n",
        "<pre>\n",
        "@inproceedings{yuan2021large,\n",
        "  title={Large-scale robust deep auc maximization: A new surrogate loss and empirical studies on medical image classification},\n",
        "  author={Yuan, Zhuoning and Yan, Yan and Sonka, Milan and Yang, Tianbao},\n",
        "  booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision},\n",
        "  pages={3040--3049},\n",
        "  year={2021}\n",
        "}\n",
        "</pre>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hTJ3ca0u4YQ4"
      },
      "source": [
        "# **Installing LibAUC**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h8iVw1kU3guh"
      },
      "outputs": [],
      "source": [
        "!pip install libauc==1.2.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nlD-4SrE4dVW"
      },
      "source": [
        "# **Downloading CheXpert**\n",
        " \n",
        "*   To request dataset access, you need to apply from CheXpert website: https://stanfordmlgroup.github.io/competitions/chexpert/\n",
        "*   In this tutorial, we use the smaller version of dataset with lower image resolution, i.e., *CheXpert-v1.0-small.zip*\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CcsJ4eoj3VST"
      },
      "outputs": [],
      "source": [
        "!cp /content/drive/MyDrive/chexpert-dataset/CheXpert-v1.0-small.zip /content/\n",
        "!mkdir CheXpert\n",
        "!unzip CheXpert-v1.0-small.zip -d /content/CheXpert/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OVvrt3ku4qpq"
      },
      "source": [
        "\n",
        "# **Importing LibAUC**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XGHWer3v4qJo"
      },
      "outputs": [],
      "source": [
        "from libauc.losses import AUCM_MultiLabel, CrossEntropyLoss\n",
        "from libauc.optimizers import PESG, Adam\n",
        "from libauc.models import densenet121 as DenseNet121\n",
        "from libauc.datasets import CheXpert\n",
        "from libauc.metrics import auc_roc_score # for multi-task\n",
        "\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "import torch \n",
        "import torchvision.transforms as transforms\n",
        "from torch.utils.data import Dataset\n",
        "import torch.nn.functional as F   "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l2swK5Mo7Kca"
      },
      "source": [
        "# **Reproducibility**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OiiT5oEp7J3C"
      },
      "outputs": [],
      "source": [
        "def set_all_seeds(SEED):\n",
        "    # REPRODUCIBILITY\n",
        "    torch.manual_seed(SEED)\n",
        "    np.random.seed(SEED)\n",
        "    torch.backends.cudnn.deterministic = True\n",
        "    torch.backends.cudnn.benchmark = False"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8Bjd5Q5wkAT7"
      },
      "source": [
        "# **Datasets, Loss and Optimizer**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KkxJ2ZNNj_uN",
        "outputId": "93c1f0c9-c94f-4a39-dcba-e22e647c606f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[0.12241724991755092, 0.32190737435022276, 0.06796421448276946, 0.31190878776298636, 0.402555659671146]\n"
          ]
        }
      ],
      "source": [
        "root = './CheXpert/CheXpert-v1.0-small/'\n",
        "# Index=-1 denotes multi-label with 5 diseases\n",
        "traindSet = CheXpert(csv_path=root+'train.csv', image_root_path=root, use_upsampling=False, use_frontal=True, image_size=224, mode='train', class_index=-1, verbose=False)\n",
        "testSet =  CheXpert(csv_path=root+'valid.csv',  image_root_path=root, use_upsampling=False, use_frontal=True, image_size=224, mode='valid', class_index=-1, verbose=False)\n",
        "trainloader =  torch.utils.data.DataLoader(traindSet, batch_size=32, num_workers=2, shuffle=True)\n",
        "testloader =  torch.utils.data.DataLoader(testSet, batch_size=32, num_workers=2, shuffle=False)\n",
        "\n",
        "# check imbalance ratio for each task\n",
        "print (traindSet.imratio_list )\n",
        "\n",
        "# paramaters\n",
        "SEED = 123\n",
        "BATCH_SIZE = 32\n",
        "lr = 0.1 \n",
        "epoch_decay = 2e-3\n",
        "weight_decay = 1e-5\n",
        "margin = 1.0\n",
        "total_epochs = 2\n",
        "\n",
        "# model\n",
        "set_all_seeds(SEED)\n",
        "model = DenseNet121(pretrained=True, last_activation=None, activations='relu', num_classes=5)\n",
        "model = model.cuda()\n",
        "\n",
        "# define loss & optimizer\n",
        "loss_fn = AUCM_MultiLabel(num_classes=5)\n",
        "optimizer = PESG(model, \n",
        "                 loss_fn=loss_fn,\n",
        "                 lr=lr, \n",
        "                 margin=margin, \n",
        "                 epoch_decay=epoch_decay, \n",
        "                 weight_decay=weight_decay)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G3Bge6KM7lBP"
      },
      "source": [
        "# **Multi-label Training**\n",
        "Optimizing Multi-label AUROC loss (e.g., 5 tasks)   \n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true,
          "base_uri": "https://localhost:8080/"
        },
        "id": "u6p0oVY8AouP",
        "outputId": "01c2b44b-7c08-4a5e-8209-fceb75aab393"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Start Training\n",
            "------------------------------\n",
            "Epoch=0, BatchID=0, Val_AUC=0.5558, Best_Val_AUC=0.5558\n",
            "Epoch=0, BatchID=400, Val_AUC=0.8283, Best_Val_AUC=0.8283\n",
            "Epoch=0, BatchID=800, Val_AUC=0.8074, Best_Val_AUC=0.8283\n",
            "Epoch=0, BatchID=1200, Val_AUC=0.8528, Best_Val_AUC=0.8528\n",
            "Epoch=0, BatchID=1600, Val_AUC=0.8337, Best_Val_AUC=0.8528\n",
            "Epoch=0, BatchID=2000, Val_AUC=0.8420, Best_Val_AUC=0.8528\n",
            "Epoch=0, BatchID=2400, Val_AUC=0.8589, Best_Val_AUC=0.8589\n",
            "Epoch=0, BatchID=2800, Val_AUC=0.8475, Best_Val_AUC=0.8589\n",
            "Epoch=0, BatchID=3200, Val_AUC=0.8702, Best_Val_AUC=0.8702\n",
            "Epoch=0, BatchID=3600, Val_AUC=0.8453, Best_Val_AUC=0.8702\n",
            "Epoch=0, BatchID=4000, Val_AUC=0.8552, Best_Val_AUC=0.8702\n",
            "Epoch=0, BatchID=4400, Val_AUC=0.8366, Best_Val_AUC=0.8702\n",
            "Epoch=0, BatchID=4800, Val_AUC=0.8603, Best_Val_AUC=0.8702\n",
            "Epoch=0, BatchID=5200, Val_AUC=0.8700, Best_Val_AUC=0.8702\n",
            "Epoch=0, BatchID=5600, Val_AUC=0.8842, Best_Val_AUC=0.8842\n",
            "Reducing learning rate to 0.01000 @ T=5970!\n",
            "Updating regularizer @ T=5970!\n",
            "Epoch=1, BatchID=0, Val_AUC=0.8587, Best_Val_AUC=0.8842\n",
            "Epoch=1, BatchID=400, Val_AUC=0.8875, Best_Val_AUC=0.8875\n",
            "Epoch=1, BatchID=800, Val_AUC=0.8899, Best_Val_AUC=0.8899\n",
            "Epoch=1, BatchID=1200, Val_AUC=0.8813, Best_Val_AUC=0.8899\n",
            "Epoch=1, BatchID=1600, Val_AUC=0.8871, Best_Val_AUC=0.8899\n",
            "Epoch=1, BatchID=2000, Val_AUC=0.8893, Best_Val_AUC=0.8899\n",
            "Epoch=1, BatchID=2400, Val_AUC=0.8685, Best_Val_AUC=0.8899\n",
            "Epoch=1, BatchID=2800, Val_AUC=0.8837, Best_Val_AUC=0.8899\n",
            "Epoch=1, BatchID=3200, Val_AUC=0.8817, Best_Val_AUC=0.8899\n",
            "Epoch=1, BatchID=3600, Val_AUC=0.8868, Best_Val_AUC=0.8899\n",
            "Epoch=1, BatchID=4000, Val_AUC=0.8857, Best_Val_AUC=0.8899\n",
            "Epoch=1, BatchID=4400, Val_AUC=0.8804, Best_Val_AUC=0.8899\n",
            "Epoch=1, BatchID=4800, Val_AUC=0.8919, Best_Val_AUC=0.8919\n",
            "Epoch=1, BatchID=5200, Val_AUC=0.8929, Best_Val_AUC=0.8929\n",
            "Epoch=1, BatchID=5600, Val_AUC=0.8809, Best_Val_AUC=0.8929\n"
          ]
        }
      ],
      "source": [
        "# training\n",
        "print ('Start Training')\n",
        "print ('-'*30)\n",
        "\n",
        "best_val_auc = 0 \n",
        "for epoch in range(total_epochs):\n",
        "    if epoch > 0:\n",
        "        optimizer.update_regularizer(decay_factor=10)    \n",
        "\n",
        "    for idx, data in enumerate(trainloader):\n",
        "      train_data, train_labels = data\n",
        "      train_data, train_labels  = train_data.cuda(), train_labels.cuda()\n",
        "      y_pred = model(train_data)\n",
        "      y_pred = torch.sigmoid(y_pred)\n",
        "      loss = loss_fn(y_pred, train_labels)\n",
        "      optimizer.zero_grad()\n",
        "      loss.backward()\n",
        "      optimizer.step()\n",
        "        \n",
        "      # validation  \n",
        "      if idx % 400 == 0:\n",
        "         model.eval()\n",
        "         with torch.no_grad():    \n",
        "              test_pred = []\n",
        "              test_true = [] \n",
        "              for jdx, data in enumerate(testloader):\n",
        "                  test_data, test_labels = data\n",
        "                  test_data = test_data.cuda()\n",
        "                  y_pred = model(test_data)\n",
        "                  y_pred = torch.sigmoid(y_pred)\n",
        "                  test_pred.append(y_pred.cpu().detach().numpy())\n",
        "                  test_true.append(test_labels.numpy())\n",
        "            \n",
        "              test_true = np.concatenate(test_true)\n",
        "              test_pred = np.concatenate(test_pred)\n",
        "              val_auc_mean = np.mean(auc_roc_score(test_true, test_pred)) \n",
        "              model.train()\n",
        "\n",
        "              if best_val_auc < val_auc_mean:\n",
        "                 best_val_auc = val_auc_mean\n",
        "                 torch.save(model.state_dict(), 'aucm_pretrained_model.pth')\n",
        "\n",
        "              print ('Epoch=%s, BatchID=%s, Val_AUC=%.4f, Best_Val_AUC=%.4f'%(epoch, idx, val_auc_mean, best_val_auc))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tFKAu80J1vzD"
      },
      "source": [
        "# **Evaluation**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F-TTAuL51u4e"
      },
      "outputs": [],
      "source": [
        "# show auc roc scores for each task \n",
        "auc_roc_score(test_true, test_pred)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "07_Optimizing_Multi_Label_AUROC_Loss_with_DenseNet121_on_CheXpert.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}