notebooks/torch_lightning_training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"../\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from functools import partial\n",
    "\n",
    "\n",
    "from ml_core.preprocessing.dataset import create_dataloader\n",
    "from ml_core.modeling.unet import UNet\n",
    "\n",
    "\n",
    "import torch\n",
    "from torch.optim import Adam\n",
    "from torch.utils.data import DataLoader, ConcatDataset\n",
    "import albumentations as A\n",
    "from pytorch_lightning import Trainer\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define dataloaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "patch_size = 256\n",
    "class_name, dataset = \"Glomerulus\", \"MultiStain\"\n",
    "# class_name, dataset = \"Glomerulus\", \"Collage\"\n",
    "# class_name, dataset = \"Artery\", \"MultiStain\"\n",
    "# class_name, dataset = \"Artery\", \"Collage\"\n",
    "# class_name, dataset = \"Tubules\", \"MultiStain\"\n",
    "# class_name, dataset = \"Tubules\", \"Collage\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "aug_transform = A.Compose([\n",
    "    A.VerticalFlip(p=0.5),\n",
    "    A.HorizontalFlip(p=0.5),\n",
    "    A.Rotate(),\n",
    "    A.OneOf([\n",
    "        A.HueSaturationValue(),\n",
    "        A.ToGray()\n",
    "    ]),\n",
    "    A.ElasticTransform(p=1, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03)\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_root = Path(f\"../../data/{dataset}/hdf5_data/patch_{patch_size}/\")\n",
    "train_fname = data_root / f\"{class_name}_train.h5\"\n",
    "val_fname = data_root / f\"{class_name}_val.h5\"\n",
    "\n",
    "test_data_root = Path(f\"../../data/AAPI/hdf5_data/patch_{patch_size}/\")\n",
    "test_fname = []\n",
    "test_fname.append(test_data_root / f\"{class_name}_test.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataloader, train_dataset = create_dataloader(train_fname, transform=aug_transform, return_dataset=True)\n",
    "val_dataloader, val_dataset = create_dataloader(val_fname, batch_size=64, shuffle=False, return_dataset=True)\n",
    "test_dataloader, test_dataset = create_dataloader(test_fname, batch_size=64, shuffle=False, return_dataset=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training starts here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_callback = ModelCheckpoint(save_top_k=1, monitor=\"val/f1_score\", mode=\"max\")\n",
    "early_stopping = EarlyStopping(monitor=\"val/f1_score\", mode=\"max\", patience=10, min_delta=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
     ]
    }
   ],
   "source": [
    "trainer = Trainer(gpus=1, \n",
    "                  default_root_dir=Path(f\"../../output/{dataset}/{class_name}\"),\n",
    "                  max_epochs=10,\n",
    "                  callbacks=[early_stopping],\n",
    "                  checkpoint_callback=checkpoint_callback,\n",
    "                  log_every_n_steps=1,\n",
    "                  flush_logs_every_n_steps=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "unet = UNet(in_channels=3,\n",
    "            n_classes=2,\n",
    "            depth = 5,# this depth may need to be modified\n",
    "            wf = 3,\n",
    "            padding = True,\n",
    "            batch_norm = True,\n",
    "            up_mode = \"upconv\",\n",
    "            optimizer = [partial(Adam, lr=1e-3)],\n",
    "            edge_weight = 1.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "  | Name      | Type             | Params\n",
      "-----------------------------------------------\n",
      "0 | down_path | ModuleList       | 296 K \n",
      "1 | up_path   | ModuleList       | 191 K \n",
      "2 | last      | Conv2d           | 18    \n",
      "3 | criterion | CrossEntropyLoss | 0     \n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='Validation sanity check'), FloatProgress(value=1.0, bar_style='info', layout=Layout…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1bf8adf6098a4f63846b5d02e1860153",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='Training'), FloatProgress(value=1.0, bar_style='info', layout=Layout(flex='2'), max…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='Validating'), FloatProgress(value=1.0, bar_style='info', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='Validating'), FloatProgress(value=1.0, bar_style='info', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='Validating'), FloatProgress(value=1.0, bar_style='info', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='Validating'), FloatProgress(value=1.0, bar_style='info', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='Validating'), FloatProgress(value=1.0, bar_style='info', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='Validating'), FloatProgress(value=1.0, bar_style='info', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='Validating'), FloatProgress(value=1.0, bar_style='info', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='Validating'), FloatProgress(value=1.0, bar_style='info', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='Validating'), FloatProgress(value=1.0, bar_style='info', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='Validating'), FloatProgress(value=1.0, bar_style='info', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.fit(unet, train_dataloader, val_dataloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fetch_best_model_path(trainer):\n",
    "    chpt_paths = (Path(trainer.logger.log_dir) / Path(\"checkpoints/\")).glob(\"*.ckpt\")\n",
    "    best_chpt = next(chpt_paths)\n",
    "    return str(best_chpt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_model = UNet.load_from_checkpoint(fetch_best_model_path(trainer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "88db0ede86924e639653550ba7cb5cd5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value='Testing'), FloatProgress(value=1.0, bar_style='info', layout=Layout(flex='2'), max=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------------------------------------\n",
      "DATALOADER:0 TEST RESULTS\n",
      "{'test/f1_score': tensor(0.2861, device='cuda:0'),\n",
      " 'test/loss': tensor(1.4463, device='cuda:0'),\n",
      " 'test/precision': tensor(0.1745, device='cuda:0'),\n",
      " 'test/recall': tensor(0.8863, device='cuda:0'),\n",
      " 'train/f1_score': tensor(0.7494, device='cuda:0'),\n",
      " 'train/loss': tensor(0.3391, device='cuda:0'),\n",
      " 'train/precision': tensor(0.7325, device='cuda:0'),\n",
      " 'train/recall': tensor(0.7670, device='cuda:0'),\n",
      " 'val/f1_score': tensor(0.6242, device='cuda:0'),\n",
      " 'val/loss': tensor(0.5620, device='cuda:0'),\n",
      " 'val/precision': tensor(0.5116, device='cuda:0'),\n",
      " 'val/recall': tensor(0.8592, device='cuda:0')}\n",
      "--------------------------------------------------------------------------------\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'train/loss': 0.339130699634552,\n",
       "  'train/f1_score': 0.7493833303451538,\n",
       "  'train/precision': 0.7325277924537659,\n",
       "  'train/recall': 0.7670329809188843,\n",
       "  'val/loss': 0.5620375275611877,\n",
       "  'val/f1_score': 0.6241845488548279,\n",
       "  'val/precision': 0.5115647912025452,\n",
       "  'val/recall': 0.8592203259468079,\n",
       "  'test/loss': 1.4463034868240356,\n",
       "  'test/f1_score': 0.2861465513706207,\n",
       "  'test/precision': 0.17450568079948425,\n",
       "  'test/recall': 0.8863489627838135}]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.test(best_model, test_dataloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset: MultiStain; Class: Glomerulus.\n"
     ]
    }
   ],
   "source": [
    "print(f\"Dataset: {dataset}; Class: {class_name}.\")\n",
    "trainer.save_checkpoint(f\"../../output/{dataset}/{class_name}_best_model.ckpt\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}