Session 2 -- Multiple Instance Learning.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "f4cc5a61",
      "metadata": {
        "id": "f4cc5a61"
      },
      "source": [
        "# Weakly-Supervised Deep Learning for Cancer Diagnosis in Computational Pathology\n",
        "\n",
        "- Presenter:\n",
        "   - Guillaume Jaume (gjaume@bwh.harvard.edu)\n",
        "   - Postdoctoral Researcher at Harvard Medical School and Brigham and Women's Hospital\n",
        "- Initally proposed and written by Richard J. Chen (richardchen@g.harvard.edu)\n",
        "\n",
        "![](https://user-images.githubusercontent.com/10300839/232984533-c22822b8-df80-4b95-80e2-93dde2409bbf.png)\n",
        "\n",
        "**Definitions:**\n",
        "\n",
        "- *Computational pathology (CPath):* Computational methods based on the microscopic analysis of cells and tissues for the study of disease.\n",
        "\n",
        "- *Digital pathology:* A set of tools and systems for the acquisition, management and diagnosis of pathology glass slides in a digital setting.\n",
        "\n",
        "- *Whole slide image (WSI):* An image obtained by digitizing a glass slide at high-resolution using a scanner.\n",
        "\n",
        "- *Hematoxylin and eosin (H&E) staining:* The reference stain for histological analysis of tissues for visualization of cell nuclei (in purple) with extracellular information and cytoplasm (in pink).\n",
        "\n",
        "\n",
        "**Background**:\n",
        "\n",
        "Computational Pathology aims to automate, assist, and augment the clinical practice of pathology using computational tools based on Artificial Intelligence.\n",
        "\n",
        "Tissue phenotyping is a fundamental problem in computational pathology (CPATH) to characterize histopathologic features for cancer diagnosis, prognosis, and prediction of treatment response. Unlike natural images, whole-slide imaging is a challenging computer vision domain in which image resolutions can be as large as $150{,}000 \\times 150{,}000$ pixels (>50 GB to load the entire image in RAM).\n",
        "\n",
        "To address this computational and memory bottleneck, the majority of state-of-the-art methods use a three-stage, weakly-supervised pipeline based on multiple instance learning (MIL):\n",
        "1. Tissue patching at a single magnification objective (\"zoom\"), e.g., 20x magnification\n",
        "\n",
        "2. Patch-level feature extraction to construct a set of patch embeddings (compress patches by a factor 100~500)\n",
        "\n",
        "3. Global pooling of embeddings to construct a slide-level representation for weak-supervision using slide-level labels (e.g., subtype, grade, stage, survival, origin).\n",
        "\n",
        "**Notebook Objective**: The following tutorial aims to distinguish Lung Adenocarcinoma (LUAD, 40% of all lung cancer) vs. Lung Squamous Cell Carcinoma (LUSC, 30% of all lung cancer) (see [Lu et al., Data-efficient and weakly supervised computational pathology on whole-slide images  Nature BME 2021](https://www.nature.com/articles/s41551-020-00682-w) and codebase [CLAM](https://github.com/mahmoodlab/CLAM). Specifically, we will:\n",
        "- Train and evaluate a \"naive\" MIL algorithm called `AverageMIL`, which takes the average of patch embeddings (as the global pooling operator).\n",
        "\n",
        "- Implement a more sophisticated algorithm called Attention-Based Multiple Instance Learning (`ABMIL`), which learns attention weights for computing a weighted average of patch embeddings.\n",
        "\n",
        "- Compare and contrast `AverageMIL` and `ABMIL`, discussing which algorithm performs better and potential limitations.\n",
        "\n",
        "**About this notebook**:\n",
        "- Model implementation and training is directly adapted from [CLAM](https://github.com/mahmoodlab/CLAM). CLAM includes many additional features (e.g. - letting users set up optimizers, model types, logging information, and other hyper-parameters) left out due to making this notebook as simple to run as possible for teaching purposes. To use all features, please see CLAM.\n",
        "\n",
        "- Though this notebook is based off of CLAM, the method-of-interest that you will be implementing is not CLAM, but a different method called ABMIL from [Ilse et al. Attention-Based Multiple Instance Learning ICML 2018](https://arxiv.org/abs/1802.04712), which CLAM is derived from.\n",
        "\n",
        "- Though pre-extracted features were generated using the CLAM codebase, the encoder was not a truncated ResNet-50 pretrained on ImageNet (dimension 1024) at 20 $\\times$ resolution. Instead, we extracted features with a much smaller CNN encoder (dimension 320) at 10 $\\times$ resolution, which shrinks the size of the dataset from ~11 GB to ~3.96 GB of storage (**download link for pre-extracted features in the cell below**). In addition, a torch.seed is set for reproducibility (all outputs should be deterministic)."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8f7ed950",
      "metadata": {
        "id": "8f7ed950"
      },
      "source": [
        "### Colab Installation, Data Download, & Dependencies\n",
        "\n",
        "- Gets clinical metadata csv for tcga-luad and tcga-lusc with predefined train/val/test splits\n",
        "- Gets pre-extracted features for tcga-luad and tcga-lusc diagnostic WSIs (1043 WSIs total, ~3.96 GB in size, ~67 seconds to download)\n",
        "\n",
        "Alternatively, you can download the data directly from Dropbox to your local computer, and run this Colab Notebook locally\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "use_drive = False\n",
        "\n",
        "if use_drive:\n",
        "  from google.colab import drive\n",
        "  drive.mount('/content/drive')\n",
        "  !mkdir -p \"/content/drive/My Drive/ai4healthsummerschool/\""
      ],
      "metadata": {
        "id": "h1fJd7bzLpZH"
      },
      "id": "h1fJd7bzLpZH",
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "QKGr936RiVFS",
      "metadata": {
        "id": "QKGr936RiVFS",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "02c12827-fced-4b80-d1fd-5ff1268ce768"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2023-07-05 17:53:04--  https://www.dropbox.com/s/5wuvu791vwntg9o/tcga_lung_splits.csv\n",
            "Resolving www.dropbox.com (www.dropbox.com)... 162.125.3.18, 2620:100:6018:18::a27d:312\n",
            "Connecting to www.dropbox.com (www.dropbox.com)|162.125.3.18|:443... connected.\n",
            "HTTP request sent, awaiting response... 302 Found\n",
            "Location: /s/raw/5wuvu791vwntg9o/tcga_lung_splits.csv [following]\n",
            "--2023-07-05 17:53:04--  https://www.dropbox.com/s/raw/5wuvu791vwntg9o/tcga_lung_splits.csv\n",
            "Reusing existing connection to www.dropbox.com:443.\n",
            "HTTP request sent, awaiting response... 302 Found\n",
            "Location: https://uc2584411338e312d2dc01a897a7.dl.dropboxusercontent.com/cd/0/inline/B_RI62q1I4vP4bVo7yBoP0o0nLUydn_k3nt5sJPm3p2prkcqhVBmRl_VY0lgTyiDQftl3nJuiQ634aBMTMeXHpwgK1S5-lNLTyiHfgyBRiCiYZJxAOMWbb0Ey7xSFm28sAbvx0tXuUgMV_5xNW3WWGxB_aVb5KTeYe7kQu4VIp8MCg/file# [following]\n",
            "--2023-07-05 17:53:04--  https://uc2584411338e312d2dc01a897a7.dl.dropboxusercontent.com/cd/0/inline/B_RI62q1I4vP4bVo7yBoP0o0nLUydn_k3nt5sJPm3p2prkcqhVBmRl_VY0lgTyiDQftl3nJuiQ634aBMTMeXHpwgK1S5-lNLTyiHfgyBRiCiYZJxAOMWbb0Ey7xSFm28sAbvx0tXuUgMV_5xNW3WWGxB_aVb5KTeYe7kQu4VIp8MCg/file\n",
            "Resolving uc2584411338e312d2dc01a897a7.dl.dropboxusercontent.com (uc2584411338e312d2dc01a897a7.dl.dropboxusercontent.com)... 162.125.3.15, 2620:100:601d:15::a27d:50f\n",
            "Connecting to uc2584411338e312d2dc01a897a7.dl.dropboxusercontent.com (uc2584411338e312d2dc01a897a7.dl.dropboxusercontent.com)|162.125.3.15|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 170283 (166K) [text/plain]\n",
            "Saving to: ‘tcga_lung_splits.csv’\n",
            "\n",
            "tcga_lung_splits.cs 100%[===================>] 166.29K  --.-KB/s    in 0.04s   \n",
            "\n",
            "2023-07-05 17:53:05 (4.58 MB/s) - ‘tcga_lung_splits.csv’ saved [170283/170283]\n",
            "\n",
            "--2023-07-05 17:53:05--  https://www.dropbox.com/s/euepd2owxvuwr7v/feats_pt.zip\n",
            "Resolving www.dropbox.com (www.dropbox.com)... 162.125.3.18, 2620:100:6018:18::a27d:312\n",
            "Connecting to www.dropbox.com (www.dropbox.com)|162.125.3.18|:443... connected.\n",
            "HTTP request sent, awaiting response... 302 Found\n",
            "Location: /s/raw/euepd2owxvuwr7v/feats_pt.zip [following]\n",
            "--2023-07-05 17:53:05--  https://www.dropbox.com/s/raw/euepd2owxvuwr7v/feats_pt.zip\n",
            "Reusing existing connection to www.dropbox.com:443.\n",
            "HTTP request sent, awaiting response... 302 Found\n",
            "Location: https://uc7c96ff4e834a7fbe7fe428b46e.dl.dropboxusercontent.com/cd/0/inline/B_RH2SZ-Ta92Wb-qrBDceeYXa-JB1eVjyWKjVeBNspWsJzpzE331m1Fg50O1tjhKdRvPNi-rY0WQRZRLPJGZ-972BLL1WAgvgaayRSOEmhtKCJckE7ep5Hilf393DzseUJZ8QZBF82mEjRu9Yc6PopDf6exI0kLJffoA63AVsw_7ZA/file# [following]\n",
            "--2023-07-05 17:53:06--  https://uc7c96ff4e834a7fbe7fe428b46e.dl.dropboxusercontent.com/cd/0/inline/B_RH2SZ-Ta92Wb-qrBDceeYXa-JB1eVjyWKjVeBNspWsJzpzE331m1Fg50O1tjhKdRvPNi-rY0WQRZRLPJGZ-972BLL1WAgvgaayRSOEmhtKCJckE7ep5Hilf393DzseUJZ8QZBF82mEjRu9Yc6PopDf6exI0kLJffoA63AVsw_7ZA/file\n",
            "Resolving uc7c96ff4e834a7fbe7fe428b46e.dl.dropboxusercontent.com (uc7c96ff4e834a7fbe7fe428b46e.dl.dropboxusercontent.com)... 162.125.3.15, 2620:100:6030:15::a27d:500f\n",
            "Connecting to uc7c96ff4e834a7fbe7fe428b46e.dl.dropboxusercontent.com (uc7c96ff4e834a7fbe7fe428b46e.dl.dropboxusercontent.com)|162.125.3.15|:443... connected.\n",
            "HTTP request sent, awaiting response... 302 Found\n",
            "Location: /cd/0/inline2/B_Tx61hoZe6I-K4BRWWflNltZKjWw5s2PfGB6DEjTpTHA25qZfXhWiDYX7RyiTvR8wnPI2WVgyOWHZ-WNNPpv0C8KoZ2K93ZrV8NQX_ASv09Apnn9BERSo9BnJfHwuSPZ2O5lMi396MkPxKGYxlYBa5LbMs9ob62nuyPzNFT15LnMvN4tFN0sDuAhexMObdwpBbSmaegpcre3nRvVYbkDK8IPPR2tovWHsMXz0mYBBHwYGfcNUVFZ18s-pUBvxNWeIR__AJ5PYZSPSih0R3g_nAZVb8JLUzwdAqRvxWt4vJO3gwMPGf9EOIwrnP8n9UbZa6pLqTN91bFmsPzMNAay_bUDTaJ14rENkTbjB7kDDRoYi4Y8bsohneP_C_d3AyIos7O38zSrm3oZWdx8g58tYEWI2VIp-82kozJhWiF32jXdQ/file [following]\n",
            "--2023-07-05 17:53:06--  https://uc7c96ff4e834a7fbe7fe428b46e.dl.dropboxusercontent.com/cd/0/inline2/B_Tx61hoZe6I-K4BRWWflNltZKjWw5s2PfGB6DEjTpTHA25qZfXhWiDYX7RyiTvR8wnPI2WVgyOWHZ-WNNPpv0C8KoZ2K93ZrV8NQX_ASv09Apnn9BERSo9BnJfHwuSPZ2O5lMi396MkPxKGYxlYBa5LbMs9ob62nuyPzNFT15LnMvN4tFN0sDuAhexMObdwpBbSmaegpcre3nRvVYbkDK8IPPR2tovWHsMXz0mYBBHwYGfcNUVFZ18s-pUBvxNWeIR__AJ5PYZSPSih0R3g_nAZVb8JLUzwdAqRvxWt4vJO3gwMPGf9EOIwrnP8n9UbZa6pLqTN91bFmsPzMNAay_bUDTaJ14rENkTbjB7kDDRoYi4Y8bsohneP_C_d3AyIos7O38zSrm3oZWdx8g58tYEWI2VIp-82kozJhWiF32jXdQ/file\n",
            "Reusing existing connection to uc7c96ff4e834a7fbe7fe428b46e.dl.dropboxusercontent.com:443.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 3961539230 (3.7G) [application/zip]\n",
            "Saving to: ‘feats_pt.zip’\n",
            "\n",
            "feats_pt.zip        100%[===================>]   3.69G  89.3MB/s    in 42s     \n",
            "\n",
            "2023-07-05 17:53:48 (90.9 MB/s) - ‘feats_pt.zip’ saved [3961539230/3961539230]\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# either download in colab (data will be deleted when re-starting) or mount your labdrive (preferred, but requires 4GB of storage)\n",
        "if use_drive:\n",
        "  !wget https://www.dropbox.com/s/5wuvu791vwntg9o/tcga_lung_splits.csv -P \"/content/drive/My Drive/ai4healthsummerschool\"\n",
        "  !wget https://www.dropbox.com/s/euepd2owxvuwr7v/feats_pt.zip\n",
        "  !unzip -q feats_pt.zip\n",
        "  !mv feats_pt \"/content/drive/My Drive/ai4healthsummerschool\"\n",
        "else:\n",
        "  !wget https://www.dropbox.com/s/5wuvu791vwntg9o/tcga_lung_splits.csv\n",
        "  !wget https://www.dropbox.com/s/euepd2owxvuwr7v/feats_pt.zip\n",
        "  !unzip -q feats_pt.zip\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "SEsPRU9wLnRl"
      },
      "id": "SEsPRU9wLnRl",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "8c5a99a0",
      "metadata": {
        "id": "8c5a99a0",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "14a5a6f3-9f29-43d1-800c-e77fac597f92"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2.0.1+cu118\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "import copy\n",
        "import matplotlib.pyplot as plt\n",
        "import seaborn\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import sklearn.metrics\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision\n",
        "print(torch.__version__)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e53517df",
      "metadata": {
        "id": "e53517df"
      },
      "source": [
        "### WSI data preprocessing for histology slides in the TCGA-Lung cohort\n",
        "\n",
        "\n",
        "![](https://user-images.githubusercontent.com/10300839/232984010-7f8e3a6f-e0c5-4847-8d0f-747460055528.png)\n",
        "\n",
        "To process WSIs, tools such as [CLAM](https://github.com/mahmoodlab/CLAM) are typically used for tissue patching and non-overlapping patch feature extraction. Though easy-to-use, using CLAM for feature processing would require downloading gigapixel WSIs (> 1000 WSIS in TCGA-LUAD and TCGA-LUSC), which exceeds over >100GB of storage space. To alleviate this issue, this problem set provides pre-extracted features (processed via CLAM, but using a much smaller vision encoder with $D = 320$). However, to still illustrate how CLAM preprocessing works, the cell below describes a high-level overview on how WSIs are formulated as `[M x D]`-dim bag of patch embeddings, where `M` is the number of tissue patches and `D` is the hidden dimension size of your encoder. Again, please use [CLAM](https://github.com/mahmoodlab/CLAM) if you are interested in re-generating these features.\n",
        "\n",
        "**Note:** This cell doesn't need to be run to train the final models!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "a94a7d83",
      "metadata": {
        "id": "a94a7d83",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c5186d2e-4342-4f45-8b6b-3fbd3d601e1f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "WSI Shape: torch.Size([2, 3, 256, 256])\n",
            "Probability Scores for ImageNet: torch.Size([1, 1000])\n",
            "Feature Embedding Shape: torch.Size([1, 576])\n",
            "Bag Shape (2, 576)\n"
          ]
        }
      ],
      "source": [
        "# Let's say we have a \"set of M [256 x 256 x 3] image patches (M = 512), which are taken from non-overlapping patches in the WSI.\n",
        "M = 2\n",
        "X = torch.randn(M, 3, 256, 256) # Arranged in (Batch, Channel, Width, Height) format or (B, C, W, H) for short\n",
        "print(\"WSI Shape:\", X.shape)\n",
        "\n",
        "# We would for instance use a CNN model (pretrained on ImageNet) as our vision encoder for pre-extracting \"compressed\" representations from each patch\n",
        "cnn = torchvision.models.mobilenet_v3_small()\n",
        "cnn.eval()\n",
        "\n",
        "# Since this model was taken from torchvision and trained on ImageNet, the output of the model are the probability scores of the ImageNet classes (1000 classes total).\n",
        "# To extract useful features from each patch, we have to use the penultimate layer(s) of the CNN, before feeding it into a linear layer.\n",
        "print(\"Probability Scores for ImageNet:\", cnn.forward(X[:1]).shape)\n",
        "\n",
        "# To extract the penultimate features, we can define a new function that returns the features\n",
        "# before giving it to internal classifier layer within the model.\n",
        "# Again, we want to use the pretrained features on ImageNet, but don't want the classification scores for \"ImageNet\" classes!\n",
        "# # See the below documentation for how the forward pass in MobileNetV3 works.\n",
        "# https://pytorch.org/vision/main/_modules/torchvision/models/mobilenetv3.html#mobilenet_v3_small\n",
        "encoder = lambda x: torch.flatten(cnn.avgpool(cnn.features(x)), 1)\n",
        "print(\"Feature Embedding Shape:\", encoder(X[:1]).shape)\n",
        "\n",
        "# We can now use our encoder to extract features for each patch.\n",
        "# Typically, the # of non-overlapping patches in a WSI is ~15,000. Thus, we often have to extract patch features in mini-batches.\n",
        "batch_size = 32\n",
        "H = []\n",
        "for bag_idx in range(0, M, batch_size):\n",
        "    H.append(encoder(X[bag_idx:(bag_idx+batch_size)]).cpu().detach().numpy())\n",
        "print(\"Bag Shape\", np.vstack(H).shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "Ohfe2sWOUD5q",
      "metadata": {
        "id": "Ohfe2sWOUD5q"
      },
      "source": [
        "### Data Exploration\n",
        "\n",
        "**Note:** This cell doesn't need to be run to train the final models!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "R1rXRyqPRaKk",
      "metadata": {
        "id": "R1rXRyqPRaKk",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 882
        },
        "outputId": "0179d85b-291a-46df-dd0e-78b66fad2d33"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "           case_id                                           slide_id  \\\n",
              "0     TCGA-73-4676  TCGA-73-4676-01Z-00-DX1.4d781bbc-a45e-4f9d-b6b...   \n",
              "1     TCGA-MP-A4T6  TCGA-MP-A4T6-01Z-00-DX1.085C4F5A-DB1B-434A-9D6...   \n",
              "2     TCGA-78-7167  TCGA-78-7167-01Z-00-DX1.f79e1a9b-a3eb-4c91-a1f...   \n",
              "3     TCGA-L9-A444  TCGA-L9-A444-01Z-00-DX1.88CF6F01-0C1F-4572-81E...   \n",
              "4     TCGA-55-8097  TCGA-55-8097-01Z-00-DX1.2f847b65-a5dc-41be-9dd...   \n",
              "...            ...                                                ...   \n",
              "1038  TCGA-21-A5DI  TCGA-21-A5DI-01Z-00-DX1.E9123261-ADE7-468C-9E9...   \n",
              "1039  TCGA-77-7465  TCGA-77-7465-01Z-00-DX1.25e4b0b4-4948-432f-801...   \n",
              "1040  TCGA-34-8454  TCGA-34-8454-01Z-00-DX1.A2308ED3-E430-4448-853...   \n",
              "1041  TCGA-77-7138  TCGA-77-7138-01Z-00-DX1.8c912762-0829-4692-92a...   \n",
              "1042  TCGA-77-8131  TCGA-77-8131-01Z-00-DX1.dcb8e2c7-0d2f-4b38-9db...   \n",
              "\n",
              "     tumor_type OncoTreeSiteCode            main_cancer_type sex project_id  \\\n",
              "0       Primary             LUNG  Non-Small Cell Lung Cancer   M  TCGA-LUAD   \n",
              "1       Primary             LUNG  Non-Small Cell Lung Cancer   F  TCGA-LUAD   \n",
              "2       Primary             LUNG  Non-Small Cell Lung Cancer   M  TCGA-LUAD   \n",
              "3       Primary             LUNG  Non-Small Cell Lung Cancer   F  TCGA-LUAD   \n",
              "4       Primary             LUNG  Non-Small Cell Lung Cancer   F  TCGA-LUAD   \n",
              "...         ...              ...                         ...  ..        ...   \n",
              "1038    Primary             LUNG  Non-Small Cell Lung Cancer   M  TCGA-LUSC   \n",
              "1039    Primary             LUNG  Non-Small Cell Lung Cancer   M  TCGA-LUSC   \n",
              "1040    Primary             LUNG  Non-Small Cell Lung Cancer   F  TCGA-LUSC   \n",
              "1041    Primary             LUNG  Non-Small Cell Lung Cancer   M  TCGA-LUSC   \n",
              "1042    Primary             LUNG  Non-Small Cell Lung Cancer   M  TCGA-LUSC   \n",
              "\n",
              "                         Diagnosis OncoTreeCode  OncoTreeCode_Binarized  split  \n",
              "0              Lung Adenocarcinoma         LUAD                       0  train  \n",
              "1              Lung Adenocarcinoma         LUAD                       0  train  \n",
              "2              Lung Adenocarcinoma         LUAD                       0  train  \n",
              "3              Lung Adenocarcinoma         LUAD                       0  train  \n",
              "4              Lung Adenocarcinoma         LUAD                       0  train  \n",
              "...                            ...          ...                     ...    ...  \n",
              "1038  Lung Squamous Cell Carcinoma         LUSC                       1   test  \n",
              "1039  Lung Squamous Cell Carcinoma         LUSC                       1   test  \n",
              "1040  Lung Squamous Cell Carcinoma         LUSC                       1   test  \n",
              "1041  Lung Squamous Cell Carcinoma         LUSC                       1   test  \n",
              "1042  Lung Squamous Cell Carcinoma         LUSC                       1   test  \n",
              "\n",
              "[1043 rows x 11 columns]"
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-4a4782bd-0a61-4e44-b4cf-8c28458fc9bf\">\n",
              "    <div class=\"colab-df-container\">\n",
              "      <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>case_id</th>\n",
              "      <th>slide_id</th>\n",
              "      <th>tumor_type</th>\n",
              "      <th>OncoTreeSiteCode</th>\n",
              "      <th>main_cancer_type</th>\n",
              "      <th>sex</th>\n",
              "      <th>project_id</th>\n",
              "      <th>Diagnosis</th>\n",
              "      <th>OncoTreeCode</th>\n",
              "      <th>OncoTreeCode_Binarized</th>\n",
              "      <th>split</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>TCGA-73-4676</td>\n",
              "      <td>TCGA-73-4676-01Z-00-DX1.4d781bbc-a45e-4f9d-b6b...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>M</td>\n",
              "      <td>TCGA-LUAD</td>\n",
              "      <td>Lung Adenocarcinoma</td>\n",
              "      <td>LUAD</td>\n",
              "      <td>0</td>\n",
              "      <td>train</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>TCGA-MP-A4T6</td>\n",
              "      <td>TCGA-MP-A4T6-01Z-00-DX1.085C4F5A-DB1B-434A-9D6...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>F</td>\n",
              "      <td>TCGA-LUAD</td>\n",
              "      <td>Lung Adenocarcinoma</td>\n",
              "      <td>LUAD</td>\n",
              "      <td>0</td>\n",
              "      <td>train</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>TCGA-78-7167</td>\n",
              "      <td>TCGA-78-7167-01Z-00-DX1.f79e1a9b-a3eb-4c91-a1f...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>M</td>\n",
              "      <td>TCGA-LUAD</td>\n",
              "      <td>Lung Adenocarcinoma</td>\n",
              "      <td>LUAD</td>\n",
              "      <td>0</td>\n",
              "      <td>train</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>TCGA-L9-A444</td>\n",
              "      <td>TCGA-L9-A444-01Z-00-DX1.88CF6F01-0C1F-4572-81E...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>F</td>\n",
              "      <td>TCGA-LUAD</td>\n",
              "      <td>Lung Adenocarcinoma</td>\n",
              "      <td>LUAD</td>\n",
              "      <td>0</td>\n",
              "      <td>train</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>TCGA-55-8097</td>\n",
              "      <td>TCGA-55-8097-01Z-00-DX1.2f847b65-a5dc-41be-9dd...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>F</td>\n",
              "      <td>TCGA-LUAD</td>\n",
              "      <td>Lung Adenocarcinoma</td>\n",
              "      <td>LUAD</td>\n",
              "      <td>0</td>\n",
              "      <td>train</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>...</th>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1038</th>\n",
              "      <td>TCGA-21-A5DI</td>\n",
              "      <td>TCGA-21-A5DI-01Z-00-DX1.E9123261-ADE7-468C-9E9...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>M</td>\n",
              "      <td>TCGA-LUSC</td>\n",
              "      <td>Lung Squamous Cell Carcinoma</td>\n",
              "      <td>LUSC</td>\n",
              "      <td>1</td>\n",
              "      <td>test</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1039</th>\n",
              "      <td>TCGA-77-7465</td>\n",
              "      <td>TCGA-77-7465-01Z-00-DX1.25e4b0b4-4948-432f-801...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>M</td>\n",
              "      <td>TCGA-LUSC</td>\n",
              "      <td>Lung Squamous Cell Carcinoma</td>\n",
              "      <td>LUSC</td>\n",
              "      <td>1</td>\n",
              "      <td>test</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1040</th>\n",
              "      <td>TCGA-34-8454</td>\n",
              "      <td>TCGA-34-8454-01Z-00-DX1.A2308ED3-E430-4448-853...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>F</td>\n",
              "      <td>TCGA-LUSC</td>\n",
              "      <td>Lung Squamous Cell Carcinoma</td>\n",
              "      <td>LUSC</td>\n",
              "      <td>1</td>\n",
              "      <td>test</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1041</th>\n",
              "      <td>TCGA-77-7138</td>\n",
              "      <td>TCGA-77-7138-01Z-00-DX1.8c912762-0829-4692-92a...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>M</td>\n",
              "      <td>TCGA-LUSC</td>\n",
              "      <td>Lung Squamous Cell Carcinoma</td>\n",
              "      <td>LUSC</td>\n",
              "      <td>1</td>\n",
              "      <td>test</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1042</th>\n",
              "      <td>TCGA-77-8131</td>\n",
              "      <td>TCGA-77-8131-01Z-00-DX1.dcb8e2c7-0d2f-4b38-9db...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>M</td>\n",
              "      <td>TCGA-LUSC</td>\n",
              "      <td>Lung Squamous Cell Carcinoma</td>\n",
              "      <td>LUSC</td>\n",
              "      <td>1</td>\n",
              "      <td>test</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "<p>1043 rows × 11 columns</p>\n",
              "</div>\n",
              "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-4a4782bd-0a61-4e44-b4cf-8c28458fc9bf')\"\n",
              "              title=\"Convert this dataframe to an interactive table.\"\n",
              "              style=\"display:none;\">\n",
              "        \n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "       width=\"24px\">\n",
              "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
              "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
              "  </svg>\n",
              "      </button>\n",
              "      \n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      flex-wrap:wrap;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "      <script>\n",
              "        const buttonEl =\n",
              "          document.querySelector('#df-4a4782bd-0a61-4e44-b4cf-8c28458fc9bf button.colab-df-convert');\n",
              "        buttonEl.style.display =\n",
              "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "        async function convertToInteractive(key) {\n",
              "          const element = document.querySelector('#df-4a4782bd-0a61-4e44-b4cf-8c28458fc9bf');\n",
              "          const dataTable =\n",
              "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                     [key], {});\n",
              "          if (!dataTable) return;\n",
              "\n",
              "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "            + ' to learn more about interactive tables.';\n",
              "          element.innerHTML = '';\n",
              "          dataTable['output_type'] = 'display_data';\n",
              "          await google.colab.output.renderOutput(dataTable, element);\n",
              "          const docLink = document.createElement('div');\n",
              "          docLink.innerHTML = docLinkHtml;\n",
              "          element.appendChild(docLink);\n",
              "        }\n",
              "      </script>\n",
              "    </div>\n",
              "  </div>\n",
              "  "
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "split  OncoTreeCode\n",
              "train  LUAD            433\n",
              "       LUSC            415\n",
              "test   LUAD             49\n",
              "       LUSC             49\n",
              "val    LUAD             49\n",
              "       LUSC             48\n",
              "dtype: int64"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Example filenames for extracted features: ['TCGA-75-6205-01Z-00-DX1.B75BC6BA-5196-4F62-BDA2-3F1D320ABD7C.pt', 'TCGA-77-8153-01Z-00-DX1.E8E40968-E7AD-4EA2-A832-8EC04D5CB7A1.pt', 'TCGA-56-8083-01Z-00-DX1.140c8d5b-f660-4fef-b8da-6bb2c119c021.pt', 'TCGA-73-A9RS-01Z-00-DX1.EDCEFE41-61E2-48C9-B8D5-28B55372E0CA.pt', 'TCGA-56-8201-01Z-00-DX1.883903fb-d70d-4c72-be76-6788b1bc3b35.pt']\n",
            "Overlap of extracted feature filenames + slide_id column: 1043\n",
            "Mean Bag Size: 3259.9090038314175\n",
            "Std Bag Size: 2133.97437395412\n"
          ]
        }
      ],
      "source": [
        "# where we downloaded the features and label csv to\n",
        "use_drive = False\n",
        "if use_drive:\n",
        "  feats_dirpath, csv_fpath = '/content/drive/My Drive/ai4healthsummerschool/feats_pt/', '/content/drive/My Drive/ai4healthsummerschool/tcga_lung_splits.csv'\n",
        "else:\n",
        "  feats_dirpath, csv_fpath = './feats_pt/', './tcga_lung_splits.csv'\n",
        "\n",
        "# label csv matches case_id (patient), slide_id (WSI image filename), and diagnosis (LUAD vs LUSC)\n",
        "# as well as pre-defined splits (train / val / test)\n",
        "df = pd.read_csv(csv_fpath)\n",
        "display(df)\n",
        "display(df[['split', 'OncoTreeCode']].value_counts())\n",
        "\n",
        "# extracted feature filenames + slide_id column match\n",
        "feats_pt_fnames = pd.Series(os.listdir(feats_dirpath))\n",
        "print(\"Example filenames for extracted features:\", list(feats_pt_fnames[:5]))\n",
        "print(\"Overlap of extracted feature filenames + slide_id column:\",\n",
        "      len(set(df['slide_id']).intersection(set(feats_pt_fnames.str[:-3]))))\n",
        "\n",
        "# statistics about the size of each bag\n",
        "bag_sizes = []\n",
        "for e in os.scandir(feats_dirpath):\n",
        "    feats_pt = torch.load(e.path)    # [M x d]-dim tensor\n",
        "    bag_sizes.append(feats_pt.shape[0])\n",
        "print('Mean Bag Size:', np.mean(bag_sizes))\n",
        "print('Std Bag Size:', np.std(bag_sizes))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "50cddf94",
      "metadata": {
        "id": "50cddf94"
      },
      "source": [
        "### Model 1:  AverageMIL\n",
        "\n",
        "Implemented is a minimalistic training setup that performs weakly-supervised learning via `AverageMIL` on LUAD vs. LUSC subtyping using 1043 diagnostic H\\&E tissue slides from the The Cancer Genome Atlas (features already pre-extracted and downloaded from installation, clinical metadata for all case and slide IDs also downloaded).\n",
        "\n",
        "You can run the cells in the Google Colab Notebook and see how well this algorithm performs in 20 epochs.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "bd9a3278",
      "metadata": {
        "id": "bd9a3278"
      },
      "outputs": [],
      "source": [
        "class AverageMIL(nn.Module):\n",
        "    def __init__(self, input_dim=320, hidden_dim=64, dropout=0.25, n_classes=2):\n",
        "        r\"\"\"\n",
        "        AverageMIL, a naive MIL algorithm that average pools all patch features.\n",
        "\n",
        "        Args:\n",
        "            input_dim (int): input feature dimension.\n",
        "            hidden_dim (int): hidden layer dimension.\n",
        "            dropout (float): Dropout probability.\n",
        "            n_classes (int): Number of classes.\n",
        "        \"\"\"\n",
        "        super(AverageMIL, self).__init__()\n",
        "        self.inst_level_fc = nn.Sequential(*[nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Dropout(dropout)]) # Fully-Connected Layer, applied \"instance-wise\" to each embedding\n",
        "        self.bag_level_classifier = nn.Linear(hidden_dim, n_classes)                                            # Bag-Level Classifier\n",
        "\n",
        "    def forward(self, H):\n",
        "        r\"\"\"\n",
        "        Takes as input a [M x D]-dim bag of patch features (representing a WSI), and outputs: 1) logits for classification, 2) un-normalized attention scores.\n",
        "\n",
        "        Args:\n",
        "            H (torch.Tensor): [M x D]-dim bag of patch features (representing a WSI)\n",
        "\n",
        "        Returns:\n",
        "            logits (torch.Tensor): [1 x n_classes]-dim tensor of un-normalized logits for classification task.\n",
        "            None (no attention scores to return)\n",
        "        \"\"\"\n",
        "        H = self.inst_level_fc(H)                   # 1. Preprocesses each \"instance-level\" embedding to be \"hidden-dim\"-dim size\n",
        "        z = H.mean(dim=0).unsqueeze(dim=0)          # 2. Average of Patch Embeddings\n",
        "        logits = self.bag_level_classifier(z)       # 3. Bag-Level Classifier\n",
        "        return logits, None\n",
        "\n",
        "\n",
        "class MILDataset(torch.utils.data.dataset.Dataset):\n",
        "    r\"\"\"\n",
        "    torch.utils.data.dataset.Dataset object that loads pre-extracted features per WSI from a CSV.\n",
        "\n",
        "    Args:\n",
        "        feats_dirpath (str): Path to pre-extracted patch features (assumes that these features are saved as a *.pt object with it's corresponding slide_id as the filename)\n",
        "        csv_fpath (str): Path to CSV file which contains: 1) Case ID, 2) Slide ID, 3) split information (train / val / test), and 4) label columns of interest for classification.\n",
        "        which_split (str): Split that is used for subsetting the CSV (choices: ['train', 'val', 'test'])\n",
        "        n_classes (int): Number of classes (default == 2 for LUAD vs LUSC subtyping)\n",
        "    \"\"\"\n",
        "    def __init__(self, feats_dirpath='./', csv_fpath='./tcga_lung_splits.csv', which_split='train', which_labelcol='OncoTreeCode_Binarized'):\n",
        "        self.feats_dirpath, self.csv, self.which_labelcol = feats_dirpath, pd.read_csv(csv_fpath), which_labelcol\n",
        "        self.csv_split = self.csv[self.csv['split']==which_split]\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        features = torch.load(os.path.join(self.feats_dirpath, self.csv_split.iloc[index]['slide_id']+'.pt'))\n",
        "        label = self.csv_split.iloc[index][self.which_labelcol]\n",
        "        return features, label\n",
        "\n",
        "    def __len__(self):\n",
        "        return self.csv_split.shape[0]\n",
        "\n",
        "\n",
        "def traineval_epoch(epoch, model, loader, optimizer=None, loss_fn=nn.CrossEntropyLoss(), split='train', device=torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\"), verbose=1, print_every=300):\n",
        "    r\"\"\"\n",
        "    Function that performs one epoch of training / evaluation with torch.nn model over torch.utils.data.DataLoader object.\n",
        "    Typically, these functions are defined separately for training and validation, but to save line space, we have combined the two.\n",
        "\n",
        "    Args:\n",
        "        epoch (int): Current epoch of training / evaluation (used for logging).\n",
        "        model (torch.nn): MIL model for processing bag of patch features.\n",
        "        loader (torch.utils.data.DataLoader): Object for getting bag of patch features per WSI.\n",
        "        loss_fn (torch.nn): Loss function.\n",
        "        split (str): Which split, used for setting up model + calculating loss + calculating gradients.\n",
        "        device (torch): Object representing the device on which a torch.Tensor will be allocated.\n",
        "        verbose (int): Whether to print summary epoch results (verbose >=1) and iteration info (verbose >=2).\n",
        "        print_every (int): How many batch iterations\n",
        "\n",
        "    Returns:\n",
        "        log_dict (dict): Dictionary for logging loss and performance for train / val / test split.\n",
        "    \"\"\"\n",
        "    model.train() if (split == 'train') else model.eval()       # turning on whether model should be used for training or evaluation\n",
        "    total_loss, Y_probs, labels = 0.0, [], []                   # tracking loss + logits/labels for performance metrics\n",
        "    for batch_idx, (X_bag, label) in enumerate(loader):\n",
        "        # Since we assume batch size == 1, we want to prevent torch from collating our bag of patch features as [1 x M x D] torch tensors.\n",
        "        X_bag, label = X_bag[0].to(device), label.to(device)\n",
        "\n",
        "        if (split == 'train'):\n",
        "            logits, A_norm = model(X_bag)\n",
        "            loss = loss_fn(logits, label)\n",
        "            loss.backward(), optimizer.step(), optimizer.zero_grad()\n",
        "        else:\n",
        "            with torch.no_grad(): logits, A_norm = model(X_bag)\n",
        "            loss = loss_fn(logits, label)\n",
        "\n",
        "        # Track total loss, logits, and current progress\n",
        "        total_loss += loss.item()\n",
        "        Y_probs.append(torch.softmax(logits, dim=-1).cpu().detach().numpy())\n",
        "        labels.append(label.cpu().detach().numpy())\n",
        "        if ((batch_idx + 1) % print_every == 0) and (verbose >= 2):\n",
        "            print(f'Epoch {epoch}:\\t Batch {batch_idx}\\t Avg Loss: {total_loss / (batch_idx+1):.04f}\\t Label: {label.item()}\\t Bag Size: {X_bag.shape[0]}')\n",
        "\n",
        "    # Compute balanced accuracy and AUC-ROC from saved logits / labels\n",
        "    Y_probs, labels = np.vstack(Y_probs), np.concatenate(labels)\n",
        "    log_dict = {f'{split} loss': total_loss/len(loader),\n",
        "                f'{split} acc': sklearn.metrics.balanced_accuracy_score(labels, Y_probs.argmax(axis=1)),\n",
        "                f'{split} auc': sklearn.metrics.roc_auc_score(labels, Y_probs[:, 1])}\n",
        "\n",
        "    # Print out end-of-epoch information\n",
        "    if (verbose >= 1):\n",
        "        print(f'### ({split.capitalize()} Summary) ###')\n",
        "        print(f'Epoch {epoch}:\\t' + f'\\t'.join([f'{k.capitalize().rjust(10)}: {log_dict[k]:.04f}' for k,v in log_dict.items()]))\n",
        "    return log_dict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "71e79d24",
      "metadata": {
        "id": "71e79d24",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "486e45ad-ce69-451a-f2e4-60a923ff7cb4"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 0:\t Batch 199\t Avg Loss: 0.7137\t Label: 1\t Bag Size: 957\n",
            "Epoch 0:\t Batch 399\t Avg Loss: 0.7084\t Label: 1\t Bag Size: 1131\n",
            "Epoch 0:\t Batch 599\t Avg Loss: 0.7093\t Label: 1\t Bag Size: 1342\n",
            "Epoch 0:\t Batch 799\t Avg Loss: 0.7075\t Label: 1\t Bag Size: 2319\n",
            "### (Train Summary) ###\n",
            "Epoch 0:\tTrain loss: 0.7067\t Train acc: 0.5178\t Train auc: 0.5205\n",
            "### (Val Summary) ###\n",
            "Epoch 0:\t  Val loss: 0.6804\t   Val acc: 0.5627\t   Val auc: 0.6947\n",
            "\n",
            "Epoch 1:\t Batch 199\t Avg Loss: 0.6892\t Label: 1\t Bag Size: 2830\n",
            "Epoch 1:\t Batch 399\t Avg Loss: 0.6933\t Label: 0\t Bag Size: 4905\n",
            "Epoch 1:\t Batch 599\t Avg Loss: 0.6951\t Label: 0\t Bag Size: 5624\n",
            "Epoch 1:\t Batch 799\t Avg Loss: 0.6919\t Label: 0\t Bag Size: 1298\n",
            "### (Train Summary) ###\n",
            "Epoch 1:\tTrain loss: 0.6945\t Train acc: 0.5530\t Train auc: 0.5598\n",
            "### (Val Summary) ###\n",
            "Epoch 1:\t  Val loss: 0.6736\t   Val acc: 0.5525\t   Val auc: 0.7105\n",
            "\n",
            "Epoch 2:\t Batch 199\t Avg Loss: 0.6909\t Label: 0\t Bag Size: 3509\n",
            "Epoch 2:\t Batch 399\t Avg Loss: 0.6808\t Label: 0\t Bag Size: 6975\n",
            "Epoch 2:\t Batch 599\t Avg Loss: 0.6799\t Label: 0\t Bag Size: 1555\n",
            "Epoch 2:\t Batch 799\t Avg Loss: 0.6784\t Label: 0\t Bag Size: 4302\n",
            "### (Train Summary) ###\n",
            "Epoch 2:\tTrain loss: 0.6773\t Train acc: 0.5668\t Train auc: 0.6022\n",
            "### (Val Summary) ###\n",
            "Epoch 2:\t  Val loss: 0.6703\t   Val acc: 0.5215\t   Val auc: 0.7105\n",
            "\n",
            "Epoch 3:\t Batch 199\t Avg Loss: 0.6628\t Label: 0\t Bag Size: 6237\n",
            "Epoch 3:\t Batch 399\t Avg Loss: 0.6645\t Label: 0\t Bag Size: 6577\n",
            "Epoch 3:\t Batch 599\t Avg Loss: 0.6678\t Label: 1\t Bag Size: 223\n",
            "Epoch 3:\t Batch 799\t Avg Loss: 0.6640\t Label: 0\t Bag Size: 3368\n",
            "### (Train Summary) ###\n",
            "Epoch 3:\tTrain loss: 0.6653\t Train acc: 0.5916\t Train auc: 0.6377\n",
            "### (Val Summary) ###\n",
            "Epoch 3:\t  Val loss: 0.6588\t   Val acc: 0.6507\t   Val auc: 0.7164\n",
            "\n",
            "Epoch 4:\t Batch 199\t Avg Loss: 0.6642\t Label: 1\t Bag Size: 1563\n",
            "Epoch 4:\t Batch 399\t Avg Loss: 0.6751\t Label: 0\t Bag Size: 1512\n",
            "Epoch 4:\t Batch 599\t Avg Loss: 0.6637\t Label: 1\t Bag Size: 8248\n",
            "Epoch 4:\t Batch 799\t Avg Loss: 0.6586\t Label: 1\t Bag Size: 933\n",
            "### (Train Summary) ###\n",
            "Epoch 4:\tTrain loss: 0.6589\t Train acc: 0.5951\t Train auc: 0.6462\n",
            "### (Val Summary) ###\n",
            "Epoch 4:\t  Val loss: 0.7084\t   Val acc: 0.5000\t   Val auc: 0.7092\n",
            "\n",
            "Epoch 5:\t Batch 199\t Avg Loss: 0.6623\t Label: 0\t Bag Size: 377\n",
            "Epoch 5:\t Batch 399\t Avg Loss: 0.6632\t Label: 0\t Bag Size: 1856\n",
            "Epoch 5:\t Batch 599\t Avg Loss: 0.6511\t Label: 1\t Bag Size: 3239\n",
            "Epoch 5:\t Batch 799\t Avg Loss: 0.6467\t Label: 0\t Bag Size: 356\n",
            "### (Train Summary) ###\n",
            "Epoch 5:\tTrain loss: 0.6434\t Train acc: 0.6376\t Train auc: 0.6816\n",
            "### (Val Summary) ###\n",
            "Epoch 5:\t  Val loss: 0.6920\t   Val acc: 0.5000\t   Val auc: 0.7083\n",
            "\n",
            "Epoch 6:\t Batch 199\t Avg Loss: 0.6496\t Label: 1\t Bag Size: 4575\n",
            "Epoch 6:\t Batch 399\t Avg Loss: 0.6349\t Label: 1\t Bag Size: 4743\n",
            "Epoch 6:\t Batch 599\t Avg Loss: 0.6317\t Label: 0\t Bag Size: 4420\n",
            "Epoch 6:\t Batch 799\t Avg Loss: 0.6311\t Label: 1\t Bag Size: 4480\n",
            "### (Train Summary) ###\n",
            "Epoch 6:\tTrain loss: 0.6323\t Train acc: 0.6517\t Train auc: 0.7069\n",
            "### (Val Summary) ###\n",
            "Epoch 6:\t  Val loss: 0.7329\t   Val acc: 0.5000\t   Val auc: 0.7083\n",
            "\n",
            "Epoch 7:\t Batch 199\t Avg Loss: 0.6297\t Label: 0\t Bag Size: 4443\n",
            "Epoch 7:\t Batch 399\t Avg Loss: 0.6245\t Label: 1\t Bag Size: 1261\n",
            "Epoch 7:\t Batch 599\t Avg Loss: 0.6233\t Label: 1\t Bag Size: 4893\n",
            "Epoch 7:\t Batch 799\t Avg Loss: 0.6257\t Label: 1\t Bag Size: 5530\n",
            "### (Train Summary) ###\n",
            "Epoch 7:\tTrain loss: 0.6244\t Train acc: 0.6617\t Train auc: 0.7160\n",
            "### (Val Summary) ###\n",
            "Epoch 7:\t  Val loss: 0.6490\t   Val acc: 0.6471\t   Val auc: 0.7117\n",
            "\n",
            "Epoch 8:\t Batch 199\t Avg Loss: 0.6155\t Label: 1\t Bag Size: 5553\n",
            "Epoch 8:\t Batch 399\t Avg Loss: 0.6293\t Label: 0\t Bag Size: 8315\n",
            "Epoch 8:\t Batch 599\t Avg Loss: 0.6184\t Label: 0\t Bag Size: 4905\n",
            "Epoch 8:\t Batch 799\t Avg Loss: 0.6141\t Label: 1\t Bag Size: 3630\n",
            "### (Train Summary) ###\n",
            "Epoch 8:\tTrain loss: 0.6119\t Train acc: 0.6746\t Train auc: 0.7349\n",
            "### (Val Summary) ###\n",
            "Epoch 8:\t  Val loss: 0.6426\t   Val acc: 0.6477\t   Val auc: 0.7122\n",
            "\n",
            "Epoch 9:\t Batch 199\t Avg Loss: 0.6157\t Label: 0\t Bag Size: 2482\n",
            "Epoch 9:\t Batch 399\t Avg Loss: 0.6002\t Label: 0\t Bag Size: 1856\n",
            "Epoch 9:\t Batch 599\t Avg Loss: 0.6062\t Label: 0\t Bag Size: 384\n",
            "Epoch 9:\t Batch 799\t Avg Loss: 0.6094\t Label: 0\t Bag Size: 579\n",
            "### (Train Summary) ###\n",
            "Epoch 9:\tTrain loss: 0.6099\t Train acc: 0.6767\t Train auc: 0.7334\n",
            "### (Val Summary) ###\n",
            "Epoch 9:\t  Val loss: 0.6621\t   Val acc: 0.6112\t   Val auc: 0.7151\n",
            "\n",
            "Epoch 10:\t Batch 199\t Avg Loss: 0.5707\t Label: 1\t Bag Size: 3912\n",
            "Epoch 10:\t Batch 399\t Avg Loss: 0.6012\t Label: 0\t Bag Size: 6484\n",
            "Epoch 10:\t Batch 599\t Avg Loss: 0.6039\t Label: 1\t Bag Size: 3609\n",
            "Epoch 10:\t Batch 799\t Avg Loss: 0.5971\t Label: 0\t Bag Size: 345\n",
            "### (Train Summary) ###\n",
            "Epoch 10:\tTrain loss: 0.5959\t Train acc: 0.6872\t Train auc: 0.7599\n",
            "### (Val Summary) ###\n",
            "Epoch 10:\t  Val loss: 0.6437\t   Val acc: 0.6577\t   Val auc: 0.7164\n",
            "\n",
            "Epoch 11:\t Batch 199\t Avg Loss: 0.5891\t Label: 0\t Bag Size: 1329\n",
            "Epoch 11:\t Batch 399\t Avg Loss: 0.6049\t Label: 0\t Bag Size: 3423\n",
            "Epoch 11:\t Batch 599\t Avg Loss: 0.5980\t Label: 1\t Bag Size: 5204\n",
            "Epoch 11:\t Batch 799\t Avg Loss: 0.5910\t Label: 1\t Bag Size: 1342\n",
            "### (Train Summary) ###\n",
            "Epoch 11:\tTrain loss: 0.5950\t Train acc: 0.6931\t Train auc: 0.7554\n",
            "### (Val Summary) ###\n",
            "Epoch 11:\t  Val loss: 0.6243\t   Val acc: 0.6913\t   Val auc: 0.7151\n",
            "Resetting early-stopping counter: inf -> 0.6243...\n",
            "\n",
            "Epoch 12:\t Batch 199\t Avg Loss: 0.5887\t Label: 1\t Bag Size: 5400\n",
            "Epoch 12:\t Batch 399\t Avg Loss: 0.5753\t Label: 0\t Bag Size: 6074\n",
            "Epoch 12:\t Batch 599\t Avg Loss: 0.5858\t Label: 0\t Bag Size: 4085\n",
            "Epoch 12:\t Batch 799\t Avg Loss: 0.5852\t Label: 0\t Bag Size: 2415\n",
            "### (Train Summary) ###\n",
            "Epoch 12:\tTrain loss: 0.5874\t Train acc: 0.7049\t Train auc: 0.7621\n",
            "### (Val Summary) ###\n",
            "Epoch 12:\t  Val loss: 0.6305\t   Val acc: 0.6482\t   Val auc: 0.7194\n",
            "Early-stopping counter updating: 0/5 -> 1/5...\n",
            "\n",
            "Epoch 13:\t Batch 199\t Avg Loss: 0.5654\t Label: 0\t Bag Size: 4486\n",
            "Epoch 13:\t Batch 399\t Avg Loss: 0.5794\t Label: 0\t Bag Size: 464\n",
            "Epoch 13:\t Batch 599\t Avg Loss: 0.5889\t Label: 0\t Bag Size: 8841\n",
            "Epoch 13:\t Batch 799\t Avg Loss: 0.5794\t Label: 0\t Bag Size: 1334\n",
            "### (Train Summary) ###\n",
            "Epoch 13:\tTrain loss: 0.5819\t Train acc: 0.6895\t Train auc: 0.7633\n",
            "### (Val Summary) ###\n",
            "Epoch 13:\t  Val loss: 0.6198\t   Val acc: 0.6913\t   Val auc: 0.7224\n",
            "Resetting early-stopping counter: 0.6243 -> 0.6198...\n",
            "\n",
            "Epoch 14:\t Batch 199\t Avg Loss: 0.5543\t Label: 0\t Bag Size: 2093\n",
            "Epoch 14:\t Batch 399\t Avg Loss: 0.5699\t Label: 0\t Bag Size: 9103\n",
            "Epoch 14:\t Batch 599\t Avg Loss: 0.5784\t Label: 1\t Bag Size: 2906\n",
            "Epoch 14:\t Batch 799\t Avg Loss: 0.5836\t Label: 1\t Bag Size: 1805\n",
            "### (Train Summary) ###\n",
            "Epoch 14:\tTrain loss: 0.5861\t Train acc: 0.6993\t Train auc: 0.7626\n",
            "### (Val Summary) ###\n",
            "Epoch 14:\t  Val loss: 0.6432\t   Val acc: 0.6412\t   Val auc: 0.7245\n",
            "Early-stopping counter updating: 0/5 -> 1/5...\n",
            "\n",
            "Epoch 15:\t Batch 199\t Avg Loss: 0.5710\t Label: 1\t Bag Size: 711\n",
            "Epoch 15:\t Batch 399\t Avg Loss: 0.5635\t Label: 1\t Bag Size: 980\n",
            "Epoch 15:\t Batch 599\t Avg Loss: 0.5850\t Label: 0\t Bag Size: 3167\n",
            "Epoch 15:\t Batch 799\t Avg Loss: 0.5800\t Label: 1\t Bag Size: 5346\n",
            "### (Train Summary) ###\n",
            "Epoch 15:\tTrain loss: 0.5789\t Train acc: 0.7088\t Train auc: 0.7732\n",
            "### (Val Summary) ###\n",
            "Epoch 15:\t  Val loss: 0.6718\t   Val acc: 0.5946\t   Val auc: 0.7287\n",
            "Early-stopping counter updating: 1/5 -> 2/5...\n",
            "\n",
            "Epoch 16:\t Batch 199\t Avg Loss: 0.5943\t Label: 0\t Bag Size: 4156\n",
            "Epoch 16:\t Batch 399\t Avg Loss: 0.5866\t Label: 1\t Bag Size: 2068\n",
            "Epoch 16:\t Batch 599\t Avg Loss: 0.5768\t Label: 1\t Bag Size: 3720\n",
            "Epoch 16:\t Batch 799\t Avg Loss: 0.5750\t Label: 1\t Bag Size: 6535\n",
            "### (Train Summary) ###\n",
            "Epoch 16:\tTrain loss: 0.5765\t Train acc: 0.6930\t Train auc: 0.7688\n",
            "### (Val Summary) ###\n",
            "Epoch 16:\t  Val loss: 0.6682\t   Val acc: 0.6050\t   Val auc: 0.7313\n",
            "Early-stopping counter updating: 2/5 -> 3/5...\n",
            "\n",
            "Epoch 17:\t Batch 199\t Avg Loss: 0.6217\t Label: 1\t Bag Size: 1594\n",
            "Epoch 17:\t Batch 399\t Avg Loss: 0.6043\t Label: 1\t Bag Size: 5042\n",
            "Epoch 17:\t Batch 599\t Avg Loss: 0.5981\t Label: 1\t Bag Size: 4271\n",
            "Epoch 17:\t Batch 799\t Avg Loss: 0.5737\t Label: 0\t Bag Size: 3418\n",
            "### (Train Summary) ###\n",
            "Epoch 17:\tTrain loss: 0.5760\t Train acc: 0.7109\t Train auc: 0.7726\n",
            "### (Val Summary) ###\n",
            "Epoch 17:\t  Val loss: 0.6160\t   Val acc: 0.6798\t   Val auc: 0.7326\n",
            "Resetting early-stopping counter: 0.6198 -> 0.6160...\n",
            "\n",
            "Epoch 18:\t Batch 199\t Avg Loss: 0.5472\t Label: 0\t Bag Size: 2924\n",
            "Epoch 18:\t Batch 399\t Avg Loss: 0.5504\t Label: 0\t Bag Size: 5320\n",
            "Epoch 18:\t Batch 599\t Avg Loss: 0.5530\t Label: 0\t Bag Size: 2522\n",
            "Epoch 18:\t Batch 799\t Avg Loss: 0.5583\t Label: 0\t Bag Size: 446\n",
            "### (Train Summary) ###\n",
            "Epoch 18:\tTrain loss: 0.5615\t Train acc: 0.6861\t Train auc: 0.7838\n",
            "### (Val Summary) ###\n",
            "Epoch 18:\t  Val loss: 0.6216\t   Val acc: 0.6405\t   Val auc: 0.7351\n",
            "Early-stopping counter updating: 0/5 -> 1/5...\n",
            "\n",
            "Epoch 19:\t Batch 199\t Avg Loss: 0.5789\t Label: 1\t Bag Size: 1063\n",
            "Epoch 19:\t Batch 399\t Avg Loss: 0.5672\t Label: 1\t Bag Size: 3959\n",
            "Epoch 19:\t Batch 599\t Avg Loss: 0.5663\t Label: 0\t Bag Size: 1455\n",
            "Epoch 19:\t Batch 799\t Avg Loss: 0.5632\t Label: 0\t Bag Size: 4642\n",
            "### (Train Summary) ###\n",
            "Epoch 19:\tTrain loss: 0.5647\t Train acc: 0.7148\t Train auc: 0.7823\n",
            "### (Val Summary) ###\n",
            "Epoch 19:\t  Val loss: 0.6092\t   Val acc: 0.6607\t   Val auc: 0.7385\n",
            "Resetting early-stopping counter: 0.6160 -> 0.6092...\n",
            "\n",
            "### (Test Summary) ###\n",
            "Epoch 19:\t Test loss: 0.5563\t  Test acc: 0.6837\t  Test auc: 0.8401\n"
          ]
        }
      ],
      "source": [
        "# Sets the random seed (for reproducibility)\n",
        "torch.manual_seed(2023)\n",
        "\n",
        "# Get data loaders for train-val-test split evaluation\n",
        "feats_dirpath, csv_fpath = './feats_pt/', './tcga_lung_splits.csv'\n",
        "loader_kwargs = {\n",
        "    'batch_size': 1,\n",
        "    'num_workers': 2,\n",
        "    'pin_memory': False\n",
        "}\n",
        "train_dataset, val_dataset, test_dataset = [MILDataset(feats_dirpath, csv_fpath, which_split=split) for split in ['train', 'val', 'test']]\n",
        "train_loader = torch.utils.data.DataLoader(train_dataset, shuffle=True, **loader_kwargs)\n",
        "val_loader = torch.utils.data.DataLoader(val_dataset, shuffle=False, **loader_kwargs)\n",
        "test_loader = torch.utils.data.DataLoader(test_dataset, shuffle=False, **loader_kwargs)\n",
        "\n",
        "# Get model, optimizer, and loss function\n",
        "device = torch.device('cpu')\n",
        "model = AverageMIL(input_dim=320, hidden_dim=64).to(device)\n",
        "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)\n",
        "loss_fn = nn.CrossEntropyLoss()\n",
        "\n",
        "# Set-up train-validation loop and early stopping\n",
        "num_epochs, min_early_stopping, patience, counter = 20, 10, 5, 0\n",
        "lowest_val_loss, best_model = np.inf, None\n",
        "all_train_logs, all_val_logs = [], []\n",
        "for epoch in range(num_epochs):\n",
        "    train_log = traineval_epoch(epoch, model, train_loader, optimizer=optimizer, split='train', device=device, verbose=2, print_every=200)\n",
        "    val_log = traineval_epoch(epoch, model, val_loader, optimizer=None, split='val', device=device, verbose=1)\n",
        "    val_loss = val_log['val loss']\n",
        "    # Early stopping: If validation loss does not go down for <patience> epochs after <min_early_stopping> epochs, stop model training early\n",
        "    if (epoch > min_early_stopping):\n",
        "        if (val_loss < lowest_val_loss):\n",
        "            print(f'Resetting early-stopping counter: {lowest_val_loss:.04f} -> {val_loss:.04f}...')\n",
        "            lowest_val_loss, counter, best_model = val_loss, 0, copy.deepcopy(model)\n",
        "        else:\n",
        "            print(f'Early-stopping counter updating: {counter}/{patience} -> {counter+1}/{patience}...')\n",
        "            counter += 1\n",
        "\n",
        "    if counter >= patience: break\n",
        "    print()\n",
        "\n",
        "# Report best model (lowest validation loss) on test split\n",
        "best_model = model if (best_model is None) else best_model\n",
        "test_log = traineval_epoch(epoch, best_model, test_loader, optimizer=None, split='test', device=device, verbose=1)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4ad43256",
      "metadata": {
        "id": "4ad43256"
      },
      "source": [
        "### Model 2. Implement Attention-Based Multiple Instance Learning (ABMIL)\n",
        "\n",
        "Following your experimentation with `AverageMIL`, you are ready to implement a more sophisticated model for LUAD vs. LUSC subtyping. Formally, let $\\mathbf{H}=\\left\\{\\mathbf{h}_1, \\ldots, \\mathbf{h}_M\\right\\} \\in \\mathbb{R}^{M \\times D}$  be a bag of $M$ patch embeddings, with each embedding having dimension size $D$. Ilse et al. 2018 proposed the following attention-based MIL pooling operation:\n",
        "\n",
        "$$\n",
        "\\mathbf{z} =\\sum_{i=1}^M a_i \\mathbf{h}_i, \\quad \\text{where} \\enspace a_i=\\frac{\\exp \\left\\{\\mathbf{w}^{\\top}\\left(\\tanh \\left(\\mathbf{V h}_{i} ^ { \\top }\\right) \\odot \\operatorname{sigm}\\left(\\mathbf{U h}_i^{\\top}\\right)\\right)\\right\\}}{\\sum_{j=1}^M \\exp \\left\\{\\mathbf{w}^{\\top}\\left(\\tanh \\left(\\mathbf{V} \\mathbf{h}_j^{\\top}\\right) \\odot \\operatorname{sigm}\\left(\\mathbf{U h}_j^{\\top}\\right)\\right)\\right\\}}\n",
        "$$\n",
        "\n",
        "where $\\mathbf{w} \\in \\mathbb{R}^{L \\times 1}$, $\\mathbf{V} \\in \\mathbb{R}^{L \\times D}$, and $\\mathbf{U} \\in \\mathbb{R}^{L \\times D}$ are learnable neural network parameters (implemented as fully-connected layers), and $\\mathbf{z} \\in \\mathbb{R}^{D}$ is the weighted average of all patch embeddings in $\\mathbf{H}$. The hyperbolic tangent $\\tanh (\\cdot)$ element-wise non-linearity and sigmoid non-linearity are utilized for proper gradient flow.\n",
        "\n",
        "Via PyTorch, the mathematical expression for computing $a_m$ is implemented as the `torch.nn` module `AttentionTanhSigmoidGating`, which we use as a layer in `ABMIL` for calculating the weighted average of patch embeddings.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "id": "09da5f05",
      "metadata": {
        "id": "09da5f05"
      },
      "outputs": [],
      "source": [
        "class AttentionTanhSigmoidGating(nn.Module):\n",
        "    def __init__(self, D=64, L=64, dropout=0.25):\n",
        "        r\"\"\"\n",
        "        Global attention pooling layer with tanh non-linearity and sigmoid gating (Ilse et al. 2018).\n",
        "\n",
        "        Args:\n",
        "            D (int): input feature dimension.\n",
        "            L (int): hidden layer dimension. Notation changed from M from Ilse et al 2018, as M is overloaded to also describe # of patch embeddings in a WSI.\n",
        "            dropout (float): Dropout probability.\n",
        "\n",
        "        Returns:\n",
        "            A_norm (torch.Tensor): [M x 1]-dim tensor of normalized attention scores (sum to 1)\n",
        "        \"\"\"\n",
        "        super(AttentionTanhSigmoidGating, self).__init__()\n",
        "        self.tanhV = nn.Sequential(*[nn.Linear(D, L), nn.Tanh(), nn.Dropout(dropout)])\n",
        "        self.sigmU = nn.Sequential(*[nn.Linear(D, L), nn.Sigmoid(), nn.Dropout(dropout)])\n",
        "        self.w = nn.Linear(L, 1)\n",
        "\n",
        "    def forward(self, H):\n",
        "        A_raw = self.w(self.tanhV(H).mul(self.sigmU(H))) # exponent term\n",
        "        A_norm = F.softmax(A_raw, dim=0)                 # apply softmax to normalize weights to 1\n",
        "        assert abs(A_norm.sum() - 1) < 1e-3              # Assert statement to check sum(A) ~= 1\n",
        "        return A_norm\n",
        "\n",
        "\n",
        "class ABMIL(nn.Module):\n",
        "    def __init__(self, input_dim=320, hidden_dim=64, dropout=0.25, n_classes=2):\n",
        "        r\"\"\"\n",
        "        Attention-Based Multiple Instance Learning (Ilse et al. 2018).\n",
        "\n",
        "        Args:\n",
        "            input_dim (int): input feature dimension.\n",
        "            hidden_dim (int): hidden layer dimension.\n",
        "            dropout (float): Dropout probability.\n",
        "            n_classes (int): Number of classes.\n",
        "        \"\"\"\n",
        "        super(ABMIL, self).__init__()\n",
        "        self.inst_level_fc = nn.Sequential(*[nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Dropout(dropout)]) # Fully-Connected Layer, applied \"instance-wise\" to each embedding\n",
        "        self.global_attn = AttentionTanhSigmoidGating(L=hidden_dim, D=hidden_dim)                              # Attention Function\n",
        "        self.bag_level_classifier = nn.Linear(hidden_dim, n_classes)                                            # Bag-Level Classifier\n",
        "\n",
        "    def forward(self, X: torch.randn(100, 320)):\n",
        "        r\"\"\"\n",
        "        Takes as input a [M x D]-dim bag of patch features (representing a WSI), and outputs: 1) logits for classification, 2) un-normalized attention scores.\n",
        "\n",
        "        Args:\n",
        "            X (torch.Tensor): [M x D]-dim bag of patch features (representing a WSI)\n",
        "\n",
        "        Returns:\n",
        "            logits (torch.Tensor): [1 x n_classes]-dim tensor of un-normalized logits for classification task.\n",
        "            A_norm (torch.Tensor): [M,]- or [M x 1]-dim tensor of attention scores.\n",
        "        \"\"\"\n",
        "        H_inst = self.inst_level_fc(X)         # 1. Process each feature embedding to be of size \"hidden-dim\"\n",
        "        A_norm = self.global_attn(H_inst)      # 2. Get normalized attention scores for each embedding (s.t. sum(A_norm) ~= 1)\n",
        "        z = torch.sum(A_norm * H_inst, dim=0)  # 3. Output of global attention pooling over the bag\n",
        "        logits = self.bag_level_classifier(z).unsqueeze(dim=0)   # 4. Get un-normalized logits for classification task\n",
        "        try:\n",
        "            assert logits.shape == (1,2)\n",
        "        except:\n",
        "            print(f\"Logit tensor shape is not formatted correctly. Should output [1 x 2] shape, but got {logits.shape} shape\")\n",
        "        return logits, A_norm\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "id": "703caf8b",
      "metadata": {
        "id": "703caf8b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "7e130c26-fb57-4cb9-e6a7-19ac828a74c1"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "        case_id                                           slide_id tumor_type  \\\n",
              "0  TCGA-73-4676  TCGA-73-4676-01Z-00-DX1.4d781bbc-a45e-4f9d-b6b...    Primary   \n",
              "1  TCGA-MP-A4T6  TCGA-MP-A4T6-01Z-00-DX1.085C4F5A-DB1B-434A-9D6...    Primary   \n",
              "2  TCGA-78-7167  TCGA-78-7167-01Z-00-DX1.f79e1a9b-a3eb-4c91-a1f...    Primary   \n",
              "3  TCGA-L9-A444  TCGA-L9-A444-01Z-00-DX1.88CF6F01-0C1F-4572-81E...    Primary   \n",
              "4  TCGA-55-8097  TCGA-55-8097-01Z-00-DX1.2f847b65-a5dc-41be-9dd...    Primary   \n",
              "5  TCGA-44-8119  TCGA-44-8119-01Z-00-DX1.1EBEBFA7-22DB-4365-9DF...    Primary   \n",
              "6  TCGA-49-AAR2  TCGA-49-AAR2-01Z-00-DX1.1F09F896-446E-4C55-8D0...    Primary   \n",
              "7  TCGA-L9-A743  TCGA-L9-A743-01Z-00-DX1.27ED2955-E8B5-4A3C-ADA...    Primary   \n",
              "8  TCGA-99-8032  TCGA-99-8032-01Z-00-DX1.7380b78f-ea25-43e0-ac9...    Primary   \n",
              "9  TCGA-55-6972  TCGA-55-6972-01Z-00-DX1.0b441ad0-c30f-4f63-849...    Primary   \n",
              "\n",
              "  OncoTreeSiteCode            main_cancer_type sex project_id  \\\n",
              "0             LUNG  Non-Small Cell Lung Cancer   M  TCGA-LUAD   \n",
              "1             LUNG  Non-Small Cell Lung Cancer   F  TCGA-LUAD   \n",
              "2             LUNG  Non-Small Cell Lung Cancer   M  TCGA-LUAD   \n",
              "3             LUNG  Non-Small Cell Lung Cancer   F  TCGA-LUAD   \n",
              "4             LUNG  Non-Small Cell Lung Cancer   F  TCGA-LUAD   \n",
              "5             LUNG  Non-Small Cell Lung Cancer   M  TCGA-LUAD   \n",
              "6             LUNG  Non-Small Cell Lung Cancer   M  TCGA-LUAD   \n",
              "7             LUNG  Non-Small Cell Lung Cancer   M  TCGA-LUAD   \n",
              "8             LUNG  Non-Small Cell Lung Cancer   M  TCGA-LUAD   \n",
              "9             LUNG  Non-Small Cell Lung Cancer   M  TCGA-LUAD   \n",
              "\n",
              "             Diagnosis OncoTreeCode  OncoTreeCode_Binarized  split  \n",
              "0  Lung Adenocarcinoma         LUAD                       0  train  \n",
              "1  Lung Adenocarcinoma         LUAD                       0  train  \n",
              "2  Lung Adenocarcinoma         LUAD                       0  train  \n",
              "3  Lung Adenocarcinoma         LUAD                       0  train  \n",
              "4  Lung Adenocarcinoma         LUAD                       0  train  \n",
              "5  Lung Adenocarcinoma         LUAD                       0  train  \n",
              "6  Lung Adenocarcinoma         LUAD                       0  train  \n",
              "7  Lung Adenocarcinoma         LUAD                       0  train  \n",
              "8  Lung Adenocarcinoma         LUAD                       0  train  \n",
              "9  Lung Adenocarcinoma         LUAD                       0  train  "
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-7a206111-ef0f-44b7-a45f-c5df042b0d8d\">\n",
              "    <div class=\"colab-df-container\">\n",
              "      <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>case_id</th>\n",
              "      <th>slide_id</th>\n",
              "      <th>tumor_type</th>\n",
              "      <th>OncoTreeSiteCode</th>\n",
              "      <th>main_cancer_type</th>\n",
              "      <th>sex</th>\n",
              "      <th>project_id</th>\n",
              "      <th>Diagnosis</th>\n",
              "      <th>OncoTreeCode</th>\n",
              "      <th>OncoTreeCode_Binarized</th>\n",
              "      <th>split</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>TCGA-73-4676</td>\n",
              "      <td>TCGA-73-4676-01Z-00-DX1.4d781bbc-a45e-4f9d-b6b...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>M</td>\n",
              "      <td>TCGA-LUAD</td>\n",
              "      <td>Lung Adenocarcinoma</td>\n",
              "      <td>LUAD</td>\n",
              "      <td>0</td>\n",
              "      <td>train</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>TCGA-MP-A4T6</td>\n",
              "      <td>TCGA-MP-A4T6-01Z-00-DX1.085C4F5A-DB1B-434A-9D6...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>F</td>\n",
              "      <td>TCGA-LUAD</td>\n",
              "      <td>Lung Adenocarcinoma</td>\n",
              "      <td>LUAD</td>\n",
              "      <td>0</td>\n",
              "      <td>train</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>TCGA-78-7167</td>\n",
              "      <td>TCGA-78-7167-01Z-00-DX1.f79e1a9b-a3eb-4c91-a1f...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>M</td>\n",
              "      <td>TCGA-LUAD</td>\n",
              "      <td>Lung Adenocarcinoma</td>\n",
              "      <td>LUAD</td>\n",
              "      <td>0</td>\n",
              "      <td>train</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>TCGA-L9-A444</td>\n",
              "      <td>TCGA-L9-A444-01Z-00-DX1.88CF6F01-0C1F-4572-81E...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>F</td>\n",
              "      <td>TCGA-LUAD</td>\n",
              "      <td>Lung Adenocarcinoma</td>\n",
              "      <td>LUAD</td>\n",
              "      <td>0</td>\n",
              "      <td>train</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>TCGA-55-8097</td>\n",
              "      <td>TCGA-55-8097-01Z-00-DX1.2f847b65-a5dc-41be-9dd...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>F</td>\n",
              "      <td>TCGA-LUAD</td>\n",
              "      <td>Lung Adenocarcinoma</td>\n",
              "      <td>LUAD</td>\n",
              "      <td>0</td>\n",
              "      <td>train</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>5</th>\n",
              "      <td>TCGA-44-8119</td>\n",
              "      <td>TCGA-44-8119-01Z-00-DX1.1EBEBFA7-22DB-4365-9DF...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>M</td>\n",
              "      <td>TCGA-LUAD</td>\n",
              "      <td>Lung Adenocarcinoma</td>\n",
              "      <td>LUAD</td>\n",
              "      <td>0</td>\n",
              "      <td>train</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>6</th>\n",
              "      <td>TCGA-49-AAR2</td>\n",
              "      <td>TCGA-49-AAR2-01Z-00-DX1.1F09F896-446E-4C55-8D0...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>M</td>\n",
              "      <td>TCGA-LUAD</td>\n",
              "      <td>Lung Adenocarcinoma</td>\n",
              "      <td>LUAD</td>\n",
              "      <td>0</td>\n",
              "      <td>train</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>7</th>\n",
              "      <td>TCGA-L9-A743</td>\n",
              "      <td>TCGA-L9-A743-01Z-00-DX1.27ED2955-E8B5-4A3C-ADA...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>M</td>\n",
              "      <td>TCGA-LUAD</td>\n",
              "      <td>Lung Adenocarcinoma</td>\n",
              "      <td>LUAD</td>\n",
              "      <td>0</td>\n",
              "      <td>train</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>8</th>\n",
              "      <td>TCGA-99-8032</td>\n",
              "      <td>TCGA-99-8032-01Z-00-DX1.7380b78f-ea25-43e0-ac9...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>M</td>\n",
              "      <td>TCGA-LUAD</td>\n",
              "      <td>Lung Adenocarcinoma</td>\n",
              "      <td>LUAD</td>\n",
              "      <td>0</td>\n",
              "      <td>train</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>9</th>\n",
              "      <td>TCGA-55-6972</td>\n",
              "      <td>TCGA-55-6972-01Z-00-DX1.0b441ad0-c30f-4f63-849...</td>\n",
              "      <td>Primary</td>\n",
              "      <td>LUNG</td>\n",
              "      <td>Non-Small Cell Lung Cancer</td>\n",
              "      <td>M</td>\n",
              "      <td>TCGA-LUAD</td>\n",
              "      <td>Lung Adenocarcinoma</td>\n",
              "      <td>LUAD</td>\n",
              "      <td>0</td>\n",
              "      <td>train</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-7a206111-ef0f-44b7-a45f-c5df042b0d8d')\"\n",
              "              title=\"Convert this dataframe to an interactive table.\"\n",
              "              style=\"display:none;\">\n",
              "        \n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "       width=\"24px\">\n",
              "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
              "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
              "  </svg>\n",
              "      </button>\n",
              "      \n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      flex-wrap:wrap;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "      <script>\n",
              "        const buttonEl =\n",
              "          document.querySelector('#df-7a206111-ef0f-44b7-a45f-c5df042b0d8d button.colab-df-convert');\n",
              "        buttonEl.style.display =\n",
              "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "        async function convertToInteractive(key) {\n",
              "          const element = document.querySelector('#df-7a206111-ef0f-44b7-a45f-c5df042b0d8d');\n",
              "          const dataTable =\n",
              "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                     [key], {});\n",
              "          if (!dataTable) return;\n",
              "\n",
              "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "            + ' to learn more about interactive tables.';\n",
              "          element.innerHTML = '';\n",
              "          dataTable['output_type'] = 'display_data';\n",
              "          await google.colab.output.renderOutput(dataTable, element);\n",
              "          const docLink = document.createElement('div');\n",
              "          docLink.innerHTML = docLinkHtml;\n",
              "          element.appendChild(docLink);\n",
              "        }\n",
              "      </script>\n",
              "    </div>\n",
              "  </div>\n",
              "  "
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 0:\t Batch 199\t Avg Loss: 0.7363\t Label: 0\t Bag Size: 507\n",
            "Epoch 0:\t Batch 399\t Avg Loss: 0.7275\t Label: 0\t Bag Size: 1163\n",
            "Epoch 0:\t Batch 599\t Avg Loss: 0.7064\t Label: 1\t Bag Size: 4417\n",
            "Epoch 0:\t Batch 799\t Avg Loss: 0.7082\t Label: 0\t Bag Size: 1598\n",
            "### (Train Summary) ###\n",
            "Epoch 0:\tTrain loss: 0.7074\t Train acc: 0.5111\t Train auc: 0.5243\n",
            "### (Val Summary) ###\n",
            "Epoch 0:\t  Val loss: 0.6894\t   Val acc: 0.4991\t   Val auc: 0.6386\n",
            "\n",
            "Epoch 1:\t Batch 199\t Avg Loss: 0.7016\t Label: 0\t Bag Size: 6939\n",
            "Epoch 1:\t Batch 399\t Avg Loss: 0.6918\t Label: 0\t Bag Size: 1229\n",
            "Epoch 1:\t Batch 599\t Avg Loss: 0.6928\t Label: 0\t Bag Size: 519\n",
            "Epoch 1:\t Batch 799\t Avg Loss: 0.6887\t Label: 1\t Bag Size: 4335\n",
            "### (Train Summary) ###\n",
            "Epoch 1:\tTrain loss: 0.6899\t Train acc: 0.5462\t Train auc: 0.5605\n",
            "### (Val Summary) ###\n",
            "Epoch 1:\t  Val loss: 0.6863\t   Val acc: 0.5000\t   Val auc: 0.6586\n",
            "\n",
            "Epoch 2:\t Batch 199\t Avg Loss: 0.6807\t Label: 1\t Bag Size: 7186\n",
            "Epoch 2:\t Batch 399\t Avg Loss: 0.6799\t Label: 1\t Bag Size: 2611\n",
            "Epoch 2:\t Batch 599\t Avg Loss: 0.6813\t Label: 0\t Bag Size: 708\n",
            "Epoch 2:\t Batch 799\t Avg Loss: 0.6780\t Label: 1\t Bag Size: 4235\n",
            "### (Train Summary) ###\n",
            "Epoch 2:\tTrain loss: 0.6780\t Train acc: 0.5782\t Train auc: 0.6022\n",
            "### (Val Summary) ###\n",
            "Epoch 2:\t  Val loss: 0.6763\t   Val acc: 0.5595\t   Val auc: 0.6739\n",
            "\n",
            "Epoch 3:\t Batch 199\t Avg Loss: 0.6632\t Label: 1\t Bag Size: 4657\n",
            "Epoch 3:\t Batch 399\t Avg Loss: 0.6718\t Label: 1\t Bag Size: 6929\n",
            "Epoch 3:\t Batch 599\t Avg Loss: 0.6729\t Label: 0\t Bag Size: 4582\n",
            "Epoch 3:\t Batch 799\t Avg Loss: 0.6697\t Label: 1\t Bag Size: 2143\n",
            "### (Train Summary) ###\n",
            "Epoch 3:\tTrain loss: 0.6716\t Train acc: 0.5818\t Train auc: 0.6167\n",
            "### (Val Summary) ###\n",
            "Epoch 3:\t  Val loss: 0.7256\t   Val acc: 0.5000\t   Val auc: 0.6705\n",
            "\n",
            "Epoch 4:\t Batch 199\t Avg Loss: 0.6751\t Label: 0\t Bag Size: 681\n",
            "Epoch 4:\t Batch 399\t Avg Loss: 0.6710\t Label: 0\t Bag Size: 4601\n",
            "Epoch 4:\t Batch 599\t Avg Loss: 0.6649\t Label: 0\t Bag Size: 3559\n",
            "Epoch 4:\t Batch 799\t Avg Loss: 0.6613\t Label: 0\t Bag Size: 287\n",
            "### (Train Summary) ###\n",
            "Epoch 4:\tTrain loss: 0.6596\t Train acc: 0.5896\t Train auc: 0.6441\n",
            "### (Val Summary) ###\n",
            "Epoch 4:\t  Val loss: 0.6605\t   Val acc: 0.6794\t   Val auc: 0.6913\n",
            "\n",
            "Epoch 5:\t Batch 199\t Avg Loss: 0.6653\t Label: 1\t Bag Size: 1096\n",
            "Epoch 5:\t Batch 399\t Avg Loss: 0.6557\t Label: 1\t Bag Size: 4743\n",
            "Epoch 5:\t Batch 599\t Avg Loss: 0.6457\t Label: 0\t Bag Size: 4152\n",
            "Epoch 5:\t Batch 799\t Avg Loss: 0.6487\t Label: 0\t Bag Size: 6785\n",
            "### (Train Summary) ###\n",
            "Epoch 5:\tTrain loss: 0.6471\t Train acc: 0.6375\t Train auc: 0.6780\n",
            "### (Val Summary) ###\n",
            "Epoch 5:\t  Val loss: 0.6603\t   Val acc: 0.5699\t   Val auc: 0.7007\n",
            "\n",
            "Epoch 6:\t Batch 199\t Avg Loss: 0.6459\t Label: 1\t Bag Size: 1111\n",
            "Epoch 6:\t Batch 399\t Avg Loss: 0.6323\t Label: 0\t Bag Size: 1265\n",
            "Epoch 6:\t Batch 599\t Avg Loss: 0.6366\t Label: 1\t Bag Size: 4928\n",
            "Epoch 6:\t Batch 799\t Avg Loss: 0.6297\t Label: 1\t Bag Size: 6929\n",
            "### (Train Summary) ###\n",
            "Epoch 6:\tTrain loss: 0.6280\t Train acc: 0.6475\t Train auc: 0.7111\n",
            "### (Val Summary) ###\n",
            "Epoch 6:\t  Val loss: 0.6431\t   Val acc: 0.6403\t   Val auc: 0.7062\n",
            "\n",
            "Epoch 7:\t Batch 199\t Avg Loss: 0.6292\t Label: 0\t Bag Size: 1555\n",
            "Epoch 7:\t Batch 399\t Avg Loss: 0.6124\t Label: 1\t Bag Size: 3609\n",
            "Epoch 7:\t Batch 599\t Avg Loss: 0.6162\t Label: 0\t Bag Size: 3900\n",
            "Epoch 7:\t Batch 799\t Avg Loss: 0.6058\t Label: 1\t Bag Size: 6125\n",
            "### (Train Summary) ###\n",
            "Epoch 7:\tTrain loss: 0.6057\t Train acc: 0.6645\t Train auc: 0.7375\n",
            "### (Val Summary) ###\n",
            "Epoch 7:\t  Val loss: 0.6388\t   Val acc: 0.7100\t   Val auc: 0.7177\n",
            "\n",
            "Epoch 8:\t Batch 199\t Avg Loss: 0.6160\t Label: 0\t Bag Size: 569\n",
            "Epoch 8:\t Batch 399\t Avg Loss: 0.6007\t Label: 1\t Bag Size: 3959\n",
            "Epoch 8:\t Batch 599\t Avg Loss: 0.5961\t Label: 0\t Bag Size: 3229\n",
            "Epoch 8:\t Batch 799\t Avg Loss: 0.5946\t Label: 0\t Bag Size: 4215\n",
            "### (Train Summary) ###\n",
            "Epoch 8:\tTrain loss: 0.5906\t Train acc: 0.6884\t Train auc: 0.7590\n",
            "### (Val Summary) ###\n",
            "Epoch 8:\t  Val loss: 0.6215\t   Val acc: 0.6607\t   Val auc: 0.7338\n",
            "\n",
            "Epoch 9:\t Batch 199\t Avg Loss: 0.5595\t Label: 0\t Bag Size: 1061\n",
            "Epoch 9:\t Batch 399\t Avg Loss: 0.5697\t Label: 0\t Bag Size: 1080\n",
            "Epoch 9:\t Batch 599\t Avg Loss: 0.5724\t Label: 1\t Bag Size: 1111\n",
            "Epoch 9:\t Batch 799\t Avg Loss: 0.5676\t Label: 1\t Bag Size: 2370\n",
            "### (Train Summary) ###\n",
            "Epoch 9:\tTrain loss: 0.5709\t Train acc: 0.7127\t Train auc: 0.7802\n",
            "### (Val Summary) ###\n",
            "Epoch 9:\t  Val loss: 0.6258\t   Val acc: 0.6996\t   Val auc: 0.7474\n",
            "\n",
            "Epoch 10:\t Batch 199\t Avg Loss: 0.5501\t Label: 0\t Bag Size: 657\n",
            "Epoch 10:\t Batch 399\t Avg Loss: 0.5604\t Label: 0\t Bag Size: 4582\n",
            "Epoch 10:\t Batch 599\t Avg Loss: 0.5507\t Label: 1\t Bag Size: 2964\n",
            "Epoch 10:\t Batch 799\t Avg Loss: 0.5445\t Label: 0\t Bag Size: 1877\n",
            "### (Train Summary) ###\n",
            "Epoch 10:\tTrain loss: 0.5479\t Train acc: 0.7466\t Train auc: 0.8090\n",
            "### (Val Summary) ###\n",
            "Epoch 10:\t  Val loss: 0.6297\t   Val acc: 0.6418\t   Val auc: 0.7721\n",
            "\n",
            "Epoch 11:\t Batch 199\t Avg Loss: 0.5339\t Label: 1\t Bag Size: 4775\n",
            "Epoch 11:\t Batch 399\t Avg Loss: 0.5165\t Label: 0\t Bag Size: 1163\n",
            "Epoch 11:\t Batch 599\t Avg Loss: 0.5147\t Label: 0\t Bag Size: 1045\n",
            "Epoch 11:\t Batch 799\t Avg Loss: 0.5211\t Label: 1\t Bag Size: 3985\n",
            "### (Train Summary) ###\n",
            "Epoch 11:\tTrain loss: 0.5204\t Train acc: 0.7333\t Train auc: 0.8233\n",
            "### (Val Summary) ###\n",
            "Epoch 11:\t  Val loss: 0.5948\t   Val acc: 0.7517\t   Val auc: 0.7874\n",
            "Resetting early-stopping counter: inf -> 0.5948...\n",
            "\n",
            "Epoch 12:\t Batch 199\t Avg Loss: 0.5097\t Label: 1\t Bag Size: 7475\n",
            "Epoch 12:\t Batch 399\t Avg Loss: 0.4951\t Label: 1\t Bag Size: 2272\n",
            "Epoch 12:\t Batch 599\t Avg Loss: 0.5076\t Label: 0\t Bag Size: 1681\n",
            "Epoch 12:\t Batch 799\t Avg Loss: 0.4966\t Label: 1\t Bag Size: 4230\n",
            "### (Train Summary) ###\n",
            "Epoch 12:\tTrain loss: 0.4967\t Train acc: 0.7786\t Train auc: 0.8461\n",
            "### (Val Summary) ###\n",
            "Epoch 12:\t  Val loss: 0.5608\t   Val acc: 0.7324\t   Val auc: 0.8014\n",
            "Resetting early-stopping counter: 0.5948 -> 0.5608...\n",
            "\n",
            "Epoch 13:\t Batch 199\t Avg Loss: 0.5048\t Label: 0\t Bag Size: 1846\n",
            "Epoch 13:\t Batch 399\t Avg Loss: 0.5039\t Label: 0\t Bag Size: 4170\n",
            "Epoch 13:\t Batch 599\t Avg Loss: 0.4992\t Label: 1\t Bag Size: 5204\n",
            "Epoch 13:\t Batch 799\t Avg Loss: 0.4776\t Label: 1\t Bag Size: 2459\n",
            "### (Train Summary) ###\n",
            "Epoch 13:\tTrain loss: 0.4785\t Train acc: 0.7714\t Train auc: 0.8557\n",
            "### (Val Summary) ###\n",
            "Epoch 13:\t  Val loss: 0.6134\t   Val acc: 0.6728\t   Val auc: 0.8134\n",
            "Early-stopping counter updating: 0/5 -> 1/5...\n",
            "\n",
            "Epoch 14:\t Batch 199\t Avg Loss: 0.4360\t Label: 1\t Bag Size: 1246\n",
            "Epoch 14:\t Batch 399\t Avg Loss: 0.4248\t Label: 1\t Bag Size: 511\n",
            "Epoch 14:\t Batch 599\t Avg Loss: 0.4275\t Label: 0\t Bag Size: 2878\n",
            "Epoch 14:\t Batch 799\t Avg Loss: 0.4469\t Label: 1\t Bag Size: 4986\n",
            "### (Train Summary) ###\n",
            "Epoch 14:\tTrain loss: 0.4499\t Train acc: 0.7972\t Train auc: 0.8744\n",
            "### (Val Summary) ###\n",
            "Epoch 14:\t  Val loss: 0.6795\t   Val acc: 0.6118\t   Val auc: 0.8355\n",
            "Early-stopping counter updating: 1/5 -> 2/5...\n",
            "\n",
            "Epoch 15:\t Batch 199\t Avg Loss: 0.4929\t Label: 0\t Bag Size: 8899\n",
            "Epoch 15:\t Batch 399\t Avg Loss: 0.4438\t Label: 1\t Bag Size: 2629\n",
            "Epoch 15:\t Batch 599\t Avg Loss: 0.4495\t Label: 1\t Bag Size: 5989\n",
            "Epoch 15:\t Batch 799\t Avg Loss: 0.4487\t Label: 1\t Bag Size: 3173\n",
            "### (Train Summary) ###\n",
            "Epoch 15:\tTrain loss: 0.4431\t Train acc: 0.8009\t Train auc: 0.8764\n",
            "### (Val Summary) ###\n",
            "Epoch 15:\t  Val loss: 0.7838\t   Val acc: 0.5508\t   Val auc: 0.8304\n",
            "Early-stopping counter updating: 2/5 -> 3/5...\n",
            "\n",
            "Epoch 16:\t Batch 199\t Avg Loss: 0.4658\t Label: 1\t Bag Size: 2192\n",
            "Epoch 16:\t Batch 399\t Avg Loss: 0.4532\t Label: 0\t Bag Size: 8059\n",
            "Epoch 16:\t Batch 599\t Avg Loss: 0.4323\t Label: 0\t Bag Size: 3019\n",
            "Epoch 16:\t Batch 799\t Avg Loss: 0.4167\t Label: 0\t Bag Size: 5014\n",
            "### (Train Summary) ###\n",
            "Epoch 16:\tTrain loss: 0.4073\t Train acc: 0.8277\t Train auc: 0.8994\n",
            "### (Val Summary) ###\n",
            "Epoch 16:\t  Val loss: 0.6598\t   Val acc: 0.6322\t   Val auc: 0.8410\n",
            "Early-stopping counter updating: 3/5 -> 4/5...\n",
            "\n",
            "Epoch 17:\t Batch 199\t Avg Loss: 0.4421\t Label: 0\t Bag Size: 2352\n",
            "Epoch 17:\t Batch 399\t Avg Loss: 0.4002\t Label: 1\t Bag Size: 502\n",
            "Epoch 17:\t Batch 599\t Avg Loss: 0.3995\t Label: 0\t Bag Size: 5177\n",
            "Epoch 17:\t Batch 799\t Avg Loss: 0.4163\t Label: 1\t Bag Size: 2739\n",
            "### (Train Summary) ###\n",
            "Epoch 17:\tTrain loss: 0.4133\t Train acc: 0.8065\t Train auc: 0.8932\n",
            "### (Val Summary) ###\n",
            "Epoch 17:\t  Val loss: 0.5700\t   Val acc: 0.6828\t   Val auc: 0.8520\n",
            "Early-stopping counter updating: 4/5 -> 5/5...\n",
            "### (Test Summary) ###\n",
            "Epoch 17:\t Test loss: 0.5071\t  Test acc: 0.7755\t  Test auc: 0.8917\n"
          ]
        }
      ],
      "source": [
        "\n",
        "# Sets the random seed (for reproducibility)\n",
        "torch.manual_seed(2023)\n",
        "\n",
        "# Get data loaders for train-val-test split evaluation\n",
        "if use_drive:\n",
        "  feats_dirpath = '/content/drive/My Drive/ai4healthsummerschool/feats_pt/'\n",
        "  csv_fpath = '/content/drive/My Drive/ai4healthsummerschool/tcga_lung_splits.csv'\n",
        "else:\n",
        "  feats_dirpath, csv_fpath = './feats_pt/', './tcga_lung_splits.csv'\n",
        "\n",
        "display(pd.read_csv(csv_fpath).head(10)) # visualize data\n",
        "loader_kwargs = {'batch_size': 1, 'num_workers': 2, 'pin_memory': False} # Batch size set to 1 due to variable bag sizes. Hard to collate.\n",
        "train_dataset, val_dataset, test_dataset = [MILDataset(feats_dirpath, csv_fpath, which_split=split) for split in ['train', 'val', 'test']]\n",
        "train_loader = torch.utils.data.DataLoader(train_dataset, shuffle=True, **loader_kwargs)\n",
        "val_loader = torch.utils.data.DataLoader(val_dataset, shuffle=False, **loader_kwargs)\n",
        "test_loader = torch.utils.data.DataLoader(test_dataset, shuffle=False, **loader_kwargs)\n",
        "\n",
        "# Get model, optimizer, and loss function\n",
        "device = torch.device('cpu')\n",
        "model = ABMIL(input_dim=320, hidden_dim=64).to(device)\n",
        "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)\n",
        "loss_fn = nn.CrossEntropyLoss()\n",
        "\n",
        "# Set-up train-validation loop and early stopping\n",
        "num_epochs, min_early_stopping, patience, counter = 20, 10, 5, 0\n",
        "lowest_val_loss, best_model = np.inf, None\n",
        "all_train_logs, all_val_logs = [], [] # TODO: do something with train_log / val_log every epoch to help visualize performance curves?\n",
        "for epoch in range(num_epochs):\n",
        "    train_log = traineval_epoch(epoch, model, train_loader, optimizer=optimizer, split='train', device=device, verbose=2, print_every=200)\n",
        "    val_log = traineval_epoch(epoch, model, val_loader, optimizer=None, split='val', device=device, verbose=1)\n",
        "    val_loss = val_log['val loss']\n",
        "\n",
        "    # Early stopping: If validation loss does not go down for <patience> epochs after <min_early_stopping> epochs, stop model training early\n",
        "    if (epoch > min_early_stopping):\n",
        "        if (val_loss < lowest_val_loss):\n",
        "            print(f'Resetting early-stopping counter: {lowest_val_loss:.04f} -> {val_loss:.04f}...')\n",
        "            lowest_val_loss, counter, best_model = val_loss, 0, copy.deepcopy(model)\n",
        "        else:\n",
        "            print(f'Early-stopping counter updating: {counter}/{patience} -> {counter+1}/{patience}...')\n",
        "            counter += 1\n",
        "\n",
        "    if counter >= patience: break\n",
        "    print()\n",
        "\n",
        "# Report best model (lowest validation loss) on test split\n",
        "best_model = model if (best_model is None) else best_model\n",
        "test_log = traineval_epoch(epoch, best_model, test_loader, optimizer=None, split='test', device=device, verbose=1)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7d450527",
      "metadata": {
        "id": "7d450527"
      },
      "source": [
        "### Discussion. Compare and Contrast AverageMIL and ABMIL\n",
        "\n",
        "Compare and contrast the **validation** and **test** performance of `AverageMIL` and `ABMIL`. In particular:\n",
        "\n",
        "2. Which model performed better on overall AUC and balanced accuracy on the **test split**? Which class (LUAD or LUSC) was more prone to mis-classification by each model?\n",
        "3. The following link at [http://clam.mahmoodlab.org](http://clam.mahmoodlab.org) visualizes high-attention heatmaps for LUAD vs LUSC subtyping via CLAM (similar to `ABMIL`) and confidence scores for each slides. If you were a clinical pathologist looking at these visualizations, what insights or concerns would you have in letting an AI algorithm assist you medical diagnoses?\n",
        "4. The experimental setup in this problem set is limited to only evaluating on data from TCGA. List three techniques used in Lu et al. 2021 (or other relevant biomedical imaging $\\times$ AI studies) that could be used in assessing 1) data efficiency, 2) generalization performance, and 3) concordance of attention-based interpretability of `ABMIL`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "id": "f56d08cc",
      "metadata": {
        "id": "f56d08cc"
      },
      "outputs": [],
      "source": [
        "# save best_model for next session\n",
        "if use_drive:\n",
        "  torch.save(best_model.state_dict(), '/content/drive/My Drive/ai4healthsummerschool/abmil.ckpt')\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "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.12"
    },
    "vscode": {
      "interpreter": {
        "hash": "ab8319496eb09321cf5597fd6abbcdcc0a1b57046a5c3edfd9bd1380f7f3d177"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}