unet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1ga7iKpg5UoI"
   },
   "source": [
    "# Cancer Detection and UNet\n",
    "---\n",
    "\n",
    "In this lab, we will learn on a dataset for a real-world task. We will be performing cancer detection: given a photo of a slide of cells, identify the areas of cancerous cells. We will be performing image segmentation: classifying each pixel as cancerous or non-cancerous.\n",
    "\n",
    "You will create a large network, using the U-net architecture described below. The main goal of this lab is for you to understand how to create complex network architectures capable of learning difficult tasks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xnHpsA3-5UoK"
   },
   "source": [
    "---\n",
    "## Grading Criteria\n",
    "* 20% Completed dataset analysis\n",
    "* 20% Correct training and validation functions\n",
    "* 10% Proper creation, training, and validation of smallest-possible network for this task\n",
    "* 40% Proper creation, training, and validation of the U-net architecture\n",
    "* 10% Reasonable validation results achieved with the U-net architecture\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wQOefmcZVgTl",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import gzip\n",
    "import tarfile\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader, Subset\n",
    "from torchvision import transforms, utils, datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2gskL3Ii5UoM"
   },
   "source": [
    "Set the global variable `device = torch.device(\"cpu\")`. You can change this to `cuda` once you know everything is working, that way you don't run out of your GPU resources on Colab.\n",
    "\n",
    "Remember to use a CPU runtime when not using the GPU. Colab has strict limits on how long you can use a GPU runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tUxZzf3E5UoN"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "49quyvIJvTq8"
   },
   "source": [
    "WARNING: You may run into an error that says `RuntimeError: CUDA out of memory.`\n",
    "\n",
    "If you get this error immediately when training, then your batch is probably larger than what the GPU is capable of. You can solve this problem by adjusting the image size or the batch size and then restarting the runtime.\n",
    "If you get this error after training for multiple steps then you probably are not freeing up computation graph, i.e. you are storing `loss` instead of `loss.item()`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Il_53HLSWPTY"
   },
   "outputs": [],
   "source": [
    "class CancerDataset(Dataset):\n",
    "  def __init__(self, root, download=True, size=512, train=True):\n",
    "    if download and not os.path.exists(os.path.join(root, 'cancer_data')):\n",
    "      datasets.utils.download_url('http://liftothers.org/cancer_data.tar.gz', root, 'cancer_data.tar.gz', None)\n",
    "      self.extract_gzip(os.path.join(root, 'cancer_data.tar.gz'))\n",
    "      self.extract_tar(os.path.join(root, 'cancer_data.tar'))\n",
    "\n",
    "    postfix = 'train' if train else 'test'\n",
    "    root = os.path.join(root, 'cancer_data', 'cancer_data')\n",
    "    self.dataset_folder = datasets.ImageFolder(os.path.join(root, 'inputs_' + postfix) ,transform = transforms.Compose([transforms.Resize(size),transforms.ToTensor()]))\n",
    "    self.label_folder = datasets.ImageFolder(os.path.join(root, 'outputs_' + postfix) ,transform = transforms.Compose([transforms.Resize(size),transforms.ToTensor()]))\n",
    "\n",
    "  @staticmethod\n",
    "  def extract_gzip(gzip_path, remove_finished=False):\n",
    "    print('Extracting {}'.format(gzip_path))\n",
    "    with open(gzip_path.replace('.gz', ''), 'wb') as out_f, gzip.GzipFile(gzip_path) as zip_f:\n",
    "      out_f.write(zip_f.read())\n",
    "    if remove_finished:\n",
    "      os.unlink(gzip_path)\n",
    "\n",
    "  @staticmethod\n",
    "  def extract_tar(tar_path):\n",
    "    print('Untarring {}'.format(tar_path))\n",
    "    z = tarfile.TarFile(tar_path)\n",
    "    z.extractall(tar_path.replace('.tar', ''))\n",
    "\n",
    "  def __getitem__(self,index):\n",
    "    img = self.dataset_folder[index]\n",
    "    label = self.label_folder[index]\n",
    "    return img[0],label[0][0]\n",
    "\n",
    "  def __len__(self):\n",
    "    return len(self.dataset_folder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_DOVQOsr5UoO"
   },
   "source": [
    "---\n",
    "\n",
    "# Part 1: Analyzing the Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zptcyjuu5UoO"
   },
   "source": [
    "First things first.\n",
    "Let's take a look at the dataset.\n",
    "\n",
    "*Note: This may take a while to download; it may be quicker to upload it to Colab from your machine.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "eOtZKw4A5UoP",
    "outputId": "4feb814a-1236-4715-ac78-d4d70ab9e11c"
   },
   "outputs": [],
   "source": [
    "train_dataset = CancerDataset(\"/tmp/Datasets/cancer\", train=True)\n",
    "val_dataset = CancerDataset(\"/tmp/Datasets/cancer\", train=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gpTt-M2J5UoP"
   },
   "source": [
    "Print out the lengths of your train and validation datasets, so you know how much data you are working with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "pfuThS5a5UoP",
    "outputId": "9928598c-34db-4d7c-84e2-2f98b81d8cb0"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JpKQHdAK5UoQ"
   },
   "source": [
    "Now you should examine a single data instance from the dataset.\n",
    "We'll use pos_test_000072.png (`val_dataset[172]`) as an example throughout training.\n",
    "Collect the image and the cancer labels from pos_test_000072.png (i.e. `val_x, val_y = val_dataset[172]`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RR-zYwwn5UoQ"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Pl47ybmY5UoQ"
   },
   "source": [
    "Let's analyze the `x` data first. Print out the shape and dtype of `val_x` below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "B3UHZ7Ka5UoQ",
    "outputId": "0fc9b212-ecfc-4f1d-89e3-532763eba6cc"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NBvCiRLf5UoR"
   },
   "source": [
    "`x` has 3 channels, which likely means that `x` is an RGB image (which it is).\n",
    "Since RGB images can be values between [0, 255] or [0, 1] we should print out the min and max values of x to see which.\n",
    "Print out the min and max values of x below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "tb-HJFyV5UoR",
    "outputId": "a9bacf7c-10ad-41c5-f3e8-cdc50ff1a1f2"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_m-Xvc4e5UoR"
   },
   "source": [
    "The values are between [0, 1], which is good because that means we don't have do any normalization.\n",
    "\n",
    "Now implement a `plot_image()` function which should plot an `x` variable.\n",
    "You can use the `plt.imshow()` function to plot an RGB image.\n",
    "\n",
    "Plot `val_x` and validate that it looks the same as the image file in \"tmp/Datasets/cancer/cancer_data/cancer_data/inputs_test/0/pos_test_000072.png\" in the dataset.\n",
    "\n",
    "*Note: Per PyTorch convention an image has its channels at the 0th dimension (ignoring the batch dimension), but matplotlib expects the channel dimension at the 2nd dimension. You will need to permute the dimensions the plot it correctly*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 435
    },
    "id": "x9CUtuoL5UoR",
    "outputId": "5bbcbd6d-247c-44b5-e101-1ef02a6d37bc"
   },
   "outputs": [],
   "source": [
    "def plot_image(x):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZG8GdOWW5UoR"
   },
   "source": [
    "Now let's have a look at the cancer labels `y`; print out the shape and dtype of `val_y`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "VGMqrnku5UoR",
    "outputId": "741f7bee-9995-4d76-b816-a474419207b1"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iAm_r6Ab5UoS"
   },
   "source": [
    "`y` has 1 channel (but it is squeezed so there is no channel dimension).\n",
    "This most likely means that `y` is a grayscale image (which it is).\n",
    "\n",
    "Implement `plot_cancer()` which should plot a `y` variable. You can use the `plt.imshow()` function to plot a grayscale image (pass in the `cmap=\"gray\"` argument to make the image grayscale).\n",
    "\n",
    "Plot `val_y` and validate that it looks like the label image \"tmp/Datasets/cancer/cancer_data/cancer_data/outputs_test/0/pos_test_000072.png\" in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 435
    },
    "id": "znjvaLMM5UoS",
    "outputId": "4f855d75-0c29-49d4-ca95-db8ec8129114"
   },
   "outputs": [],
   "source": [
    "def plot_cancer(y):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Q9lL_gRp5UoS"
   },
   "source": [
    "The cancer is white which means that our cancerous label is 1 and our noncancerous label is 0.\n",
    "Interestingly, the dtype of `y` is a float, but we were expecting cancer labels (i.e. cancerous vs noncancerous).\n",
    "Use `torch.unique()` below to see all the unique values of `y`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "P3Pvpqlb5UoS",
    "outputId": "3da6f1a2-e75f-410f-e4cb-e2388e738b18"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Qf48qbkV5UoS"
   },
   "source": [
    "While 0 and 1 are in `y`, we have other values inbetween. For classification with cross entropy, we need `y` to be integers only. The original image was only black and white, so these non-integer values are due to the image being resized when it was loaded in. We need to think about the dataset and what we want to accomplish to determine how to treat these values: do we round, truncate, or ceiling?\n",
    "\n",
    "Because this is due to the resizing, these values only occur on the outer edge of the cancer label blobs. For this task, we want to identify regions with cancer, but we don't really need pixel accuracy - we just want the doctors to know to look there for cancer. So, we can truncate these non-one values by calling `y.long()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CEMY9kGC5UoT"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should also add a batch dimension to `val_x` and put it on the `device` so we can pass it through the network later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0JWlfiOj5UoT"
   },
   "source": [
    "Now retrieve some baseline accuracy information about the train and validation datasets.\n",
    "- How many images have cancerous cells in the dataset\n",
    "- On average how likely is a pixel going to be cancerous\n",
    "\n",
    "Print this information below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "CREFXUHo5UoT",
    "outputId": "b2fd2b47-dab6-4121-f57e-5bd500cfa908"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mfWMlPap5UoT"
   },
   "source": [
    "This dataset is very skewed towards non-cancerous data. This makes sense: most cells are non-cancerous, and images that contain cancer will have non-cancerous areas. For our model to be successful, we need to get higher than this baseline accuracy, so our validation accuracy should be over 90%."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ghxmQAotb9H1"
   },
   "source": [
    "## Network outputs\n",
    "\n",
    "Now that we understand our input images and our label images, what should our network's predictions look like?\n",
    "\n",
    "Because we are doing pixelwise classification, we want to output logits (unnormalized probabilities) for each pixel in our input. So the spatial dimension, i.e. HxW, of the network's output will be the same as those of our input.\n",
    "\n",
    "#### Loss\n",
    "If you recall from lab 1, when we implemented cross entropy, our predicted logits had a shape of `(B, K)`, where `K` was the number of classes, and our target classes had a shape of `(B,)`, which was used to index into our predicted logits. \n",
    "This will work similarly for pixelwise cross entropy, the predicted logits $\\hat y$ will have a shape of `(B, K, H, W)`, where `K=2` for our dataset, and the target $y$ will have a shape of `(B, H, W)`. \n",
    "You can use PyTorch's `F.cross_entropy()` function to compute the cross entropy loss. \n",
    "\n",
    "#### Accuracy\n",
    "Implement an accuracy function which compute the pixelwise accuracy between $\\hat y$ and $y$. \n",
    "\n",
    "*Note: Don't forget to use `torch.no_grad()` and since we are just interested in number, return the accuracy as a number, using `.item()`.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "E3fKUDffAg80"
   },
   "outputs": [],
   "source": [
    "def cancer_detection_accuracy(y_hat, y):\n",
    "    pass "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IApIBEZiAnYL"
   },
   "source": [
    "Now create a test function for `cancer_detection_accuracy`. Your function should:\n",
    "- Call `cancer_detection_accuracy()` twice, with two different predictions and accuracies.\n",
    "- Be simple, but nontrivial. Don't compare all (non)cancerous predictions with all (non)cancerous targets.\n",
    "\n",
    "Don't forget to add a batch dimension, even if it is just a batch size of 1.\n",
    "\n",
    "*Note: We recommend you create small images since they are easier to visualize and you have different values for each dimension, except for maybe H and W, because `B, C, H, W = 2, 2, 2, 2` can be harder to understand than `B, C, H, W = 1, 2, 3, 3`.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VL5O8zDOAn9e"
   },
   "outputs": [],
   "source": [
    "def test_cancer_detection_accuracy():\n",
    "    pass\n",
    "\n",
    "test_cancer_detection_accuracy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vJ0e84yN5UoT"
   },
   "source": [
    "## Dataloaders\n",
    "\n",
    "Now that we have the datasets, put them into dataloaders for training. Define a batch size (should be pretty small because we'll use a huge network). You should always shuffle your training dataset, but this time it is even more important because the dataset has all negative examples first. If you don't shuffle, your network will learn to predict all pixels as non-cancerous."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SmqtqXl5BdOs"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pqMjbQX9-ND9"
   },
   "source": [
    "---\n",
    "\n",
    "# Part 2: Training and Validation\n",
    "\n",
    "Now that we understand the data, we'll create training and validation functions BEFORE creating our network.\n",
    "\n",
    "We'll start with validation, since its simpler. Create a validation function that takes in a network and validation dataloader and returns the mean cross-entropy loss (just the value, not the tensor) and pixel-wise accuracy over the dataset. Don't forget to use `torch.no_grad()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dJIUcQxE_I7J"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "npvXwNkjBHEx"
   },
   "source": [
    "Because this is a very visual task, we also want to see how our model's predictions change over time. We'll use a single image: `val_dataset[172]` and store our model's prediction on it each time we perform validation. At the end, we can print out the images to see how our model improves over time. Create a function that takes in an image tensor of shape `(1, C, H, W)` (batch size of 1) and returns a prediction image of shape `(H, W)`, i.e remove the batch dimension and argmax the channel dimension to get the pixel predictions. Don't forget to use `torch.no_grad()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_hOw61JBBIV1"
   },
   "outputs": [],
   "source": [
    "def get_prediction(net, x):\n",
    "    pass "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nwGOr8vYFtdR"
   },
   "source": [
    "Create a function(s) to plot your training/validation losses on the same plot and your training/validation accuracies on the same plot. \n",
    "Label your axes and add a legend to make your plots legible. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WIOlpIwbF5W5"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create another function to plot 5 image predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0h8o9ncjDSv_"
   },
   "source": [
    "---\n",
    "\n",
    "Now create your training function. This should be similar to what we've done in lab 3, but remember that your data and labels are different sizes from lab 3's dataset.\n",
    "\n",
    "Store your training accuracy and losses every `log_interval` steps and validate every `val_interval` steps. Also store 5 predictions of `val_x` dispersed evenly throughout training. \n",
    "Return the statistics you've gathered (losses, accuracies, and val image predictions)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "d89dRpiADSAC"
   },
   "outputs": [],
   "source": [
    "def train(net, optimizer, train_dataloader, val_dataloader, n_minibatch_steps, log_interval, val_interval, val_im_delay, val_x):\n",
    "    # Train net on the training data, keeping track of loss, accuracy, and predictions on val_x\n",
    "    # Return losses and accuracies on both training and validation datasets, and 5 images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uFSbxQqs_IQs"
   },
   "source": [
    "---\n",
    "\n",
    "# Part 3: Training Smallest-Possible Network\n",
    "\n",
    "While we want to use the UNet on this problem, we should make sure our training and validation runs with a smaller network. Doing this first makes sure that we can focus on bugs in the training and validation loop before focusing on debugging the network itself. And this makes sure that we understand what our inputs and outputs should be.\n",
    "\n",
    "Before reading further, answer the following question:\n",
    "\n",
    "What is the smallest-possible convolutional network that can solve this task?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tSHaDax8EnaO"
   },
   "outputs": [],
   "source": [
    "# CREATE SMALLEST NETWORK\n",
    "small_net = nn.Conv2d(3, 2, kernel_size=1).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hZGa-ervEvZR"
   },
   "source": [
    "---\n",
    "\n",
    "You should have created a one-layer network with 3 in-channels and 2 out-channels, with a 1x1 convolution. This network obviously won't learn well, but it will be useful for quickly testing and debugging our training and validation functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Z6QNhv0yEu0u",
    "outputId": "03df0ae5-52b9-4a9a-a805-d2879f45d070"
   },
   "outputs": [],
   "source": [
    "# Create an optimizer for your network\n",
    "\n",
    "# Train for a few training steps: enough to gather some stats\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "dq8ro4LJFXHJ",
    "outputId": "6b71998c-c821-4d9d-c33e-81e45649d116"
   },
   "outputs": [],
   "source": [
    "# Create your loss and accuracy graphs, and print out the 5 validation images:\n",
    "# We do not expect this network to have learned well.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_ZfLz8iJI4XX"
   },
   "outputs": [],
   "source": [
    "# We don't need this network anymore, so we'll delete it to free up GPU space\n",
    "del small_net, small_optimizer, train_losses, train_accs, val_losses, val_accs, val_images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5nLBjLkM5UoU"
   },
   "source": [
    "---\n",
    "\n",
    "# Part 4: Implementing the UNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PoPIQ5qZ5UoU"
   },
   "source": [
    "Use the “Deep Convolution U-Net” from this paper: [(U-Net: Convolutional Networks for Biomedical Image Segmentation)](https://arxiv.org/pdf/1505.04597.pdf)\n",
    "\n",
    "![(Figure 1)](https://lh3.googleusercontent.com/qnHiB3B2KRxC3NjiSDtY08_DgDGTDsHcO6PP53oNRuct-p2QXCR-gyLkDveO850F2tTAhIOPC5Ha06NP9xq1JPsVAHlQ5UXA5V-9zkUrJHGhP_MNHFoRGnjBz1vn1p8P2rMWhlAb6HQ=w2400)\n",
    "<!-- <img src=\"https://lh3.googleusercontent.com/qnHiB3B2KRxC3NjiSDtY08_DgDGTDsHcO6PP53oNRuct-p2QXCR-gyLkDveO850F2tTAhIOPC5Ha06NP9xq1JPsVAHlQ5UXA5V-9zkUrJHGhP_MNHFoRGnjBz1vn1p8P2rMWhlAb6HQ=w2400\" width=\"1600\" height=\"1200\"/> -->\n",
    "\n",
    "You will implement the UNet with some simplifying adjustments.\n",
    "The original UNet does not pad their convolutions, so their output is smaller than their input. \n",
    "Your input and output are the same size so you should pad your convolutions. \n",
    "The only time height and width should change are from \"max pool\" and the \"up-conv\". \n",
    "This means you no longer need to \"crop\" the cross connections, since both sides of the U should have the same spatial dimensions. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yJhqQGTB5UoU"
   },
   "source": [
    "This is a large complex network structure and due to the cross connections, you can't just stick a bunch of layers into `nn.Sequential`. \n",
    "This network does have a clear repeating pattern of blocks that move down (left side of U) and blocks that move up (right side of U). \n",
    "Just like how we use functions to made our code more modular, we'll create two modules to capture that repeating structure to simplify our final network.\n",
    "\n",
    "Implement `DownBlock` (left side of U): The `DownBlock` incorporates 2 sequential convolution layers (with nonlinearities) followed by a downsampling (`nn.MaxPool2d`). \n",
    "It should return the downsampled image and the cross connection output.\n",
    "You should NOT use `copy` or `clone` on the cross connection output to do the \"copy and crop\" operation; just return the same tensor you passed into `MaxPool2D` layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hfX_D3uZ5UoY"
   },
   "outputs": [],
   "source": [
    "# You are welcome (and encouraged) to use the built-in batch normalization, nn.BatchNorm2d, and dropout layers, nn.Dropout.\n",
    "\n",
    "class DownBlock(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super().__init__()\n",
    "        # TODO\n",
    "\n",
    "    def forward(self, x):\n",
    "        # TODO\n",
    "        return"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Implement `UpBlock` (right side of U): The `UpBlock` incorporates an upsampling (i.e. `nn.Conv2DTranspose`) followed by 2 sequential convolution layers (with nonlinearities). Its input should be a downsampled image, `x`, and the cross connection output, `cross_output`, from a `DownBlock`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class UpBlock(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels) -> None:\n",
    "        super().__init__()\n",
    "        # TODO\n",
    "\n",
    "    def forward(self, x, cross_output):\n",
    "        # TODO\n",
    "        return"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yWYrQI-k5UoZ"
   },
   "source": [
    "Test below that your `DownBlock` and `UpBlock` work correctly.\n",
    "Create a function `test_updown_block()` which instantiates a `DownBlock` and `UpBlock` and passes in a dummy image `x = torch.zeros((2, 4, 16, 16))`.\n",
    "The downsampled output of your `DownBlock` should be passed through the provided `mid_conv` layer (representing the bottom layer in the UNet) before being passed to the `DownBlock`. \n",
    "You will assert that the the cross-connection output of your `DownBlock` and the output of your `UpBlock` have the shape of `(2, 8, 16, 16)`.\n",
    "Remember to use `torch.no_grad()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-86Ovwb95UoZ"
   },
   "outputs": [],
   "source": [
    "def test_updown_block():\n",
    "    mid_conv = nn.Sequential(nn.Conv2d(8, 16, 3, 1), nn.ReLU(), nn.Conv2d(16, 16, 3, 1), nn.ReLU())\n",
    "    pass\n",
    "\n",
    "test_updown_block()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Implement the `UNet` using your `UpBlock` and `DownBlock` modules.\n",
    "You will additionally need to add layers for the bottom of the U and the final output layer.\n",
    "You do not need to use the same number of channels as the UNet paper. It is recommended you use fewer channels for faster training and less memory consumption. `[8, 16, 32, 64]` are a good choice for channel sizes.\n",
    "\n",
    "*Note: If you want to store `UpBlock`s and `DownBlock`s in a list then use the `nn.ModuleList()` which is a module that operates like a list, but exposes the modules to PyTorch.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jq22IyKanxo_"
   },
   "outputs": [],
   "source": [
    "class UNet(nn.Module):\n",
    "  def __init__(self, channels=(8, 16, 32, 64)) -> None:\n",
    "    super().__init__()\n",
    "    # Initialize all modules for the full UNet\n",
    "    # Use the `channels` parameter to specify the number of channels on each level of the U\n",
    "\n",
    "  def forward(self, x):\n",
    "    # TODO\n",
    "    return"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test `UNet` by making sure that it takes a batch of input images `x` and returns an appropriately-shaped output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "laudR91z5UoZ"
   },
   "outputs": [],
   "source": [
    "def test_UNet_shapes():\n",
    "    pass \n",
    "    \n",
    "test_UNet_shapes()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MhuXf9ky5UoZ"
   },
   "source": [
    "---\n",
    "\n",
    "# Training and Testing UNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train `UNet` below and get a better than baseline accuracy. For Colab/Kaggle users, we recommend training with a small `UNet` first on the CPU and then scaling up the `UNet` and training on the GPU. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "RkieTbwlYWPS",
    "outputId": "1ead0f53-927e-42ae-a43a-14bf7ae2f3eb"
   },
   "outputs": [],
   "source": [
    "# Here are some example hyperparameters: feel free to play around with them for better results\n",
    "n_minibatch_steps = 1000 # number of optimization steps\n",
    "train_log_interval = 10 # how many steps between saving training loss and accuracy\n",
    "val_interval = 200 # how many steps between performing the validation loop\n",
    "val_im_delay = 50 # how many steps before we start storing validation image predictions\n",
    "\n",
    "# TODO: create UNet and train it, storing statistics from training/validating"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot your training/validation losses and accuracies. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 927
    },
    "id": "mTg1jyIsYVZN",
    "outputId": "52694f46-f73c-4d56-cf27-8eba0494392a"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "S4s92S2_jQOG"
   },
   "source": [
    "Visualize the 5 predictions of `val_x`. Your images should start to look like the `plot_cancer()` image you made above (signs of cancer in the top left corner). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "XXfG3wClh8an",
    "outputId": "17591b6c-1268-42d6-d074-d5403d1d790a",
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "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.10.13"
  },
  "nteract": {
   "version": "0.28.0"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}