exercise.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dd2abfa6",
   "metadata": {},
   "source": [
    "# Semantic Segmentation\n",
    "\n",
    "<hr style=\"height:2px;\">\n",
    "\n",
    "In this notebook, we adapt our 2D U-Net for better nuclei segmentations in the Kaggle Nuclei dataset.\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "\n",
    "**Specifically, you will:**\n",
    "\n",
    "1. Prepare the 2D U-Net baseline model and validation dataset.\n",
    "2. Implement and use the Dice coefficient as an evaluation metric for the baseline model.\n",
    "3. Improve metrics by experimenting with:\n",
    "    - Data augmentations\n",
    "    - Loss functions\n",
    "    - (bonus) Group Normalization, U-Net architecture\n",
    "</div>\n",
    "Written by William Patton, Valentyna Zinchenko, and Constantin Pape."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b25f192",
   "metadata": {},
   "source": [
    "Our goal is to produce a model that can take an image as input and produce a segmentation as shown in this table.\n",
    "\n",
    "| Image | Mask | Prediction |\n",
    "| :-: | :-: | :-: |\n",
    "| ![image](static/img_0.png) | ![mask](static/mask_0.png) | ![pred](static/pred_0.png) |\n",
    "| ![image](static/img_1.png) | ![mask](static/mask_1.png) | ![pred](static/pred_1.png) |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ddcf625",
   "metadata": {},
   "source": [
    "<hr style=\"height:2px;\">\n",
    "\n",
    "## The libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "699d517b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import subprocess\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "import torch.nn as nn\n",
    "\n",
    "import torchvision.transforms.v2 as transforms_v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8d2042d",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4ec5fac",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# make sure gpu is available. Please call a TA if this cell fails\n",
    "assert torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51b863c1",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## Section 0: What we have so far\n",
    "You have already implemented a U-Net architecture in the previous exercise. We will use it as a starting point for this exercise.\n",
    "You should also alredy have the dataset and the dataloader implemented, along with a simple train loop with MSELoss.\n",
    "Lets go ahead and visualize some of the data along with some predictions to see how we are doing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6b758ab",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "from local import (\n",
    "    NucleiDataset,\n",
    "    show_random_dataset_image,\n",
    "    show_random_dataset_image_with_prediction,\n",
    "    show_random_augmentation_comparison,\n",
    "    train,\n",
    ")\n",
    "from dlmbl_unet import ConvBlock, UNet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4f77c18",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "\n",
    "*Note*: We are artificially making our validation data worse. This dataset\n",
    "was chosen to be reasonable to segment in the amount of time it takes to\n",
    "run this exercise. However this means that some techniques like augmentations\n",
    "aren't as useful as they would be on a more complex dataset. So we are\n",
    "artificially adding noise to the validation data to make it more challenging."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3f52a97",
   "metadata": {},
   "outputs": [],
   "source": [
    "def salt_and_pepper_noise(image, amount=0.05):\n",
    "    \"\"\"\n",
    "    Add salt and pepper noise to an image\n",
    "    \"\"\"\n",
    "    out = image.clone()\n",
    "    num_salt = int(amount * image.numel() * 0.5)\n",
    "    num_pepper = int(amount * image.numel() * 0.5)\n",
    "\n",
    "    # Add Salt noise\n",
    "    coords = [\n",
    "        torch.randint(0, i - 1, [num_salt]) if i > 1 else [0] * num_salt\n",
    "        for i in image.shape\n",
    "    ]\n",
    "    out[coords] = 1\n",
    "\n",
    "    # Add Pepper noise\n",
    "    coords = [\n",
    "        torch.randint(0, i - 1, [num_pepper]) if i > 1 else [0] * num_pepper\n",
    "        for i in image.shape\n",
    "    ]\n",
    "    out[coords] = 0\n",
    "\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "403511d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = NucleiDataset(\"nuclei_train_data\", transforms_v2.RandomCrop(256))\n",
    "train_loader = DataLoader(train_data, batch_size=5, shuffle=True, num_workers=8)\n",
    "val_data = NucleiDataset(\n",
    "    \"nuclei_val_data\",\n",
    "    transforms_v2.RandomCrop(256),\n",
    "    img_transform=transforms_v2.Lambda(salt_and_pepper_noise),\n",
    ")\n",
    "val_loader = DataLoader(val_data, batch_size=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3420cccd",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "unet = UNet(depth=4, in_channels=1, out_channels=1, num_fmaps=2).to(device)\n",
    "loss = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(unet.parameters())\n",
    "\n",
    "for epoch in range(10):\n",
    "    train(unet, train_loader, optimizer, loss, epoch, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92723f6c",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "# Show some predictions on the train data\n",
    "show_random_dataset_image(train_data)\n",
    "show_random_dataset_image_with_prediction(train_data, unet, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b851a7a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show some predictions on the validation data\n",
    "show_random_dataset_image(val_data)\n",
    "show_random_dataset_image_with_prediction(val_data, unet, device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34df0ddc",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "    <p><b>Task 0.1</b>: Are the predictions good enough? Take some time to try to think about\n",
    "    what could be improved and how that could be addressed. If you have time try training a second\n",
    "    model and see which one is better</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09f042c5",
   "metadata": {},
   "source": [
    "Write your answers here:\n",
    "<ol>\n",
    "    <li></li>\n",
    "    <li></li>\n",
    "    <li></li>\n",
    "</ol>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a773609",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">\n",
    "<h2> Checkpoint 0 </h2>\n",
    "<p>We will go over the steps up to this point soon. By this point you should have imported and re-used\n",
    "code from previous exercises to train a basic UNet.</p>\n",
    "<p>The rest of this exercise will focus on tailoring our network to semantic segmentation to improve\n",
    "performance. The main areas we will tackle are:</p>\n",
    "<ol>\n",
    "  <li> Evaluation\n",
    "  <li> Augmentation\n",
    "  <li> Activations/Loss Functions\n",
    "</ol>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b57172c",
   "metadata": {},
   "source": [
    "<hr style=\"height:2px;\">\n",
    "\n",
    "## Section 1: Evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a3e5a2b",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "source": [
    "One of the most important parts of training a model is evaluating it. We need to know how well our model is doing and if it is improving.\n",
    "We will start by implementing a metric to evaluate our model. Evaluation is always specific to the task, in this case semantic segmentation.\n",
    "We will use the [Dice Coefficient](https://en.wikipedia.org/wiki/S%C3%B8rensen%E2%80%93Dice_coefficient) to evaluate the network predictions.\n",
    "We can use it for validation if we interpret set $a$ as predictions and $b$ as labels. It is often used to evaluate segmentations with sparse\n",
    "foreground, because the denominator normalizes by the number of foreground pixels.\n",
    "The Dice Coefficient is closely related to Jaccard Index / Intersection over Union."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc324791",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Task 1.1</b>: Fill in implementation details for the Dice Coefficient\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "369a3b07",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sorensen Dice Coefficient implemented in torch\n",
    "# the coefficient takes values in two discrete arrays\n",
    "# with values in {0, 1}, and produces a score in [0, 1]\n",
    "# where 0 is the worst score, 1 is the best score\n",
    "class DiceCoefficient(nn.Module):\n",
    "    def __init__(self, eps=1e-6):\n",
    "        super().__init__()\n",
    "        self.eps = eps\n",
    "\n",
    "    # the dice coefficient of two sets represented as vectors a, b can be\n",
    "    # computed as (2 *|a b| / (a^2 + b^2))\n",
    "    def forward(self, prediction, target):\n",
    "        intersection = ...\n",
    "        union = ...\n",
    "        return 2 * intersection / union.clamp(min=self.eps)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16269191",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    Test your Dice Coefficient here, are you getting the right scores?\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db6f6548",
   "metadata": {},
   "outputs": [],
   "source": [
    "dice = DiceCoefficient()\n",
    "target = torch.tensor([0.0, 1.0])\n",
    "good_prediction = torch.tensor([0.0, 1.0])\n",
    "bad_prediction = torch.tensor([0.0, 0.0])\n",
    "wrong_prediction = torch.tensor([1.0, 0.0])\n",
    "\n",
    "assert dice(good_prediction, target) == 1.0, dice(good_prediction, target)\n",
    "assert dice(bad_prediction, target) == 0.0, dice(bad_prediction, target)\n",
    "assert dice(wrong_prediction, target) == 0.0, dice(wrong_prediction, target)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3ef161b",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Task 1.2</b>: What happens if your predictions are not discrete elements of {0,1}?\n",
    "    <ol>\n",
    "        <li>What if the predictions are in range (0,1)?</li>\n",
    "        <li>What if the predictions are in range ($-\\infty$,$\\infty$)?</li>\n",
    "    </ol>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "778c4989",
   "metadata": {},
   "source": [
    "Answer:\n",
    "1) ...\n",
    "\n",
    "2) ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd6d7987",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">\n",
    "    <h2>Checkpoint 1.1 </h2>\n",
    "\n",
    "This is a good place to stop for a moment. If you have extra time look into some extra\n",
    "evaluation functions or try to implement your own without hints.\n",
    "Some popular alternatives to the Dice Coefficient are the Jaccard Index and Balanced F1 Scores.\n",
    "You may even have time to compute the evaluation score between some of your training and\n",
    "validation predictions to their ground truth using our previous models.\n",
    "\n",
    "</div>\n",
    "\n",
    "<hr style=\"height:2px;\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8cef2e6",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>Task 1.3</b>: Fix in all the TODOs to make the validate function work. If confused, you can use this\n",
    "<a href=\"https://pytorch.org/tutorials/beginner/basics/optimization_tutorial.html\">PyTorch tutorial</a> as a template\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "330e6aa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run validation after training epoch\n",
    "def validate(\n",
    "    model,\n",
    "    loader,\n",
    "    loss_function,\n",
    "    metric,\n",
    "    step=None,\n",
    "    tb_logger=None,\n",
    "    device=None,\n",
    "):\n",
    "    if device is None:\n",
    "        # You can pass in a device or we will default to using\n",
    "        # the gpu. Feel free to try training on the cpu to see\n",
    "        # what sort of performance difference there is\n",
    "        if torch.cuda.is_available():\n",
    "            device = torch.device(\"cuda\")\n",
    "        else:\n",
    "            device = torch.device(\"cpu\")\n",
    "\n",
    "    # set model to eval mode\n",
    "    model.eval()\n",
    "    model.to(device)\n",
    "\n",
    "    # running loss and metric values\n",
    "    val_loss = 0\n",
    "    val_metric = 0\n",
    "\n",
    "    # disable gradients during validation\n",
    "    with torch.no_grad():\n",
    "        # iterate over validation loader and update loss and metric values\n",
    "        for x, y in loader:\n",
    "            x, y = x.to(device), y.to(device)\n",
    "            # TODO: evaluate this example with the given loss and metric\n",
    "            prediction = ...\n",
    "            # We *usually* want the target to be the same type as the prediction\n",
    "            # however this is very dependent on your choice of loss function and\n",
    "            # metric. If you get errors such as \"RuntimeError: Found dtype Float but expected Short\"\n",
    "            # then this is where you should look.\n",
    "            if y.dtype != prediction.dtype:\n",
    "                y = y.type(prediction.dtype)\n",
    "            val_loss += ...\n",
    "            val_metric += ...\n",
    "\n",
    "    # normalize loss and metric\n",
    "    val_loss /= len(loader)\n",
    "    val_metric /= len(loader)\n",
    "\n",
    "    if tb_logger is not None:\n",
    "        assert (\n",
    "            step is not None\n",
    "        ), \"Need to know the current step to log validation results\"\n",
    "        tb_logger.add_scalar(tag=\"val_loss\", scalar_value=val_loss, global_step=step)\n",
    "        tb_logger.add_scalar(\n",
    "            tag=\"val_metric\", scalar_value=val_metric, global_step=step\n",
    "        )\n",
    "        # we always log the last validation images\n",
    "        tb_logger.add_images(tag=\"val_input\", img_tensor=x.to(\"cpu\"), global_step=step)\n",
    "        tb_logger.add_images(tag=\"val_target\", img_tensor=y.to(\"cpu\"), global_step=step)\n",
    "        tb_logger.add_images(\n",
    "            tag=\"val_prediction\", img_tensor=prediction.to(\"cpu\"), global_step=step\n",
    "        )\n",
    "\n",
    "    print(\n",
    "        \"\\nValidate: Average loss: {:.4f}, Average Metric: {:.4f}\\n\".format(\n",
    "            val_loss, val_metric\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "616a6834",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>Task 1.4</b>: Evaluate your first model using the Dice Coefficient. How does it perform? If you trained two models,\n",
    "    do the scores agree with your visual determination of which model was better?\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ad64cbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Evaluate your model here\n",
    "validate(...)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b7422a0",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">\n",
    "    <h2>Checkpoint 1.2</h2>\n",
    "\n",
    "We have finished writing the evaluation function. We will go over the code up to this point soon.\n",
    "Next we will work on augmentations to improve the generalization of our model.\n",
    "\n",
    "</div>\n",
    "\n",
    "<hr style=\"height:2px;\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a052b72d",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## Section 2: Augmentation\n",
    "Often our models will perform better on the evaluation dataset if we augment our training data.\n",
    "This is because the model will be exposed to a wider variety of data that will hopefully help\n",
    "cover the full distribution of data in the validation set. We will use the `torchvision.transforms`\n",
    "to augment our data."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f43e31f4",
   "metadata": {},
   "source": [
    "PS: PyTorch already has quite a few possible data transforms, so if you need one, check\n",
    "[here](https://pytorch.org/vision/stable/transforms.html#transforms-on-pil-image-and-torch-tensor).\n",
    "The biggest problem with them is that they are clearly separated into transforms applied to PIL\n",
    "images (remember, we initially load the images as PIL.Image?) and torch.tensors (remember, we\n",
    "converted the images into tensors by calling transforms.ToTensor()?). This can be incredibly\n",
    "annoying if for some reason you might need to transorm your images to tensors before applying any\n",
    "other transforms or you don't want to use PIL library at all."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84a60211",
   "metadata": {},
   "source": [
    "Here is an example augmented dataset. Use it to see how it affects your data, then play around with at least\n",
    "2 other augmentations.\n",
    "There are two types of augmentations: `transform` and `img_transform`. The first one is applied to both the\n",
    "image and the mask, the second is only applied to the image. This is useful if you want to apply augmentations\n",
    "that spatially distort your data and you want to make sure the same distortion is applied to the mask and image.\n",
    "`img_transform` is useful for augmentations that don't make sense to apply to the mask, like blurring."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe57ee09",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = NucleiDataset(\"nuclei_train_data\", transforms_v2.RandomCrop(256))\n",
    "\n",
    "# Note this augmented data uses extreme augmentations for visualization. It will not train well\n",
    "example_augmented_data = NucleiDataset(\n",
    "    \"nuclei_train_data\",\n",
    "    transforms_v2.Compose(\n",
    "        [transforms_v2.RandomRotation(45), transforms_v2.RandomCrop(256)]\n",
    "    ),\n",
    "    img_transform=transforms_v2.Compose([transforms_v2.GaussianBlur(21, sigma=10.0)]),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0e0f307",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_random_augmentation_comparison(train_data, example_augmented_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b326e87",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>Task 2.1</b>: Now create an augmented dataset with an augmentation of your choice.\n",
    "      **hint**: Using the same augmentation as was applied to the validation data will\n",
    "     likely be optimal. Bonus points if you can get good results without the custom noise.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91ed16e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "augmented_data = ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "641e017b",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>Task 2.2</b>: Now retrain your model with your favorite augmented dataset. Did your model improve?\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "491b816a",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "unet = UNet(depth=4, in_channels=1, out_channels=1, num_fmaps=2).to(device)\n",
    "loss = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(unet.parameters())\n",
    "augmented_loader = DataLoader(augmented_data, batch_size=5, shuffle=True, num_workers=8)\n",
    "\n",
    "..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2ba112e",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>Task 2.3</b>: Now evaluate your model. Did your model improve?\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c43986b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "validate(...)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9334b1a9",
   "metadata": {},
   "source": [
    "<hr style=\"height:2px;\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ef67dde",
   "metadata": {},
   "source": [
    "## Section 3: Loss Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6a2d56e",
   "metadata": {},
   "source": [
    "The next step to do would be to improve our loss function - the metric that tells us how\n",
    "close we are to the desired output. This metric should be differentiable, since this\n",
    "is the value to be backpropagated. The are\n",
    "[multiple losses](https://lars76.github.io/2018/09/27/loss-functions-for-segmentation.html)\n",
    "we could use for the segmentation task.\n",
    "\n",
    "Take a moment to think which one is better to use. If you are not sure, don't forget\n",
    "that you can always google! Before you start implementing the loss yourself, take a look\n",
    "at the [losses](https://pytorch.org/docs/stable/nn.html#loss-functions) already implemented\n",
    "in PyTorch. You can also look for implementations on GitHub."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d85f676",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>Task 3.1</b>: Implement your loss (or take one from pytorch):\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f42ee25",
   "metadata": {},
   "outputs": [],
   "source": [
    "# implement your loss here or initialize the one of your choice from PyTorch\n",
    "loss_function: torch.nn.Module = ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82cb365e",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    Test your loss function here, is it behaving as you'd expect?\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df895867",
   "metadata": {},
   "outputs": [],
   "source": [
    "target = torch.tensor([0.0, 1.0])\n",
    "good_prediction = torch.tensor([0.01, 0.99])\n",
    "bad_prediction = torch.tensor([0.4, 0.6])\n",
    "wrong_prediction = torch.tensor([0.9, 0.1])\n",
    "\n",
    "good_loss = loss_function(good_prediction, target)\n",
    "bad_loss = loss_function(bad_prediction, target)\n",
    "wrong_loss = loss_function(wrong_prediction, target)\n",
    "\n",
    "assert good_loss < bad_loss\n",
    "assert bad_loss < wrong_loss\n",
    "\n",
    "# Can your loss function handle predictions outside of (0, 1)?\n",
    "# Some loss functions will be perfectly happy with this which may\n",
    "# make them easier to work with, but predictions outside the expected\n",
    "# range will not work well with our soon to be discussed evaluation metric.\n",
    "out_of_bounds_prediction = torch.tensor([-0.1, 1.1])\n",
    "\n",
    "try:\n",
    "    oob_loss = loss_function(out_of_bounds_prediction, target)\n",
    "    print(\"Your loss supports out-of-bounds predictions.\")\n",
    "except RuntimeError as e:\n",
    "    print(e)\n",
    "    print(\"Your loss does not support out-of-bounds predictions\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c85cd931",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "Pay close attention to whether your loss function can handle predictions outside of the range (0, 1).\n",
    "If it can't, theres a good chance that the activation function requires a specific activation before\n",
    "being passed into the loss function. This is a common source of bugs in DL models. For example, trying\n",
    "to use the `torch.nn.BCEWithLogitsLoss` loss function with a model that has a sigmoid activation will\n",
    "result in abysmal performance, wheras using the `torch.nn.BCELoss` loss function with a model that has\n",
    "no activation function will likely error out and fail to train."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be9d799f",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# Now lets start experimenting. Start a tensorboard logger to keep track of experiments.\n",
    "# start a tensorboard writer\n",
    "logger = SummaryWriter(\"runs/Unet\")\n",
    "\n",
    "\n",
    "# Function to find an available port and launch TensorBoard on the browser\n",
    "def launch_tensorboard(log_dir):\n",
    "    import socket\n",
    "\n",
    "    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:\n",
    "        s.bind((\"\", 0))\n",
    "        port = s.getsockname()[1]\n",
    "\n",
    "    tensorboard_cmd = f\"tensorboard --logdir={log_dir} --port={port}\"\n",
    "    process = subprocess.Popen(tensorboard_cmd, shell=True)\n",
    "    print(\n",
    "        f\"TensorBoard started at http://localhost:{port}. \\n\"\n",
    "        \"If you are using VSCode remote session, forward the port using the PORTS tab next to TERMINAL.\"\n",
    "    )\n",
    "    return process\n",
    "\n",
    "\n",
    "launch_tensorboard(\"runs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88e964c3",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# Use the unet you expect to work the best!\n",
    "model = UNet(\n",
    "    depth=4,\n",
    "    in_channels=1,\n",
    "    out_channels=1,\n",
    "    num_fmaps=2,\n",
    "    final_activation=torch.nn.Sigmoid(),\n",
    ").to(device)\n",
    "\n",
    "# use adam optimizer\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "\n",
    "# build the dice coefficient metric\n",
    "metric = DiceCoefficient()\n",
    "\n",
    "# train for $25$ epochs\n",
    "# during the training you can inspect the\n",
    "# predictions in the tensorboard\n",
    "n_epochs = 25\n",
    "for epoch in range(n_epochs):\n",
    "    # train\n",
    "    train(\n",
    "        model,\n",
    "        train_loader,\n",
    "        optimizer=optimizer,\n",
    "        loss_function=loss_function,\n",
    "        epoch=epoch,\n",
    "        log_interval=25,\n",
    "        tb_logger=logger,\n",
    "        device=device,\n",
    "    )\n",
    "    step = epoch * len(train_loader)\n",
    "    # validate\n",
    "    validate(model, val_loader, loss_function, metric, step=step, tb_logger=logger)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09b13426",
   "metadata": {},
   "source": [
    "Your validation metric was probably around 85% by the end of the training. That sounds good enough,\n",
    "but an equally important thing to check is: Open the Images tab in your Tensorboard and compare\n",
    "predictions to targets. Do your predictions look reasonable? Are there any obvious failure cases?\n",
    "If nothing is clearly wrong, let's see if we can still improve the model performance by changing\n",
    "the model or the loss\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c284be3e",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">\n",
    "    <h2>Checkpoint 3</h2>\n",
    "\n",
    "This is the end of the guided exercise. We will go over all of the code up until this point shortly.\n",
    "While you wait you are encouraged to try alternative loss functions, evaluation metrics, augmentations,\n",
    "and networks. After this come additional exercises if you are interested and have the time.\n",
    "\n",
    "</div>\n",
    "<hr style=\"height:2px;\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c2f4187",
   "metadata": {},
   "source": [
    "## Additional Exercises\n",
    "\n",
    "1. Modify and evaluate the following architecture variants of the U-Net:\n",
    "    * use [GroupNorm](https://pytorch.org/docs/stable/nn.html#torch.nn.GroupNorm) to normalize convolutional group inputs\n",
    "    * use more layers in your U-Net.\n",
    "\n",
    "2. Use the Dice Coefficient as loss function. Before we only used it for validation, but it is differentiable\n",
    "and can thus also be used as loss. Compare to the results from exercise 2.\n",
    "Hint: The optimizer we use finds minima of the loss, but the minimal value for the Dice coefficient corresponds\n",
    "to a bad segmentation. How do we need to change the Dice Coefficient to use it as loss nonetheless?\n",
    "\n",
    "3. Compare the results of these trainings to the first one. If any of the modifications you've implemented show\n",
    "better results, combine them (e.g. add both GroupNorm and one more layer) and run trainings again.\n",
    "What is the best result you could get?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "617d238f",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>Task BONUS.1</b>: Modify the ConvBlockGN class in bonus_unet.py to include GroupNorm layers. Then update the UNetGN class to use the modified ConvBlock\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa8eb08d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# See the original U-Net for an example of how to build the convolutional block\n",
    "# We want operation -> activation -> normalization (2x)\n",
    "# Hint: Group norm takes a \"num_groups\" argument. Use 2 to match the solution\n",
    "# Task: Modify the bonus_unet.py file as needed and save the changes before you run this cell\n",
    "\n",
    "from bonus_unet import ConvBlockGN, UNetGN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7fd9357",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "model = UNetGN(\n",
    "    depth=4,\n",
    "    in_channels=1,\n",
    "    out_channels=1,\n",
    "    num_fmaps=2,\n",
    "    final_activation=torch.nn.Sigmoid(),\n",
    ").to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "\n",
    "metric = DiceCoefficient()\n",
    "\n",
    "logger = SummaryWriter(\"runs/UNetGN\")\n",
    "\n",
    "\n",
    "# train for 40 epochs\n",
    "# during the training you can inspect the\n",
    "# predictions in the tensorboard\n",
    "n_epochs = 40\n",
    "for epoch in range(n_epochs):\n",
    "    train(\n",
    "        model,\n",
    "        train_loader,\n",
    "        optimizer=optimizer,\n",
    "        loss_function=loss_function,\n",
    "        epoch=epoch,\n",
    "        log_interval=5,\n",
    "        tb_logger=logger,\n",
    "        device=device,\n",
    "    )\n",
    "    step = epoch * len(train_loader)\n",
    "    validate(\n",
    "        model,\n",
    "        val_loader,\n",
    "        loss_function,\n",
    "        metric,\n",
    "        step=step,\n",
    "        tb_logger=logger,\n",
    "        device=device,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a6f7516",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>Task BONUS.2</b>: More Layers\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24d9418e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Experiment with more layers. For example UNet with depth 5\n",
    "\n",
    "model = ...\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "\n",
    "metric = DiceCoefficient()\n",
    "\n",
    "loss = torch.nn.BCELoss()\n",
    "\n",
    "logger = SummaryWriter(\"runs/UNet5layers\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b3efb4a",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# train for 25 epochs\n",
    "# during the training you can inspect the\n",
    "# predictions in the tensorboard\n",
    "n_epochs = 25\n",
    "for epoch in range(n_epochs):\n",
    "    train(\n",
    "        model,\n",
    "        train_loader,\n",
    "        optimizer=optimizer,\n",
    "        loss_function=loss,\n",
    "        epoch=epoch,\n",
    "        log_interval=5,\n",
    "        tb_logger=logger,\n",
    "        device=device,\n",
    "    )\n",
    "    step = epoch * len(train_loader)\n",
    "    validate(\n",
    "        model, val_loader, loss, metric, step=step, tb_logger=logger, device=device\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8c10429",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>Task BONUS.3</b>: Dice Loss\n",
    "    Dice Loss is a simple inversion of the Dice Coefficient.\n",
    "    We already have a Dice Coefficient implementation, so now we just\n",
    "    need a layer that can invert it.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55614406",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DiceLoss(nn.Module):\n",
    "    \"\"\" \"\"\"\n",
    "\n",
    "    def __init__(self, offset: float = 1):\n",
    "        super().__init__()\n",
    "        self.dice_coefficient = DiceCoefficient()\n",
    "\n",
    "    def forward(self, x, y): ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2877dc6",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# Now combine the Dice Coefficient layer with the Invert layer to make a Dice Loss\n",
    "dice_loss = ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bfcc07a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Experiment with Dice Loss\n",
    "net = ...\n",
    "optimizer = ...\n",
    "metric = ...\n",
    "loss_func = ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97fd4f67",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "logger = SummaryWriter(\"runs/UNet_diceloss\")\n",
    "\n",
    "n_epochs = 40\n",
    "for epoch in range(n_epochs):\n",
    "    train(\n",
    "        net,\n",
    "        train_loader,\n",
    "        optimizer=optimizer,\n",
    "        loss_function=loss_func,\n",
    "        epoch=epoch,\n",
    "        log_interval=5,\n",
    "        tb_logger=logger,\n",
    "        device=device,\n",
    "    )\n",
    "    step = epoch * len(train_loader)\n",
    "    validate(\n",
    "        net, val_loader, loss_func, metric, step=step, tb_logger=logger, device=device\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c430a598",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>Task BONUS.4</b>: Group Norm + Dice\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26d9910a",
   "metadata": {},
   "outputs": [],
   "source": [
    "net = ...\n",
    "optimizer = ...\n",
    "metric = ...\n",
    "loss_func = ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7437268",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "logger = SummaryWriter(\"runs/UNetGN_diceloss\")\n",
    "\n",
    "n_epochs = 40\n",
    "for epoch in range(n_epochs):\n",
    "    train(\n",
    "        net,\n",
    "        train_loader,\n",
    "        optimizer=optimizer,\n",
    "        loss_function=loss_func,\n",
    "        epoch=epoch,\n",
    "        log_interval=5,\n",
    "        tb_logger=logger,\n",
    "        device=device,\n",
    "    )\n",
    "    step = epoch * len(train_loader)\n",
    "    validate(\n",
    "        net, val_loader, loss_func, metric, step=step, tb_logger=logger, device=device\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95d4f82b",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>Task BONUS.5</b>: Group Norm + Dice + U-Net 5 Layers\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3da401ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "net = ...\n",
    "optimizer = ...\n",
    "metric = ...\n",
    "loss_func = ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8eed5161",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "all"
  },
  "kernelspec": {
   "display_name": "03-semantic-segmentation",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}