notebooks/psemix_walkthrough_step2_pseudo_bag_mixup.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "59017d5b",
   "metadata": {},
   "source": [
    "# Introduction: PseMix Walkthrough\n",
    "\n",
    "PseMix (pseudo-bag mixup) contains two key steps:\n",
    "- **generating pseudo-bags**（`Step 1`; [its notebook](https://github.com/liupei101/PseMix/blob/main/notebooks/psemix_walkthrough_step1_pseudo_bag_generation.ipynb)),\n",
    "- **mixing pseudo-bags**（`Step 2`; [its notebook](https://github.com/liupei101/PseMix/blob/main/notebooks/psemix_walkthrough_step2_pseudo_bag_mixup.ipynb)).\n",
    "\n",
    "This notebook aims to help you get started with ***Step 2: Pseudo-bag Mixup***. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae385f5d",
   "metadata": {},
   "source": [
    "First of all, load required packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cdfe0e85",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/home/liup/repo/PseMix'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "import os.path as osp\n",
    "import numpy as np\n",
    "from sklearn.manifold import TSNE\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import matplotlib as mpl\n",
    "from matplotlib import pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "os.chdir(\"..\")\n",
    "os.getcwd()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85bc47ef",
   "metadata": {},
   "source": [
    "# Step 2: Pseudo-bag Mixup\n",
    "\n",
    "This step mainly shows how to obtained **mixed bags** (or **mixup-augmented bags**).\n",
    "\n",
    "Two WSI bags will be taken as the example for illustation: \n",
    "- Bag A: `./wsi_feats/feat_wsi_A_TCGA_3C_AALI.pt`.\n",
    "- Bag B: `./wsi_feats/feat_wsi_B_TCGA_BH_A2L8.pt`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f44080b",
   "metadata": {},
   "source": [
    "We first obtain the pseudo-bags of bag A and bag B, as explained in [the tutorial of the first step of PseMix](https://github.com/liupei101/PseMix/blob/main/notebooks/psemix_walkthrough_step1_pseudo_bag_generation.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "96ff2331",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[setup] seed: 42\n",
      "Prototype-based pseudo-bag dividing: n = 30, l = 8\n",
      "[info] Bag A: it has 5584 instances.\n",
      "[info] Bag A: its first pseudo-bag has 185 instances.\n",
      "[info] Bag A: its second pseudo-bag has 187 instances.\n",
      "\n",
      "\n",
      "[info] Bag B: it has 6544 instances.\n",
      "[info] Bag B: its first pseudo-bag has 219 instances.\n",
      "[info] Bag B: its second pseudo-bag has 219 instances.\n"
     ]
    }
   ],
   "source": [
    "from utils.io import read_patch_data\n",
    "from utils.core import PseudoBag\n",
    "from utils.func import seed_everything\n",
    "seed_everything(42)\n",
    "\n",
    "NUM_CLUSTER = 8 # the number of clusters\n",
    "NUM_PSEB = 30 # the number of pseudo-bags\n",
    "NUM_FT = 8 # # fine-tuning times\n",
    "\n",
    "PB = PseudoBag(NUM_PSEB, NUM_CLUSTER, proto_method='mean', iter_fine_tuning=NUM_FT)\n",
    "\n",
    "# Bag A:\n",
    "# load WSI features\n",
    "bag_feats_A = read_patch_data(\"./wsi_feats/feat_wsi_A_TCGA_3C_AALI.pt\", dtype='torch').to(torch.float)\n",
    "label_A = torch.LongTensor([1])\n",
    "# label_pseudo_bag: the indicator of pseudo-bags\n",
    "label_pseudo_bag_A = PB.divide(bag_feats_A, ret_pseudo_bag=False)\n",
    "print(f\"[info] Bag A: it has {bag_feats_A.shape[0]} instances.\")\n",
    "print(f\"[info] Bag A: its first pseudo-bag has {(label_pseudo_bag_A == 0).sum()} instances.\")\n",
    "print(f\"[info] Bag A: its second pseudo-bag has {(label_pseudo_bag_A == NUM_PSEB - 1).sum()} instances.\")\n",
    "\n",
    "# Bag B:\n",
    "# load WSI features\n",
    "bag_feats_B = read_patch_data(\"./wsi_feats/feat_wsi_B_TCGA_BH_A2L8.pt\", dtype='torch').to(torch.float)\n",
    "label_B = torch.LongTensor([0])\n",
    "# label_pseudo_bag: the indicator of pseudo-bags\n",
    "label_pseudo_bag_B = PB.divide(bag_feats_B, ret_pseudo_bag=False)\n",
    "print(\"\\n\")\n",
    "print(f\"[info] Bag B: it has {bag_feats_B.shape[0]} instances.\")\n",
    "print(f\"[info] Bag B: its first pseudo-bag has {(label_pseudo_bag_B == 0).sum()} instances.\")\n",
    "print(f\"[info] Bag B: its second pseudo-bag has {(label_pseudo_bag_B == NUM_PSEB - 1).sum()} instances.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "133c70af",
   "metadata": {},
   "source": [
    "### (1) Generating a random Mixup coefficient from Beta distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "28f3376e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[info] current Mixup coefficient is 0.5000386523859661\n"
     ]
    }
   ],
   "source": [
    "PARAM_ALPHA = 1.0 # the parameter of Beta distribution\n",
    "NUM_ITER = 1\n",
    "for i in range(NUM_ITER):\n",
    "    lam = np.random.beta(PARAM_ALPHA, PARAM_ALPHA)\n",
    "print(f\"[info] current Mixup coefficient is {lam}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe46bc84",
   "metadata": {},
   "source": [
    "Here we map the Mixup coefficient to an integer that represents the number of pseudo-bags needed to sample from bag A or bag B.\n",
    "\n",
    "Further discussions:\n",
    "- This integer is uniformly distributed, one of `(0, 1,..., NUM_PSEB)`. \n",
    "- In fact, the probability of Pseudo-bag Mixup is `(n - 2) / n`, as there is no Mixup when the interger is `0` or `NUM_PSEB`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1ad0f88b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[info] current Mixup coefficient (integer) is 15\n"
     ]
    }
   ],
   "source": [
    "lam_temp = lam if lam != 1.0 else lam - 1e-5\n",
    "lam_int  = int(lam_temp * (NUM_PSEB + 1))\n",
    "print(f\"[info] current Mixup coefficient (integer) is {lam_int}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8231906c",
   "metadata": {},
   "source": [
    "### (2) Mixing pseudo-bags according to the generated Mixup coefficient\n",
    "\n",
    "at first, we fetch the pseudo-bags from bag A and bag B:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "99940134",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fetch_pseudo_bags(X, ind_X, n:int, n_parts:int):\n",
    "    \"\"\"\n",
    "    X: bag features, usually with a shape of [N, d]\n",
    "    ind_X: pseudo-bag indicator, usually with a shape of [N, ]\n",
    "    n: pseudo-bag number, int\n",
    "    n_parts: the pseudo-bag number to fetch, int\n",
    "    \"\"\"\n",
    "    if len(X.shape) > 2:\n",
    "        X = X.squeeze(0)\n",
    "    assert n_parts <= n, 'the pseudo-bag number to fetch is invalid.'\n",
    "    if n_parts == 0:\n",
    "        return None\n",
    "\n",
    "    ind_fetched = torch.randperm(n)[:n_parts]\n",
    "    X_fetched = torch.cat([X[ind_X == ind] for ind in ind_fetched], dim=0)\n",
    "\n",
    "    return X_fetched"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "200cb4c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[info] After fetching 15 pseudo-bags from bag A, there are 2796 instances left.\n",
      "[info] After fetching 15 pseudo-bags from bag B, there are 3273 instances left.\n"
     ]
    }
   ],
   "source": [
    "bag_A = fetch_pseudo_bags(bag_feats_A, label_pseudo_bag_A, NUM_PSEB, lam_int)\n",
    "print(f\"[info] After fetching {lam_int} pseudo-bags from bag A, there are {bag_A.shape[0]} instances left.\")\n",
    "bag_B = fetch_pseudo_bags(bag_feats_B, label_pseudo_bag_B, NUM_PSEB, NUM_PSEB - lam_int)\n",
    "print(f\"[info] After fetching {NUM_PSEB - lam_int} pseudo-bags from bag B, there are {bag_B.shape[0]} instances left.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "faffa0e2",
   "metadata": {},
   "source": [
    "Next, we directly mix the two masked bags by concatenating.\n",
    "\n",
    "In fact, in our PseMix implementation we introduce a special `Random Mixup` mechanism, which is different from vanilla Mixup. A further discussions for this is as follows:\n",
    "- It could enhance the diversity of training samples.\n",
    "- It could make the model efficiently learn from vicinity samples (mixed bags), as stated in our PseMix paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "080ac08c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[info] New bag: it has 6069 instances.\n",
      "[info] New bag: its Mixup ratio is 0.5.\n"
     ]
    }
   ],
   "source": [
    "PROB_MIXUP = 0.98 # the probability to mix the pseudo-bags from two bags \n",
    "\n",
    "if np.random.rand() <= PROB_MIXUP: # our Random Mixup mechanism\n",
    "    new_bag = torch.cat([bag_A, bag_B], dim=0) # instance-axis concat\n",
    "    mixup_ratio = lam_int / NUM_PSEB\n",
    "else:\n",
    "    new_bag = bag_A\n",
    "    mixup_ratio = 1.0\n",
    "\n",
    "print(f\"[info] New bag: it has {new_bag.shape[0]} instances.\")\n",
    "print(f\"[info] New bag: its Mixup ratio is {mixup_ratio}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "893ced8b",
   "metadata": {},
   "source": [
    "### (3) training MIL models using new mixed bags (pseudo-bag-augmented bags)\n",
    "\n",
    "At this point, the new mixed bag and its label can be expressed as follows:\n",
    "```python\n",
    "Mixup_sample = new_bag # obtained above\n",
    "Mixup_label = mixup_ratio * label_A + (1 - mixup_ratio) * label_B\n",
    "```\n",
    "\n",
    "This new sample (`Mixup_sample`) and its label (`Mixup_label`) can be utilized to supervise the model training. \n",
    "\n",
    "In implementation, actually, the `Mixup_label` is not directly used for training; instead, a weighted loss is often leveraged as follows:\n",
    "\n",
    "Pseudo-code:\n",
    "```python\n",
    "# forward inference\n",
    "pred = MIL_network(new_bag)\n",
    "\n",
    "# predictive loss weighted by the `mixup_ratio`\n",
    "clf_loss = mixup_ratio * BCE_loss(pred, label_A) + (1 - mixup_ratio) * BCE_loss(pred, label_B)\n",
    "\n",
    "# backward gradients and update networks\n",
    "clf_loss.backward()\n",
    "optimizer.step()\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}