notebooks/lightning_segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Segmentation with `rising` and `PytorchLightning`\n",
    "\n",
    "This example will show you how to build a proper training pipeline with [PyTorch Lightning](https://github.com/PytorchLightning/pytorch-lightning) and [rising](https://github.com/PhoenixDL/rising).\n",
    "But first let's configure this notebook correctly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and install all our dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --upgrade --quiet pytorch-lightning # for training\n",
    "!pip install --upgrade --quiet git+https://github.com/PhoenixDL/rising # for data handling\n",
    "!pip install --upgrade --quiet SimpleITK # for loading medical data\n",
    "!pip install --upgrade --quiet tensorboard # for monitoring training\n",
    "!pip install --upgrade --quiet gdown # to download data cross platform"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data\n",
    "\n",
    "Once this is done, we need to take care of our training data. To show risings full capabilities, we will be using 3D data from [medical decathlon](http://medicaldecathlon.com/) (specifically Task 4: Hippocampus). \n",
    "\n",
    "### Download\n",
    "We will use the data provided on Google Drive and download it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import SimpleITK as sitk\n",
    "import json\n",
    "import tempfile\n",
    "import numpy as np\n",
    "import tarfile\n",
    "import time\n",
    "import gdown\n",
    "\n",
    "\n",
    "temp_dir = tempfile.mkdtemp()\n",
    "\n",
    "# generate dummy data for ci/cd\n",
    "if 'CI' in os.environ:\n",
    "    data_dir = os.path.join(temp_dir, 'DummyData')\n",
    "    os.makedirs(os.path.join(data_dir, 'training'), exist_ok=True)\n",
    "    data_paths = []\n",
    "    \n",
    "    for idx in range(50):\n",
    "        img = np.random.randint(-500, 500, (32, 64, 32), np.int16)\n",
    "        mask = np.random.randint(0, 1, (32, 64, 32), np.int16)\n",
    "\n",
    "        img_path = os.path.join(data_dir, 'training', 'img_%03d.nii.gz' % idx)\n",
    "        mask_path = os.path.join(data_dir, 'training', 'mask_%03d.nii.gz' % idx)\n",
    "        sitk.WriteImage(sitk.GetImageFromArray(img), img_path)\n",
    "        sitk.WriteImage(sitk.GetImageFromArray(mask), mask_path)\n",
    "\n",
    "        data_paths.append({'image': img_path, 'label': mask_path})\n",
    "\n",
    "    with open(os.path.join(data_dir, 'dataset.json'), 'w') as f:\n",
    "        json.dump({'training': data_paths}, f, sort_keys=True, indent=4)\n",
    "\n",
    "else:\n",
    "\n",
    "    data_url = \"https://drive.google.com/uc?export=download&id=1RzPB1_bqzQhlWvU-YGvZzhx2omcDh38C\"\n",
    "\n",
    "    data_dir = os.path.join(temp_dir, 'Task04_Hippocampus')\n",
    "    download_path = os.path.join(temp_dir, 'data.tar')\n",
    "\n",
    "    gdown.download(data_url, download_path)\n",
    "\n",
    "    tarfile.TarFile(download_path).extractall(temp_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great! We got our data. Now we can work on loading it.\n",
    "For loading data, `rising` follows the same principle as [PyTorch](https://pytorch.org): It separates the dataset, which provides the logic of loading a single sample, from the dataloader for automatted handling of parallel loading and batching.\n",
    "\n",
    "In fact we at `rising` thought that there is no need to reinvent the wheel. This is why we internally use PyTorch's data structure and just extend it a bit. We'll come to these extensions later.\n",
    "\n",
    "### Dataset\n",
    "\n",
    "Our dataset is fairly simple. It just loads the Nifti Data we downloaded before and returns each sample as a dict:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import SimpleITK as sitk\n",
    "import json\n",
    "from rising import loading\n",
    "from rising.loading import Dataset\n",
    "import torch\n",
    "class NiiDataset(Dataset):\n",
    "    def __init__(self, train: bool, data_dir: str):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            train: whether to use the training or the validation split\n",
    "            data_dir: directory containing the data\n",
    "        \"\"\"\n",
    "        with open(os.path.join(data_dir, 'dataset.json')) as f:\n",
    "            content = json.load(f)['training']\n",
    "\n",
    "            num_train_samples = int(len(content) * 0.9)\n",
    "            \n",
    "            # Split train data into training and validation, \n",
    "            # since test data contains no ground truth\n",
    "            if train:\n",
    "                data = content[:num_train_samples]\n",
    "            else:\n",
    "                data = content[num_train_samples:]\n",
    "            \n",
    "            self.data = data\n",
    "            self.data_dir = data_dir\n",
    "\n",
    "    def __getitem__(self, item: int) -> dict:\n",
    "        \"\"\"\n",
    "        Loads and Returns a single sample\n",
    "        \n",
    "        Args:\n",
    "            item: index specifying which item to load\n",
    "            \n",
    "        Returns:\n",
    "            dict: the loaded sample\n",
    "        \"\"\"\n",
    "        sample = self.data[item]\n",
    "        img = sitk.GetArrayFromImage(\n",
    "            sitk.ReadImage(os.path.join(self.data_dir, sample['image'])))\n",
    "\n",
    "        # add channel dim if necesary\n",
    "        if img.ndim == 3:\n",
    "            img = img[None]\n",
    "\n",
    "        label = sitk.GetArrayFromImage(\n",
    "            sitk.ReadImage(os.path.join(self.data_dir, sample['label'])))\n",
    "        \n",
    "        # convert multiclass to binary task by combining all positives\n",
    "        label = label > 0\n",
    "        \n",
    "        # add channel dim if necessary\n",
    "        if label.ndim == 3:\n",
    "            label = label[None]\n",
    "        return {'data': torch.from_numpy(img).float(), \n",
    "                'label': torch.from_numpy(label).float()}\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        \"\"\"\n",
    "        Adds a length to the dataset\n",
    "        \n",
    "        Returns:\n",
    "            int: dataset's length\n",
    "        \"\"\"\n",
    "        return len(self.data)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For compatibility each `rising` dataset must hold the same attributes as a `PyTorch` dataset. This basically comes down to be indexeable. This means, each Sequence-like data (e.g. lists, tuples, tensors or arrays) could also directly be used as a dataset. Ideally each dataset also has a length, since the dataloader tries to use this length to calculate/estimate its own length.\n",
    "\n",
    "## Integration With PyTorch Lightning: Model and Training\n",
    "\n",
    "\n",
    "After obtaining our data and implementing a way to load it, we now need a model we can train. For this, we will use a fairly simple implementation of the [U-Net](https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/), which basically is an encoder-decoder network with skip connections. In `Lightning` all modules should be derived from a `LightningModule`, which itself is a subclass of the `torch.nn.Module`. For further details on the `LightningModule` please refer to the [project itself](https://github.com/PyTorchLightning/pytorch-lightning) or it's [documentation](https://pytorch-lightning.readthedocs.io/en/stable/).\n",
    "\n",
    "### Model\n",
    "\n",
    "For now we will only define the network's logic and omit the training logic, which we'll add later.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytorch_lightning as pl\n",
    "import torch\n",
    "\n",
    "class Unet(pl.LightningModule):\n",
    "    \"\"\"Simple U-Net without training logic\"\"\"\n",
    "    def __init__(self, hparams: dict):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            hparams: the hyperparameters needed to construct the network.\n",
    "                Specifically these are:\n",
    "                * start_filts (int)\n",
    "                * depth (int)\n",
    "                * in_channels (int)\n",
    "                * num_classes (int)\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        # 4 downsample layers\n",
    "        out_filts = hparams.get('start_filts', 16)\n",
    "        depth = hparams.get('depth', 3)\n",
    "        in_filts = hparams.get('in_channels', 1)\n",
    "        num_classes = hparams.get('num_classes', 2)\n",
    "\n",
    "        for idx in range(depth):\n",
    "            down_block = torch.nn.Sequential(\n",
    "                torch.nn.Conv3d(in_filts, out_filts, kernel_size=3, padding=1), \n",
    "                torch.nn.ReLU(inplace=True),\n",
    "                torch.nn.Conv3d(out_filts, out_filts, kernel_size=3, padding=1), \n",
    "                torch.nn.ReLU(inplace=True)\n",
    "            )\n",
    "            in_filts = out_filts\n",
    "            out_filts *= 2\n",
    "\n",
    "            setattr(self, 'down_block_%d' % idx, down_block)\n",
    "\n",
    "        out_filts = out_filts // 2\n",
    "        in_filts = in_filts // 2\n",
    "        out_filts, in_filts = in_filts, out_filts\n",
    "\n",
    "        for idx in range(depth-1):\n",
    "            up_block = torch.nn.Sequential(\n",
    "                torch.nn.Conv3d(in_filts + out_filts, out_filts, kernel_size=3, padding=1), \n",
    "                torch.nn.ReLU(inplace=True),\n",
    "                torch.nn.Conv3d(out_filts, out_filts, kernel_size=3, padding=1), \n",
    "                torch.nn.ReLU(inplace=True)\n",
    "            )\n",
    "\n",
    "            in_filts = out_filts\n",
    "            out_filts = out_filts // 2\n",
    "\n",
    "            setattr(self, 'up_block_%d' % idx, up_block)\n",
    "\n",
    "        self.final_conv = torch.nn.Conv3d(in_filts, num_classes, kernel_size=1)\n",
    "        self.max_pool = torch.nn.MaxPool3d(2, stride=2)\n",
    "        self.up_sample = torch.nn.Upsample(scale_factor=2)\n",
    "        self.hparams = hparams\n",
    "    \n",
    "    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Forwards the :attr`input_tensor` through the network to obtain a prediction\n",
    "        \n",
    "        Args:\n",
    "            input_tensor: the network's input\n",
    "            \n",
    "        Returns:\n",
    "            torch.Tensor: the networks output given the :attr`input_tensor`\n",
    "        \"\"\"\n",
    "        depth = self.hparams.get('depth', 3)\n",
    "\n",
    "        intermediate_outputs = []\n",
    "\n",
    "        # Compute all the encoder blocks' outputs\n",
    "        for idx in range(depth):\n",
    "            intermed = getattr(self, 'down_block_%d' % idx)(input_tensor)\n",
    "            if idx < depth - 1:\n",
    "                # store intermediate values for usage in decoder\n",
    "                intermediate_outputs.append(intermed)\n",
    "                input_tensor = getattr(self, 'max_pool')(intermed)\n",
    "            else:\n",
    "                input_tensor = intermed\n",
    "\n",
    "        # Compute all the decoder blocks' outputs\n",
    "        for idx in range(depth-1):\n",
    "            input_tensor = getattr(self, 'up_sample')(input_tensor)\n",
    "            \n",
    "            # use intermediate values from encoder\n",
    "            from_down = intermediate_outputs.pop(-1)\n",
    "            intermed = torch.cat([input_tensor, from_down], dim=1)\n",
    "            input_tensor = getattr(self, 'up_block_%d' % idx)(intermed)\n",
    "\n",
    "        return getattr(self, 'final_conv')(input_tensor)\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay, that was easy, right? Now let's just check if everything in our network is fine:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Unet({'num_classes': 2, 'in_channels': 1, 'depth': 2, 'start_filts': 2})\n",
    "print(net(torch.rand(1, 1, 16, 16, 16)).shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So what did we do here? We initialized a network accepting input images with one channel. This network will then predict a segmentation map for 2 classes (of which one is the background class). It does so with 3 resolution stages.\n",
    "\n",
    "---\n",
    "> **Note:** PyTorch expects 3D  inputs of a layer to be a 5D Tensor, that has a batch dimension (the first one) and a channel dimension (the second one) along with the three spatial dimensions (the last ones).\n",
    "---\n",
    "\n",
    "When we tested the network, we forwarded a tensor with random values of size `(1, 1, 16, 16, 16)` through it. The first `1` here is the batch dim, the second `1` the channel dim (as we specified one input channel) and the three `16` are the spatial dimension (depth, height and width).\n",
    "\n",
    "The output has the same dimensons except the channel dimension now holding `2` channels (one per class).\n",
    "\n",
    "### Training Criterions and Metrics\n",
    "\n",
    "For training we will use the combination of [CrossEntropyLoss](https://pytorch.org/docs/stable/nn.html#torch.nn.CrossEntropyLoss) and the SoftDiceLoss (see below).\n",
    "\n",
    "For more details on this, I'd recommend [Jeremy Jordan's Blog on semantic segmentation](https://www.jeremyjordan.me/semantic-segmentation/).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rising\n",
    "from typing import Sequence, Optional, Union\n",
    "import torch\n",
    "\n",
    "# Taken from https://github.com/justusschock/dl-utils/blob/master/dlutils/losses/soft_dice.py\n",
    "class SoftDiceLoss(torch.nn.Module):\n",
    "    \"\"\"Soft Dice Loss\"\"\"\n",
    "    def __init__(self, square_nom: bool = False, \n",
    "                 square_denom: bool = False, \n",
    "                 weight: Optional[Union[Sequence, torch.Tensor]] = None,\n",
    "                 smooth: float = 1.):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            square_nom: whether to square the nominator\n",
    "            square_denom: whether to square the denominator\n",
    "            weight: additional weighting of individual classes\n",
    "            smooth: smoothing for nominator and denominator\n",
    "    \n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.square_nom = square_nom\n",
    "        self.square_denom = square_denom\n",
    "\n",
    "        self.smooth = smooth\n",
    "\n",
    "        if weight is not None:\n",
    "            if not isinstance(weight, torch.Tensor):\n",
    "                weight = torch.tensor(weight)\n",
    "                \n",
    "            self.register_buffer(\"weight\", weight)\n",
    "        else:\n",
    "            self.weight = None\n",
    "\n",
    "    def forward(self, predictions: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Computes SoftDice Loss\n",
    "        \n",
    "        Args:\n",
    "            predictions: the predictions obtained by the network\n",
    "            targets: the targets (ground truth) for the :attr:`predictions`\n",
    "            \n",
    "        Returns:\n",
    "            torch.Tensor: the computed loss value\n",
    "        \"\"\"\n",
    "        # number of classes for onehot\n",
    "        n_classes = predictions.shape[1]\n",
    "        with torch.no_grad():\n",
    "            targets_onehot = rising.transforms.functional.channel.one_hot_batch(\n",
    "                targets.unsqueeze(1), num_classes=n_classes)\n",
    "        # sum over spatial dimensions\n",
    "        dims = tuple(range(2, predictions.dim()))\n",
    "\n",
    "        # compute nominator\n",
    "        if self.square_nom:\n",
    "            nom = torch.sum((predictions * targets_onehot.float()) ** 2, dim=dims)\n",
    "        else:\n",
    "            nom = torch.sum(predictions * targets_onehot.float(), dim=dims)\n",
    "        nom = 2 * nom + self.smooth\n",
    "\n",
    "        # compute denominator\n",
    "        if self.square_denom:\n",
    "            i_sum = torch.sum(predictions ** 2, dim=dims)\n",
    "            t_sum = torch.sum(targets_onehot ** 2, dim=dims)\n",
    "        else:\n",
    "            i_sum = torch.sum(predictions, dim=dims)\n",
    "            t_sum = torch.sum(targets_onehot, dim=dims)\n",
    "\n",
    "        denom = i_sum + t_sum.float() + self.smooth\n",
    "\n",
    "        # compute loss\n",
    "        frac = nom / denom\n",
    "\n",
    "        # apply weight for individual classesproperly\n",
    "        if self.weight is not None:\n",
    "            frac = self.weight * frac\n",
    "\n",
    "        # average over classes\n",
    "        frac = - torch.mean(frac, dim=1)\n",
    "\n",
    "        return frac\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay, now that we are able to properly calculate the loss function, we still lack a metric to monitor, that describes our performance. For segmentation tasks, this usually comes down to the [dice coefficient](https://en.wikipedia.org/wiki/S%C3%B8rensen%E2%80%93Dice_coefficient). So let's implement this one as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Taken from https://github.com/justusschock/dl-utils/blob/master/dlutils/metrics/dice.py\n",
    "def binary_dice_coefficient(pred: torch.Tensor, gt: torch.Tensor,\n",
    "                            thresh: float = 0.5, smooth: float = 1e-7) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    computes the dice coefficient for a binary segmentation task\n",
    "    \n",
    "    Args:\n",
    "        pred: predicted segmentation (of shape Nx(Dx)HxW)\n",
    "        gt: target segmentation (of shape NxCx(Dx)HxW)\n",
    "        thresh: segmentation threshold\n",
    "        smooth: smoothing value to avoid division by zero\n",
    "        \n",
    "    Returns:\n",
    "        torch.Tensor: dice score\n",
    "    \"\"\"\n",
    "    \n",
    "    assert pred.shape == gt.shape\n",
    "    \n",
    "    pred_bool = pred > thresh\n",
    "\n",
    "    intersec = (pred_bool * gt).float()\n",
    "    return 2 * intersec.sum() / (pred_bool.float().sum()\n",
    "                                 + gt.float().sum() + smooth)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Neat! So far we defined all criterions and metrics necessary for proper training and monitoring. But there are still two major parts of our pipeline missing:\n",
    "\n",
    "1.) Data Preprocessing and Augmentation\n",
    "\n",
    "2.) what to do for parameter update\n",
    "\n",
    "Let's deal with the first point now.\n",
    "\n",
    "### Data Preprocessing\n",
    "\n",
    "Since all samples in our dataset are of different size, we cannot collate them to a batch directly. Instead we need to resize them. Frameworks like [torchvision](https://github.com/pytorch/vision) do this inside the dataset. With `rising` however, we opted for moving this part outside the dataset (but still apply it on each sample separately before batching) for these reasons.\n",
    "\n",
    "1.) The dataset get's more reusable for different settings\n",
    "\n",
    "2.) The transforms don't have to be implemented into each dataset, which means it is easier to switch datasets without code duplication\n",
    "\n",
    "3.) Applying different transforms is as easy as changing an argument of the loader; no need to deal with this manually in the dataset\n",
    "\n",
    "This kind of transforms kann be passed to the dataloader with `sample_transforms`. If you have an implementation that usually works on batched data, we got you. All you need to do is specifying `pseudo_batch_dim` and we will take care of the rest. We will then automatically add a pseudo batch dim to all kind of data (tensors, arrays and all kind of built-in python containers containing a mixture thereof) before applying these transforms and remove it afterwards.\n",
    "\n",
    "For now, we use our batched implementation of native torch resizing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rising.transforms import Compose, ResizeNative\n",
    "\n",
    "def common_per_sample_trafos():\n",
    "        return Compose(ResizeNative(size=(32, 64, 32), keys=('data',), mode='trilinear'),\n",
    "                       ResizeNative(size=(32, 64, 32), keys=('label',), mode='nearest'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "> **Note:** We used different interpolation modes for image ('data' key) and mask ('label' key). Otherwise we could have also specified this with a single transform.\n",
    "---\n",
    "\n",
    "### Data Augmentation\n",
    "\n",
    "Now that we have defined our preprocessing, let's come to data augmentation. To enrich our dataset, we randomly apply an affine. While `rising` already contains an implementation of Affine transforms that can also handle random inputs pretty well, we will implement a basic random parameter sampling by ourselves, since this also serves as an educational example.\n",
    "\n",
    "Basically this is really straight forward. We just derive the `BaseAffine` class, overwrite the way the matrix is assembled by adding the sampling before we call the actual assembly method. We leave the rest to the already defined class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "outputPrepend"
    ]
   },
   "outputs": [],
   "source": [
    "from rising.transforms.affine import BaseAffine\n",
    "import random\n",
    "from typing import Optional, Sequence\n",
    "\n",
    "class RandomAffine(BaseAffine):\n",
    "    \"\"\"Base Affine with random parameters for scale, rotation and translation\"\"\"\n",
    "    def __init__(self, scale_range: Optional[tuple] = None, \n",
    "                 rotation_range: Optional[tuple] = None, \n",
    "                 translation_range: Optional[tuple] = None,\n",
    "                 degree: bool = True,\n",
    "                 image_transform: bool = True,\n",
    "                 keys: Sequence = ('data',),\n",
    "                 grad: bool = False,\n",
    "                 output_size: Optional[tuple] = None,\n",
    "                 adjust_size: bool = False,\n",
    "                 interpolation_mode: str = 'nearest',\n",
    "                 padding_mode: str = 'zeros',\n",
    "                 align_corners: bool = False,\n",
    "                 reverse_order: bool = False,\n",
    "                 **kwargs,):\n",
    "        \n",
    "        \"\"\"\n",
    "        Args:\n",
    "            scale_range: tuple containing minimum and maximum values for scale. \n",
    "                Actual values will be sampled from uniform distribution with these \n",
    "                constraints.\n",
    "            rotation_range: tuple containing minimum and maximum values for rotation. \n",
    "                Actual values will be sampled from uniform distribution with these \n",
    "                constraints.\n",
    "            translation_range: tuple containing minimum and maximum values for translation. \n",
    "                Actual values will be sampled from uniform distribution with these \n",
    "                constraints.\n",
    "            keys: keys which should be augmented\n",
    "            grad: enable gradient computation inside transformation\n",
    "            degree: whether the given rotation(s) are in degrees.\n",
    "                Only valid for rotation parameters, which aren't passed\n",
    "                as full transformation matrix.\n",
    "            output_size: if given, this will be the resulting image size.\n",
    "                Defaults to ``None``\n",
    "            adjust_size: if True, the resulting image size will be\n",
    "                calculated dynamically to ensure that the whole image fits.\n",
    "            interpolation_mode: interpolation mode to calculate output values\n",
    "                ``'bilinear'`` | ``'nearest'``. Default: ``'bilinear'``\n",
    "            padding_mode: padding mode for outside grid values\n",
    "                ``'zeros'`` | ``'border'`` | ``'reflection'``.\n",
    "                Default: ``'zeros'``\n",
    "            align_corners: Geometrically, we consider the pixels of the\n",
    "                input as squares rather than points. If set to True,\n",
    "                the extrema (-1 and 1)  are considered as referring to the\n",
    "                center points of the input’s corner pixels. If set to False,\n",
    "                they are instead considered as referring to the corner points\n",
    "                of the input’s corner pixels, making the sampling more\n",
    "                resolution agnostic.\n",
    "            reverse_order: reverses the coordinate order of the\n",
    "                transformation to conform to the pytorch convention:\n",
    "                transformation params order [W,H(,D)] and\n",
    "                batch order [(D,)H,W]\n",
    "            **kwargs: additional keyword arguments passed to the\n",
    "                affine transf\n",
    "        \"\"\"\n",
    "        super().__init__(scale=None, rotation=None, translation=None, \n",
    "                         degree=degree,\n",
    "                         image_transform=image_transform, \n",
    "                         keys=keys, \n",
    "                         grad=grad, \n",
    "                         output_size=output_size, \n",
    "                         adjust_size=adjust_size, \n",
    "                         interpolation_mode=interpolation_mode, \n",
    "                         padding_mode=padding_mode, \n",
    "                         align_corners=align_corners, \n",
    "                         reverse_order=reverse_order, \n",
    "                         **kwargs)\n",
    "        \n",
    "        self.scale_range = scale_range\n",
    "        self.rotation_range = rotation_range\n",
    "        self.translation_range = translation_range\n",
    "        \n",
    "    def assemble_matrix(self, **data) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Samples Parameters for scale, rotation and translation\n",
    "        before actual matrix assembly.\n",
    "        \n",
    "        Args:\n",
    "            **data: dictionary containing a batch\n",
    "            \n",
    "        Returns:\n",
    "            torch.Tensor: assembled affine matrix\n",
    "        \"\"\"\n",
    "        ndim = data[self.keys[0]].ndim - 2\n",
    "        \n",
    "        if self.scale_range is not None:\n",
    "            self.scale = [random.uniform(*self.scale_range) for _ in range(ndim)]\n",
    "            \n",
    "        if self.translation_range is not None:\n",
    "            self.translation = [random.uniform(*self.translation_range) for _ in range(ndim)]\n",
    "            \n",
    "        if self.rotation_range is not None:\n",
    "            if ndim == 3:\n",
    "                self.rotation = [random.uniform(*self.rotation_range) for _ in range(ndim)]\n",
    "            elif ndim == 1:\n",
    "                self.rotation = random.uniform(*self.rotation_range)\n",
    "            \n",
    "        return super().assemble_matrix(**data)\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Also not that hard... So, now we have a custom implementation of a randomly parametrized affine transformation. This is all we will use as data augmentation for now.\n",
    "\n",
    "Batched Transforms that shall be executed on CPU in a multiprocessed way should be specified to the dataloader as `batch_transforms`. If they should be executed on GPU, you can pass them as `gpu_transforms`. Unfortnuately it is not possible to add GPU transforms in a multiprocessing environment. Thus the internal computation order is like this:\n",
    "\n",
    "1.) Extract sample from dataset\n",
    "\n",
    "2.) Apply per-sample transforms to it (with or without pseudo batch dim)\n",
    "\n",
    "3.) Collate to batch\n",
    "\n",
    "4.) Apply batch transforms\n",
    "\n",
    "5.) Apply GPU transforms\n",
    "\n",
    "Steps 1.-4. can be executed in a multiprocessing environment. If this is the case, the results will be synced back to the main process before applying GPU transforms.\n",
    "\n",
    "### Training Logic\n",
    "The only remaining step is now to integrate this to the training logic of `PyTorchLightning`.\n",
    "\n",
    "The only things we did not yet discuss is how to setup optimizers, logging and train/validation step.\n",
    "\n",
    "The optimizer setup is done by a function `configure_optimizers` that should return the created optimizers.\n",
    "\n",
    "Logging can either be done automatically (all values for the key `log` in the dict returned from `validation_epoch_end` and `training_epoch_end` will autoamtically be logged) or manually (explicitly calling the logger in any of these functions). We show both examples below.\n",
    "\n",
    "For setting up the actual training logic we need to specify `training_step` (and `validation_step` for validation).\n",
    "The complete example is below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rising.transforms import NormZeroMeanUnitStd\n",
    "from rising.loading import DataLoader\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "\n",
    "class TrainableUNet(Unet):\n",
    "    \"\"\"A trainable UNet (extends the base class by training logic)\"\"\"\n",
    "    def __init__(self, hparams: Optional[dict] = None):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            hparams: the hyperparameters needed to construct and train the network. \n",
    "                Specifically these are:\n",
    "                * start_filts (int)\n",
    "                * depth (int)\n",
    "                * in_channels (int)\n",
    "                * num_classes (int)\n",
    "                * min_scale (float)\n",
    "                * max_scale (float)\n",
    "                * min_rotation (int, float)\n",
    "                * max_rotation (int, float)\n",
    "                * batch_size (int)\n",
    "                * num_workers(int)\n",
    "                * learning_rate (float)\n",
    "                \n",
    "                For all of them exist usable default parameters.\n",
    "        \"\"\"\n",
    "        if hparams is None:\n",
    "            hparams = {}\n",
    "        super().__init__(hparams)\n",
    "        \n",
    "        # define loss functions\n",
    "        self.dice_loss = SoftDiceLoss(weight=[0., 1.])\n",
    "        self.ce_loss = torch.nn.CrossEntropyLoss()\n",
    "        \n",
    "    def train_dataloader(self) -> DataLoader:\n",
    "        \"\"\"\n",
    "        Specifies the train dataloader\n",
    "        \n",
    "        Returns:\n",
    "            DataLoader: the train dataloader\n",
    "        \"\"\"\n",
    "        # construct dataset\n",
    "        dataset = NiiDataset(train=True, data_dir=data_dir)\n",
    "        \n",
    "        # specify batch transforms\n",
    "        batch_transforms = Compose([\n",
    "            RandomAffine(scale_range=(self.hparams.get('min_scale', 0.9), self.hparams.get('max_scale', 1.1)),\n",
    "                         rotation_range=(self.hparams.get('min_rotation', -10), self.hparams.get('max_rotation', 10)),\n",
    "                        keys=('data', 'label')),\n",
    "            NormZeroMeanUnitStd(keys=('data',))\n",
    "        ])\n",
    "        \n",
    "        # construct loader\n",
    "        dataloader = DataLoader(dataset,\n",
    "                                batch_size=self.hparams.get('batch_size', 1),\n",
    "                                batch_transforms=batch_transforms,\n",
    "                                shuffle=True, \n",
    "                                sample_transforms=common_per_sample_trafos(),\n",
    "                                pseudo_batch_dim=True,\n",
    "                                num_workers=self.hparams.get('num_workers', 4))\n",
    "        return dataloader\n",
    "    \n",
    "    def val_dataloader(self) -> DataLoader:\n",
    "        # construct dataset\n",
    "        dataset = NiiDataset(train=False, data_dir=data_dir)\n",
    "        \n",
    "        # specify batch transforms (no augmentation here)\n",
    "        batch_transforms = NormZeroMeanUnitStd(keys=('data',))\n",
    "        \n",
    "        # construct loader\n",
    "        dataloader = DataLoader(dataset,\n",
    "                                batch_size=self.hparams.get('batch_size', 1),\n",
    "                                batch_transforms=batch_transforms,\n",
    "                                shuffle=False, \n",
    "                                sample_transforms=common_per_sample_trafos(),\n",
    "                                pseudo_batch_dim=True,\n",
    "                                num_workers=self.hparams.get('num_workers', 4))\n",
    "        \n",
    "        return dataloader\n",
    "    \n",
    "    def configure_optimizers(self) -> torch.optim.Optimizer:\n",
    "        \"\"\"\n",
    "        Configures the optimier to use for training\n",
    "        \n",
    "        Returns:\n",
    "            torch.optim.Optimier: the optimizer for updating the model's parameters\n",
    "        \"\"\"\n",
    "        return torch.optim.Adam(self.parameters(), lr=self.hparams.get('learning_rate', 1e-3))\n",
    "    \n",
    "    def training_step(self, batch: dict, batch_idx: int) -> dict:\n",
    "        \"\"\"\n",
    "        Defines the training logic\n",
    "        \n",
    "        Args:\n",
    "            batch: contains the data (inputs and ground truth)\n",
    "            batch_idx: the number of the current batch\n",
    "            \n",
    "        Returns:\n",
    "            dict: the current loss value\n",
    "        \"\"\"\n",
    "        x, y = batch['data'], batch['label']\n",
    "        \n",
    "        # remove channel dim from gt (was necessary for augmentation)\n",
    "        y = y[:, 0].long()\n",
    "        \n",
    "        # obtain predictions\n",
    "        pred = self(x)\n",
    "        softmaxed_pred = torch.nn.functional.softmax(pred, dim=1)\n",
    "        \n",
    "        # Calculate losses\n",
    "        ce_loss = self.ce_loss(pred, y)\n",
    "        dice_loss = self.dice_loss(softmaxed_pred, y)\n",
    "        total_loss = (ce_loss + dice_loss) / 2\n",
    "        \n",
    "        # calculate dice coefficient\n",
    "        dice_coeff = binary_dice_coefficient(torch.argmax(softmaxed_pred, dim=1), y)\n",
    "        \n",
    "        # log values\n",
    "        self.logger.experiment.add_scalar('Train/DiceCoeff', dice_coeff)\n",
    "        self.logger.experiment.add_scalar('Train/CE', ce_loss)\n",
    "        self.logger.experiment.add_scalar('Train/SoftDiceLoss', dice_loss)\n",
    "        self.logger.experiment.add_scalar('Train/TotalLoss', total_loss)\n",
    "        \n",
    "        return {'loss': total_loss}\n",
    "    \n",
    "    def validation_step(self, batch: dict, batch_idx: int) -> dict:\n",
    "        \"\"\"\n",
    "        Defines the validation logic\n",
    "        \n",
    "        Args:\n",
    "            batch: contains the data (inputs and ground truth)\n",
    "            batch_idx: the number of the current batch\n",
    "            \n",
    "        Returns:\n",
    "            dict: the current loss and metric values\n",
    "        \"\"\"\n",
    "        x, y = batch['data'], batch['label']\n",
    "        \n",
    "        # remove channel dim from gt (was necessary for augmentation)\n",
    "        y = y[:, 0].long()\n",
    "        \n",
    "        # obtain predictions\n",
    "        pred = self(x)\n",
    "        softmaxed_pred = torch.nn.functional.softmax(pred, dim=1)\n",
    "        \n",
    "        # calculate losses\n",
    "        ce_loss = self.ce_loss(pred, y)\n",
    "        dice_loss = self.dice_loss(softmaxed_pred, y)\n",
    "        total_loss = (ce_loss + dice_loss) / 2\n",
    "        \n",
    "        # calculate dice coefficient\n",
    "        dice_coeff = binary_dice_coefficient(torch.argmax(softmaxed_pred, dim=1), y)\n",
    "        \n",
    "        # log values\n",
    "        self.logger.experiment.add_scalar('Val/DiceCoeff', dice_coeff)\n",
    "        self.logger.experiment.add_scalar('Val/CE', ce_loss)\n",
    "        self.logger.experiment.add_scalar('Val/SoftDiceLoss', dice_loss)\n",
    "        self.logger.experiment.add_scalar('Val/TotalLoss', total_loss)\n",
    "        \n",
    "        return {'val_loss': total_loss, 'dice': dice_coeff}\n",
    "    \n",
    "    def validation_epoch_end(self, outputs: list) -> dict:\n",
    "        \"\"\"Aggregates data from each validation step\n",
    "        \n",
    "        Args:\n",
    "            outputs: the returned values from each validation step\n",
    "            \n",
    "        Returns:\n",
    "            dict: the aggregated outputs\n",
    "        \"\"\"\n",
    "        mean_outputs = {}\n",
    "        for k in outputs[0].keys():\n",
    "            mean_outputs[k] = torch.stack([x[k] for x in outputs]).mean()\n",
    "        \n",
    "        tqdm.write('Dice: \\t%.3f' % mean_outputs['dice'].item())\n",
    "        return mean_outputs\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Most of this stuff is relevant for `PyTorch Lightning`. But the dataloader setup nicely shows the integration of `rising` with any existing framework working on `PyTorch Dataloaders` (like `PyTorch Lightning` or `PyTorch Ignite`) for batched and sample transforms.\n",
    "\n",
    "---\n",
    "> **Note:** When playing around, you can also try to switch these transforms from CPU to GPU\n",
    "---\n",
    "\n",
    "## Training\n",
    "\n",
    "We've finally finished all the pipeline definition. Now let's just load the tensorboard extension to monitor our training. For this we will define a common output dir for lightning:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir = 'logs'\n",
    "os.makedirs(output_dir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start tensorboard.\n",
    "\n",
    "%reload_ext tensorboard\n",
    "%tensorboard --logdir {output_dir}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now it's finally time to train!\n",
    "\n",
    "---\n",
    "> **Note:** If you're in a colab environment, remember to enable GPU runtime.\n",
    "---\n",
    "\n",
    "On a GPU in colab, training takes approximately 40 seconds per epoch, which is a total of 33 minutes (2000 seconds) for training, if early stopping doesn't kick in. For me it kicks in after 25 Epochs which takes around 16 Minutes on a colab GPU\n",
    "\n",
    "---\n",
    "> **Note:** for proper visualization, you should switch the horizontal axis in tensorboard to \"relative\" once the first values are showing up.\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_lightning.callbacks import EarlyStopping\n",
    "from pytorch_lightning import Trainer\n",
    "\n",
    "early_stop_callback = EarlyStopping(monitor='dice', min_delta=0.001, patience=10, verbose=False, mode='max')\n",
    "\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    gpus = 1\n",
    "else:\n",
    "    gpus = None\n",
    "    \n",
    "nb_epochs = 50\n",
    "num_start_filts = 16\n",
    "num_workers = 4\n",
    "\n",
    "if 'CI' in os.environ:\n",
    "    nb_epochs = 1\n",
    "    num_start_filts = 2\n",
    "    num_workers = 0\n",
    "\n",
    "model = TrainableUNet({'start_filts': num_start_filts, 'num_workers': num_workers})\n",
    "\n",
    "trainer = Trainer(gpus=gpus, default_save_path=output_dir, early_stop_callback=early_stop_callback, max_nb_epochs=nb_epochs)\n",
    "trainer.fit(model)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the end, you should see a dice coefficient of 0.88 after 25 Epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6.10 64-bit ('rising': conda)",
   "language": "python",
   "name": "python361064bitrisingconda39a2201ad26e4c96be348e851446357a"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}