06_fast_training_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fast training with MONAI features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Disclaimer\n",
    "\n",
    "This tutorial was adapted from the official MONAI [\"Fast training tutorial\"](https://github.com/Project-MONAI/tutorials/blob/main/acceleration/fast_training_tutorial.ipynb).]\n",
    "\n",
    "This tutorial shows a regular PyTorch training program and a MONAI optimized training program, and compares the performance.  \n",
    "\n",
    "Depending on the hardware this tutorial is run on, speed-ups on the order of 10-20x or more can be obtained per epoch on a single-GPU. With additional algorithmic improvements, a desired target Dice overlap (e.g. 0.94) on the MSD Spleen-CT dataset can be obtained 100-200x times faster.\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/nvahmadi/NVIDIA_CAMP_Workshop/blob/main/06_fast_training_tutorial.ipynb)\n",
    "\n",
    "#### Learning goals\n",
    "First, run this notebook as-is, up until the section \"**Enable determinism and execute MONAI optimized training**\" to experience the training speed in a \"naive\" implementation.\n",
    "\n",
    "For MONAI fast training, your task is to introduce the following 6 acceleration features into the `train_process()` function below.\n",
    "\n",
    "Search (Ctrl+F) for the tag TODO-N to quick-jump to the correct location(s) where to solve the task.\n",
    "\n",
    "* TODO-1: `AMP` (auto mixed precision): AMP is an important feature in PyTorch since v1.6, and NVIDIA CUDA 11 added strong support for AMP and significantly improved training speed. Check the [AMP website](https://developer.nvidia.com/automatic-mixed-precision) (section \"PyTorch\") for the 7 lines of code that are required to use AMP and gradient scaling in your optimization loop (note: the mentioned `autocast()` function is part of the torch module `torch.cuda.amp.autocast()`). (Solution __[here](step1_solution.py)__).\n",
    "* TODO-2: `CacheDataset`: The [CacheDataset](https://docs.monai.io/en/stable/data.html#cachedataset) class provides a cache mechanism that can load and cache all or parts of the training data, including the results from all deterministic transforms. (Solution __[here](step2_solution.py)__).\n",
    "* TODO-3: `ToDeviced`: The [ToDeviced](https://docs.monai.io/en/stable/transforms.html#todeviced) class moves batch data to GPU. Combined with `CacheDataset`, all following random transforms are executed on GPU directly. This avoids CPU -> GPU syncs in every epoch. Please note that not all the MONAI transforms support GPU operation so far, this is a continuous work-in-progress. (Solution __[here](step3_solution.py)__).\n",
    "* TODO-4: `ThreadDataLoader`: The [ThreadDataLoader](https://docs.monai.io/en/stable/data.html#threaddataloader) class uses multi-threads instead of multi-processing, which yields faster results than `DataLoader` in light-weight task. We can use it here as we already cached the results of most computation, and remaining transforms are executed on the GPU. (Solution __[here](step4_solution.py)__).\n",
    "* TODO-5: `DiceCE` loss function: The [DiceCE](https://docs.monai.io/en/stable/losses.html#diceceloss) class computes both Dice loss and Cross Entropy Loss, and returns a weighted sum of these two losses. (Solution __[here](step5_solution.py)__).\n",
    "* TODO-6: `Novograd` optimizer: The [Novograd](https://docs.monai.io/en/stable/optimizers.html#monai.optimizers.Novograd) optimizer is based on the paper \"Stochastic Gradient Methods with Layer-wise Adaptive Moments for Training of Deep Networks\" ([arXiv](https://arxiv.org/pdf/1905.11286.pdf)). It can simply replace Adam, with an increased learning rate (e.g. 10x). (Solution __[here](step6_solution.py)__).\n",
    "\n",
    "Implement these steps one by one. Once done, run e.g. 5 epochs of training in the cell below to evaluate the acceleration. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -c \"import monai\" || pip install -q \"monai-weekly[nibabel, tqdm]\"\n",
    "!python -c \"import matplotlib\" || pip install -q matplotlib\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2020 MONAI Consortium\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "\n",
    "import glob\n",
    "import math\n",
    "import os\n",
    "import shutil\n",
    "import tempfile\n",
    "import time\n",
    "from pprint import pprint\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from torch.optim import Adam\n",
    "from monai.apps import download_and_extract\n",
    "from monai.config import print_config\n",
    "from monai.data import (\n",
    "    CacheDataset,\n",
    "    DataLoader,\n",
    "    ThreadDataLoader,\n",
    "    Dataset,\n",
    "    decollate_batch,\n",
    ")\n",
    "from monai.inferers import sliding_window_inference\n",
    "from monai.losses import DiceLoss, DiceCELoss\n",
    "from monai.metrics import DiceMetric\n",
    "from monai.networks.layers import Norm\n",
    "from monai.networks.nets import UNet\n",
    "from monai.optimizers import Novograd\n",
    "from monai.transforms import (\n",
    "    EnsureChannelFirstd,\n",
    "    AsDiscrete,\n",
    "    Compose,\n",
    "    CropForegroundd,\n",
    "    FgBgToIndicesd,\n",
    "    LoadImaged,\n",
    "    Orientationd,\n",
    "    RandCropByPosNegLabeld,\n",
    "    ScaleIntensityRanged,\n",
    "    Spacingd,\n",
    "    EnsureTyped,\n",
    "    EnsureType,\n",
    ")\n",
    "from monai.utils import get_torch_version_tuple, set_determinism\n",
    "\n",
    "print_config()\n",
    "\n",
    "if get_torch_version_tuple() < (1, 6):\n",
    "    raise RuntimeError(\n",
    "        \"AMP feature only exists in PyTorch version greater than v1.6.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup data directory\n",
    "\n",
    "You can specify a directory with the `MONAI_DATA_DIRECTORY` environment variable.  \n",
    "This allows you to save results and reuse downloads.  \n",
    "If not specified a temporary directory will be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#directory = os.environ.get(\"MONAI_DATA_DIRECTORY\")\n",
    "#root_dir = tempfile.mkdtemp() if directory is None else directory\n",
    "root_dir = './data' #os.getenv(\"WORKDIR\")\n",
    "print(f\"root dir is: {root_dir}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download dataset\n",
    "\n",
    "Downloads and extracts the Decathlon Spleen dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\"https://msd-for-monai.s3-us-west-2.amazonaws.com/Task09_Spleen.tar\"\n",
    "md5 = \"410d4a301da4e5b2f6f86ec3ddba524e\"\n",
    "compressed_file = os.path.join(root_dir, \"Task09_Spleen.tar\")\n",
    "data_root = os.path.join(root_dir, \"Task09_Spleen\")\n",
    "if not os.path.exists(data_root):\n",
    "    download_and_extract(compressed_file, compressed_file, root_dir, md5)\n",
    "    print(\"Data downloaded and extracted to 'Task09_Spleen' directory.\")\n",
    "else:\n",
    "    print(\"'Task09_Spleen' directory exists. Data download skipped.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set MSD Spleen dataset path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_images = sorted(\n",
    "    glob.glob(os.path.join(data_root, \"imagesTr\", \"*.nii.gz\"))\n",
    ")\n",
    "train_labels = sorted(\n",
    "    glob.glob(os.path.join(data_root, \"labelsTr\", \"*.nii.gz\"))\n",
    ")\n",
    "data_dicts = [\n",
    "    {\"image\": image_name, \"label\": label_name}\n",
    "    for image_name, label_name in zip(train_images, train_labels)\n",
    "]\n",
    "train_files, val_files = data_dicts[:-9], data_dicts[-9:]\n",
    "\n",
    "print('\\nExample dicts in training dataset (N=3):\\n')\n",
    "pprint(train_files[:3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup transforms for training and validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def transformations(fast=False, device='cuda:0'):\n",
    "    train_transforms = [\n",
    "        LoadImaged(keys=[\"image\", \"label\"]),\n",
    "        EnsureChannelFirstd(keys=[\"image\", \"label\"]),\n",
    "        Orientationd(keys=[\"image\", \"label\"], axcodes=\"RAS\"),\n",
    "        Spacingd(\n",
    "            keys=[\"image\", \"label\"],\n",
    "            pixdim=(1.5, 1.5, 2.0),\n",
    "            mode=(\"bilinear\", \"nearest\"),\n",
    "        ),\n",
    "        ScaleIntensityRanged(\n",
    "            keys=[\"image\"],\n",
    "            a_min=-57,\n",
    "            a_max=164,\n",
    "            b_min=0.0,\n",
    "            b_max=1.0,\n",
    "            clip=True,\n",
    "        ),\n",
    "        CropForegroundd(keys=[\"image\", \"label\"], source_key=\"image\"),\n",
    "        # pre-compute foreground and background indexes\n",
    "        # and cache them to accelerate training\n",
    "        FgBgToIndicesd(\n",
    "            keys=\"label\",\n",
    "            fg_postfix=\"_fg\",\n",
    "            bg_postfix=\"_bg\",\n",
    "            image_key=\"image\",\n",
    "        ),\n",
    "    ]\n",
    "    if fast:\n",
    "        # TODO-3-Part1/2: Augmentation on the GPU\n",
    "        # NB: move the data to GPU using EnsureTyped() transform (set: track_meta=False)\n",
    "        #     Do not forget the validation transforms below.\n",
    "        #train_transforms.append(\n",
    "        #    ... # Fill here\n",
    "        #)\n",
    "        \n",
    "        pass\n",
    "        \n",
    "    train_transforms.append(\n",
    "        # randomly crop out patch samples from big\n",
    "        # image based on pos / neg ratio\n",
    "        # the image centers of negative samples\n",
    "        # must be in valid image area\n",
    "        RandCropByPosNegLabeld(\n",
    "            keys=[\"image\", \"label\"],\n",
    "            label_key=\"label\",\n",
    "            spatial_size=(96, 96, 96),\n",
    "            pos=1,\n",
    "            neg=1,\n",
    "            num_samples=4,\n",
    "            fg_indices_key=\"label_fg\",\n",
    "            bg_indices_key=\"label_bg\",\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    val_transforms = [\n",
    "        LoadImaged(keys=[\"image\", \"label\"]),\n",
    "        EnsureChannelFirstd(keys=[\"image\", \"label\"]),\n",
    "        Orientationd(keys=[\"image\", \"label\"], axcodes=\"RAS\"),\n",
    "        Spacingd(\n",
    "            keys=[\"image\", \"label\"],\n",
    "            pixdim=(1.5, 1.5, 2.0),\n",
    "            mode=(\"bilinear\", \"nearest\"),\n",
    "        ),\n",
    "        ScaleIntensityRanged(\n",
    "            keys=[\"image\"],\n",
    "            a_min=-57,\n",
    "            a_max=164,\n",
    "            b_min=0.0,\n",
    "            b_max=1.0,\n",
    "            clip=True,\n",
    "        ),\n",
    "        CropForegroundd(keys=[\"image\", \"label\"], source_key=\"image\"),\n",
    "    ]\n",
    "    if fast:\n",
    "        \n",
    "        # TODO-3-Part2/2: Augmentation on the GPU\n",
    "        # NB: move the data to GPU using ToDeviced() transform\n",
    "        #val_transforms.append(\n",
    "        #    ... # Fill here\n",
    "        #)\n",
    "        \n",
    "        pass\n",
    "\n",
    "    return Compose(train_transforms), Compose(val_transforms)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the training progress\n",
    "The cell below implements a typical PyTorch training loop, with regular classes like `Dataset`, `DataLoader`, `Adam` optimizer and `Dice` loss to train the model.\n",
    "\n",
    "There are several locations where your job will be to implement various forms of acceleration (see task description in cells below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def train_process(fast=False, max_epochs=100):\n",
    "    learning_rate = 2e-4\n",
    "    val_interval = 1  # do validation for every epoch\n",
    "\n",
    "    train_trans, val_trans = transformations(fast=fast)\n",
    "    # set CacheDataset, ThreadDataLoader and DiceCE loss for MONAI fast training\n",
    "    if fast:\n",
    "        # TODO-2: Replace regular `DataLoader` with `ThreadDataLoader`\n",
    "        # NB: as `RandCropByPosNegLabeld` crops from the cached content and `deepcopy`\n",
    "        # the crop area instead of modifying the cached value, we can set `copy_cache=False`\n",
    "        # to avoid unnecessary deepcopy of cached content in `CacheDataset`\n",
    "        #train_ds = ... # Fill here\n",
    "        #val_ds = ... # Fill here\n",
    "        \n",
    "        # TODO-4: Instead of regular `DataLoader` use `ThreadDataLoader`\n",
    "        # NB: no multi-workers argument, because `ThreadDataLoader` works with multi-threads\n",
    "        #train_loader = ... # Fill here...\n",
    "        #val_loader = ... # ...and here\n",
    "        \n",
    "        # TODO-5: Use DiceCE loss instead of regular Dice loss.\n",
    "        #loss_function = ... # Fill here\n",
    "        \n",
    "        pass\n",
    "    \n",
    "    else:\n",
    "        train_ds = Dataset(data=train_files, transform=train_trans)\n",
    "        val_ds = Dataset(data=val_files, transform=val_trans)\n",
    "        # num_worker=4 is the best parameter according to the test\n",
    "        train_loader = DataLoader(train_ds, batch_size=4, shuffle=True, num_workers=4)\n",
    "        val_loader = DataLoader(val_ds, batch_size=1, num_workers=4)\n",
    "        loss_function = DiceLoss(to_onehot_y=True, softmax=True)\n",
    "\n",
    "    device = torch.device(\"cuda:0\")\n",
    "    model = UNet(\n",
    "        spatial_dims=3,\n",
    "        in_channels=1,\n",
    "        out_channels=2,\n",
    "        channels=(16, 32, 64, 128, 256),\n",
    "        strides=(2, 2, 2, 2),\n",
    "        num_res_units=2,\n",
    "        norm=Norm.BATCH,\n",
    "    ).to(device)\n",
    "\n",
    "    post_pred = Compose([EnsureType(), AsDiscrete(argmax=True, to_onehot=2)])\n",
    "    post_label = Compose([EnsureType(), AsDiscrete(to_onehot=2)])\n",
    "\n",
    "    dice_metric = DiceMetric(include_background=True, reduction=\"mean\", get_not_nans=False)\n",
    "\n",
    "    # set Novograd optimizer for MONAI training\n",
    "    if fast:\n",
    "        # TODO-1-Part1/2: Automatic Mixed Precision (AMP)\n",
    "        # NB: Initialize the AMP scaler\n",
    "        #scaler = ... # Fill here\n",
    "        \n",
    "        # TODO-6: Replace regular `Adam` optimizer with `Novograd`\n",
    "        # NB: Novograd paper suggests to use a bigger LR than Adam (use: learning_rate * 10),\n",
    "        #     because Adam does normalization by element-wise second moments\n",
    "        # optimizer = ... # Fill Here\n",
    "        \n",
    "        pass\n",
    "        \n",
    "    else:\n",
    "        optimizer = Adam(model.parameters(), learning_rate)\n",
    "\n",
    "    best_metric = -1\n",
    "    best_metric_epoch = -1\n",
    "    best_metrics_epochs_and_time = [[], [], []]\n",
    "    epoch_loss_values = []\n",
    "    metric_values = []\n",
    "    epoch_times = []\n",
    "    total_start = time.time()\n",
    "    for epoch in range(max_epochs):\n",
    "        epoch_start = time.time()\n",
    "        print(\"-\" * 10)\n",
    "        print(f\"epoch {epoch + 1}/{max_epochs}\")\n",
    "        model.train()\n",
    "        epoch_loss = 0\n",
    "        step = 0\n",
    "        for batch_data in train_loader:\n",
    "            step_start = time.time()\n",
    "            step += 1\n",
    "            inputs, labels = (\n",
    "                batch_data[\"image\"].to(device),\n",
    "                batch_data[\"label\"].to(device),\n",
    "            )\n",
    "            optimizer.zero_grad()\n",
    "            # set AMP for MONAI training\n",
    "            if fast:\n",
    "                \n",
    "                # TODO-1-Part2/2: Automatic Mixed Precision (AMP)\n",
    "                #... # Fill here\n",
    "                \n",
    "                pass\n",
    "                \n",
    "            else:\n",
    "                outputs = model(inputs)\n",
    "                loss = loss_function(outputs, labels)\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "            epoch_loss += loss.item()\n",
    "            epoch_len = math.ceil(len(train_ds) / train_loader.batch_size)\n",
    "            print(\n",
    "                f\"{step}/{epoch_len}, train_loss: {loss.item():.4f}\"\n",
    "                f\" step time: {(time.time() - step_start):.4f}\"\n",
    "            )\n",
    "        epoch_loss /= step\n",
    "        epoch_loss_values.append(epoch_loss)\n",
    "        print(f\"epoch {epoch + 1} average loss: {epoch_loss:.4f}\")\n",
    "\n",
    "        if (epoch + 1) % val_interval == 0:\n",
    "            model.eval()\n",
    "            with torch.no_grad():\n",
    "                for val_data in val_loader:\n",
    "                    val_inputs, val_labels = (\n",
    "                        val_data[\"image\"].to(device),\n",
    "                        val_data[\"label\"].to(device),\n",
    "                    )\n",
    "                    roi_size = (160, 160, 160)\n",
    "                    sw_batch_size = 4\n",
    "                    # set AMP for MONAI validation\n",
    "                    if fast:\n",
    "                        with torch.cuda.amp.autocast():\n",
    "                            val_outputs = sliding_window_inference(\n",
    "                                val_inputs, roi_size, sw_batch_size, model\n",
    "                            )\n",
    "                    else:\n",
    "                        val_outputs = sliding_window_inference(\n",
    "                            val_inputs, roi_size, sw_batch_size, model\n",
    "                        )\n",
    "                    val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]\n",
    "                    val_labels = [post_label(i) for i in decollate_batch(val_labels)]\n",
    "                    dice_metric(y_pred=val_outputs, y=val_labels)\n",
    "\n",
    "                metric = dice_metric.aggregate().item()\n",
    "                dice_metric.reset()\n",
    "                metric_values.append(metric)\n",
    "                if metric > best_metric:\n",
    "                    best_metric = metric\n",
    "                    best_metric_epoch = epoch + 1\n",
    "                    best_metrics_epochs_and_time[0].append(best_metric)\n",
    "                    best_metrics_epochs_and_time[1].append(best_metric_epoch)\n",
    "                    best_metrics_epochs_and_time[2].append(\n",
    "                        time.time() - total_start\n",
    "                    )\n",
    "                    torch.save(model.state_dict(), \"best_metric_model.pth\")\n",
    "                    print(\"saved new best metric model\")\n",
    "                print(\n",
    "                    f\"current epoch: {epoch + 1} current\"\n",
    "                    f\" mean dice: {metric:.4f}\"\n",
    "                    f\" best mean dice: {best_metric:.4f}\"\n",
    "                    f\" at epoch: {best_metric_epoch}\"\n",
    "                )\n",
    "        print(\n",
    "            f\"time consuming of epoch {epoch + 1} is:\"\n",
    "            f\" {(time.time() - epoch_start):.4f}\"\n",
    "        )\n",
    "        epoch_times.append(time.time() - epoch_start)\n",
    "\n",
    "    total_time = time.time() - total_start\n",
    "    print(\n",
    "        f\"train completed, best_metric: {best_metric:.4f}\"\n",
    "        f\" at epoch: {best_metric_epoch}\"\n",
    "        f\" total time: {total_time:.4f}\"\n",
    "    )\n",
    "    return (\n",
    "        max_epochs,\n",
    "        epoch_loss_values,\n",
    "        metric_values,\n",
    "        epoch_times,\n",
    "        best_metrics_epochs_and_time,\n",
    "        total_time,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Enable determinism and execute regular PyTorch training\n",
    "\n",
    "Run the cell below for a few epochs of training (e.g. `max_epochs = 3`). Observe and note down the training times per epoch.\n",
    "\n",
    "Once there is time (i.e. after the workshop), feel free to run a full convergence (i.e. max_epochs = 300) to lay a baseline for full training duration and epoch durations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "max_epochs = 3\n",
    "set_determinism(seed=0)\n",
    "regular_start = time.time()\n",
    "(\n",
    "    epoch_num,\n",
    "    epoch_loss_values,\n",
    "    metric_values,\n",
    "    epoch_times,\n",
    "    best,\n",
    "    train_time,\n",
    ") = train_process(fast=False, max_epochs=max_epochs)\n",
    "total_time = time.time() - regular_start\n",
    "print(\n",
    "    f\"total time of {epoch_num} epochs with regular PyTorch training: {total_time:.4f}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Enable determinism and execute MONAI optimized training\n",
    "\n",
    "Run optimized training by setting the training `train_process(fast=True, ...)`.\n",
    "\n",
    "Once there is time (i.e. after the workshop), feel free to run a full convergence (e.g. `max_epochs = 300`) to be able to compare the final outcomes.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "max_epochs = 3\n",
    "set_determinism(seed=0)\n",
    "monai_start = time.time()\n",
    "(\n",
    "    epoch_num,\n",
    "    m_epoch_loss_values,\n",
    "    m_metric_values,\n",
    "    m_epoch_times,\n",
    "    m_best,\n",
    "    m_train_time,\n",
    ") = train_process(fast=True, max_epochs=max_epochs)\n",
    "m_total_time = time.time() - monai_start\n",
    "print(\n",
    "    f\"total time of {epoch_num} epochs with MONAI fast training: {m_train_time:.4f},\"\n",
    "    f\" time of preparing cache: {(m_total_time - m_train_time):.4f}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot training loss and validation metrics\n",
    "\n",
    "This cell plots the training metrics and ideally should be executed after a full training run (with e.g. `max_epochs = 300`).\n",
    "\n",
    "You can see examples of all following plots at the bottom of the official [Fast Training Tutorial](https://github.com/Project-MONAI/tutorials/blob/master/acceleration/fast_training_tutorial.ipynb), obtained after 300 epochs on a V100 GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(\"train\", (12, 12))\n",
    "plt.subplot(2, 2, 1)\n",
    "plt.title(\"Regular Epoch Average Loss\")\n",
    "x = [i + 1 for i in range(len(epoch_loss_values))]\n",
    "y = epoch_loss_values\n",
    "plt.xlabel(\"epoch\")\n",
    "plt.grid(alpha=0.4, linestyle=\":\")\n",
    "plt.plot(x, y, color=\"red\")\n",
    "\n",
    "plt.subplot(2, 2, 2)\n",
    "plt.title(\"Regular Val Mean Dice\")\n",
    "x = [i + 1 for i in range(len(metric_values))]\n",
    "y = metric_values\n",
    "plt.xlabel(\"epoch\")\n",
    "plt.ylim(0, 1)\n",
    "plt.grid(alpha=0.4, linestyle=\":\")\n",
    "plt.plot(x, y, color=\"red\")\n",
    "\n",
    "plt.subplot(2, 2, 3)\n",
    "plt.title(\"Fast Epoch Average Loss\")\n",
    "x = [i + 1 for i in range(len(m_epoch_loss_values))]\n",
    "y = m_epoch_loss_values\n",
    "plt.xlabel(\"epoch\")\n",
    "plt.grid(alpha=0.4, linestyle=\":\")\n",
    "plt.plot(x, y, color=\"green\")\n",
    "\n",
    "plt.subplot(2, 2, 4)\n",
    "plt.title(\"Fast Val Mean Dice\")\n",
    "x = [i + 1 for i in range(len(m_metric_values))]\n",
    "y = m_metric_values\n",
    "plt.xlabel(\"epoch\")\n",
    "plt.ylim(0, 1)\n",
    "plt.grid(alpha=0.4, linestyle=\":\")\n",
    "plt.plot(x, y, color=\"green\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot total time and every epoch time\n",
    "\n",
    "This cell plots the training times and, as above, ideally should be executed after a full training run (with e.g. `max_epochs = 300`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(\"train\", (12, 6))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.title(\"Total Train Time(300 epochs)\")\n",
    "plt.bar(\n",
    "    \"regular PyTorch\", total_time, 1, label=\"Regular training\", color=\"red\"\n",
    ")\n",
    "plt.bar(\"Fast\", m_total_time, 1, label=\"Fast training\", color=\"green\")\n",
    "plt.ylabel(\"secs\")\n",
    "plt.grid(alpha=0.4, linestyle=\":\")\n",
    "plt.legend(loc=\"best\")\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(\"Epoch Time\")\n",
    "x = [i + 1 for i in range(max_epochs)]\n",
    "plt.xlabel(\"epoch\")\n",
    "plt.ylabel(\"secs\")\n",
    "plt.plot(x, epoch_times, label=\"Regular training\", color=\"red\")\n",
    "plt.plot(x, m_epoch_times, label=\"Fast training\", color=\"green\")\n",
    "plt.grid(alpha=0.4, linestyle=\":\")\n",
    "plt.legend(loc=\"best\")\n",
    "plt.show()\n",
    "\n",
    "print('-------------------------------')\n",
    "print(f'Total speed-up: {total_time/m_total_time:.3f} x')\n",
    "print('-------------------------------')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot total time to achieve metrics\n",
    "\n",
    "This cell plots the times and speed-ups until reaching target Dice values (i.e. 0.9/0.93/0.95/0.97). As above, this should be ideally executed after a full training run (with e.g. `max_epochs = 300`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_best_metric_time(threshold, best_values):\n",
    "    for i, v in enumerate(best_values[0]):\n",
    "        if v > threshold:\n",
    "            return best_values[2][i]\n",
    "    return -1\n",
    "\n",
    "\n",
    "def get_best_metric_epochs(threshold, best_values):\n",
    "    for i, v in enumerate(best_values[0]):\n",
    "        if v > threshold:\n",
    "            return best_values[1][i]\n",
    "    return -1\n",
    "\n",
    "\n",
    "def get_label(index):\n",
    "    if index == 0:\n",
    "        return \"Regular training\"\n",
    "    elif index == 1:\n",
    "        return \"Fast training\"\n",
    "    else:\n",
    "        return None\n",
    "\n",
    "\n",
    "plt.figure(\"train\", (18, 6))\n",
    "plt.subplot(1, 3, 1)\n",
    "plt.title(\"Metrics Time\")\n",
    "plt.xlabel(\"secs\")\n",
    "plt.ylabel(\"best mean_dice\")\n",
    "plt.plot(best[2], best[0], label=\"Regular training\", color=\"red\")\n",
    "plt.plot(m_best[2], m_best[0], label=\"Fast training\", color=\"green\")\n",
    "plt.grid(alpha=0.4, linestyle=\":\")\n",
    "plt.legend(loc=\"best\")\n",
    "\n",
    "plt.subplot(1, 3, 2)\n",
    "plt.title(\"Typical Metrics Time\")\n",
    "plt.xlabel(\"best mean_dice\")\n",
    "plt.ylabel(\"secs\")\n",
    "labels = [\"0.90\", \"0.90 \", \"0.93\", \"0.93 \", \"0.95\", \"0.95 \", \"0.97\", \"0.97 \"]\n",
    "x_values = [0.9, 0.9, 0.93, 0.93, 0.95, 0.95, 0.97, 0.97]\n",
    "for i, (l, x) in enumerate(zip(labels, x_values)):\n",
    "    value = int(get_best_metric_time(x, best if i % 2 == 0 else m_best))\n",
    "    color = \"red\" if i % 2 == 0 else \"green\"\n",
    "    plt.bar(l, value, 0.5, label=get_label(i), color=color)\n",
    "    plt.text(l, value, \"%s\" % value, ha=\"center\", va=\"bottom\")\n",
    "plt.grid(alpha=0.4, linestyle=\":\")\n",
    "plt.legend(loc=\"best\")\n",
    "\n",
    "plt.subplot(1, 3, 3)\n",
    "plt.title(\"Typical Metrics Epochs\")\n",
    "plt.xlabel(\"best mean_dice\")\n",
    "plt.ylabel(\"epochs\")\n",
    "for i, (l, x) in enumerate(zip(labels, x_values)):\n",
    "    value = int(get_best_metric_epochs(x, best if i % 2 == 0 else m_best))\n",
    "    color = \"red\" if i % 2 == 0 else \"green\"\n",
    "    plt.bar(l, value, 0.5, label=get_label(i), color=color)\n",
    "    plt.text(l, value, \"%s\" % value, ha=\"center\", va=\"bottom\")\n",
    "plt.grid(alpha=0.4, linestyle=\":\")\n",
    "plt.legend(loc=\"best\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cleanup data directory\n",
    "\n",
    "Remove directory if a temporary was used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment if you want to delete data.\n",
    "#if directory is None:\n",
    "#    shutil.rmtree(root_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In the above tutorial, you went through several steps of accelerating your segmentation model training. Running a full training until convergence (e.g. `max_epochs = 300`) probably takes too long for the scope of this workshop. \n",
    "\n",
    "In the official MONAI [\"Fast training tutorial\"](https://github.com/Project-MONAI/tutorials/blob/master/acceleration/fast_training_tutorial.ipynb), it is stated that with a V100 GPU, the developers are able to achieve training convergence at a validation mean dice of `0.95` within `1 minute`. This corresponds to an approximately `200x` speedup compared with the Pytorch regular implementation when achieving same metric. And every epoch is `20x` faster than regular training.\n",
    "\n",
    "There are more tips and tricks for accelerated training with MONAI. For more details, please visit the official [\"Fast Model Training Guide\"](https://github.com/Project-MONAI/tutorials/blob/master/acceleration/fast_model_training_guide.md) in the [MONAI tutorials](https://github.com/Project-MONAI/tutorials) GitHub repository. \n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}