part3.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "dlim_part3.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7DtEVhuIt30V"
      },
      "source": [
        "# **Deep Learning with PyTorch** -  Introductory Lab\n",
        "\n",
        "## **Part 3:** Deep Learning with PyTorch - Image Denoising Example\n",
        "\n",
        "* In this part we do a simple deep learning example.\n",
        "* The aim it to train a Convolutional Neural Network (CNN) for image denoising.\n",
        "\n",
        "* This notebook is designed to be run on Colab.\n",
        "* But it can easily be modified to be run locally.\n",
        "\n",
        "### Setup\n",
        "* **IMPORTANT!** Activate GPU acceleration by clicking: `Runtime -> Change runtime type -> Hardware Accelerator -> GPU`.\n",
        "* Unless you already did, you need to make shortcut to the shared data folder in your google drive (see instructions in **part 2** of the lab).\n",
        "* Run the code in the cell below and follow the instructions to mount your drive to Colab."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RMn4vmsStzfg"
      },
      "source": [
        "# Connect your google drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wY48y3DPvWwW"
      },
      "source": [
        "# Set up path to your dataset\n",
        "# Use this path if you have copied the data to your Google drive\n",
        "# If you put the data somewhere else, you need to modify this\n",
        "dataset_path = '/content/gdrive/My Drive/DIV2K_train_small/'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "idAY2B4dvXfd"
      },
      "source": [
        "# Import libraries we will need\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from PIL import Image"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VS4vWDgXvM5T"
      },
      "source": [
        "# Functions we used in part 2\n",
        "\n",
        "# Read an image with the given name and convert it to torch\n",
        "def imread(image_file, base_path=None):\n",
        "    if base_path is None:\n",
        "        base_path = dataset_path\n",
        "    im_pil = Image.open(base_path + image_file)\n",
        "    im_np = np.array(im_pil, copy=False)\n",
        "    im_torch = torch.from_numpy(im_np).permute(2, 0, 1)\n",
        "    return im_torch.float()/255\n",
        "\n",
        "# Show a PyTorch image tensor\n",
        "def imshow(im, normalize=False):\n",
        "    # Detatch from graph and move to CPU\n",
        "    im = im.detach().cpu()\n",
        "\n",
        "    # Fit the image to the [0, 1] range if normalize is True\n",
        "    if normalize:\n",
        "        im = (im - im.min()) / (im.max() - im.min())\n",
        "\n",
        "    # Remove redundant dimensions \n",
        "    im = im.squeeze()    # Mini excersize: check in the documentation what this function does\n",
        "\n",
        "    is_color = (im.dim() == 3)\n",
        "\n",
        "    # If there is a color channel dimension, move it to the end\n",
        "    if is_color:\n",
        "        im = im.permute(1, 2, 0)\n",
        "\n",
        "    im_np = im.numpy().clip(0,1)    # Convert to numpy and ensure the values in the range [0, 1]\n",
        "    if is_color:\n",
        "        plt.imshow(im_np)\n",
        "    else:\n",
        "        plt.imshow(im_np, cmap='gray')\n",
        "    plt.axis('off')\n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pfDn2cPjwwuv"
      },
      "source": [
        "### Defining a Neural Network\n",
        "* Lets first create a CNN that just consists of a single conv layer with no nonlinearities.\n",
        "* When reading the code, also check the [doc for the Conv2d layer](https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html#conv2d)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uMVxDeY2v_fy"
      },
      "source": [
        "# In PyTorch, a neural network is represented by a python class.\n",
        "# All neural networks and network modules must inherit from the nn.Module class.\n",
        "\n",
        "class LinearConvNet(nn.Module):\n",
        "    \"\"\" A CNN consisting of a single conv layer.\n",
        "    args:\n",
        "        kernel_size: Size of the filter. Must be odd.\n",
        "        bias: Use bias or not.\n",
        "        init_std: Initialize weights randomly with this standard deviation.\"\"\"\n",
        "\n",
        "    # In the constructor we can define the modules and layers that we want the network to have.\n",
        "    # We should also initialize the weights of the layers.\n",
        "\n",
        "    def __init__(self, kernel_size=5, num_channels=1):\n",
        "        super().__init__()   # Always call the constructor of the parent\n",
        "\n",
        "        # Define a convolutional layer\n",
        "        # We talked about in/out channels and the kernel size in the previous part of the lab\n",
        "        # The padding parameter controls the ammount of zeros to implicitly add around the image border.\n",
        "        # Setting padding=kernel_size//2 will give use the same output size as input size, if kernel_size is odd.\n",
        "        self.conv = nn.Conv2d(in_channels=num_channels, \n",
        "                              out_channels=num_channels, \n",
        "                              kernel_size=kernel_size,\n",
        "                              padding=kernel_size//2)\n",
        "\n",
        "        # torch.nn.init implements many different initialization strategies.\n",
        "        # dirac_() simply initializes the filter to the identity mapping.\n",
        "        torch.nn.init.dirac_(self.conv.weight)\n",
        "\n",
        "\n",
        "    # In the forward function, you implement what the network module does.\n",
        "    # In this case, we simply apply the convolutional layer and return its output.\n",
        "\n",
        "    def forward(self, im):\n",
        "        return self.conv(im)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_u2P2nJJ2MVq"
      },
      "source": [
        "* In the forward function above, note that `self.conv(im)` is essentially equivalent to the functional form we used in the previous part: `F.conv2d(im, self.conv.weight)`\n",
        "* The `self.conv` object internally stores the convolution weights we want to learn, along with the other settings of the convolution (such as the padding).\n",
        "* We can therefore apply it to an input image with the simple call `self.conv(im)`.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MFe7b5LH3Yv4"
      },
      "source": [
        "### Applying a Neural Network\n",
        "* Lets apply the network to an image.\n",
        "* First we need to create an instance of the network."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hvBSNVdz3Gqo"
      },
      "source": [
        "# Create a network object\n",
        "net = LinearConvNet(kernel_size=9, num_channels=1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7AQRZiyp33he"
      },
      "source": [
        "# Load an example image as in the previous part\n",
        "im = imread('0705x4.png')\n",
        "im_gray = im.mean(dim=0)\n",
        "imshow(im_gray)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-Pu87hpV4RAm"
      },
      "source": [
        "# Run the network on the image\n",
        "# This will call the forward function of the network automatically\n",
        "# As before, we first need to resize the image to be 4D\n",
        "\n",
        "im_out = net(im_gray.view(1, 1, *im_gray.shape))\n",
        "\n",
        "imshow(im_out)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OdGJ3LmIYinH"
      },
      "source": [
        "* Note that the output of the network is identical to the input.\n",
        "* That is because we initialize the only convolutional layer to be identity (a dirac).\n",
        "* We can manually change the weights to an averaging filter and see the effect ..."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "68L67i_2X26w"
      },
      "source": [
        "# change the weights to an averaging filter\n",
        "net.conv.weight[:] = 1 / net.conv.weight.nelement()\n",
        "\n",
        "im_out = net(im_gray.view(1, 1, *im_gray.shape))\n",
        "\n",
        "imshow(im_out)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CtcIsfm-Zuk9"
      },
      "source": [
        "### Towards Deep Image Denoising\n",
        "* The goal in the next steps is to train the network to do image denoising.\n",
        "\n",
        "#### Generating Noise\n",
        "* To start with, we need a function that can generate noisy images from clean ones."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u2cjZTKrZi6m"
      },
      "source": [
        "# Lets make a simple function that adds noise.\n",
        "def add_noise(im, std=0.1):\n",
        "    return im + std * torch.randn(im.shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4mkn4a_qb_zT"
      },
      "source": [
        "# Lets test it\n",
        "im_noisy = add_noise(im)\n",
        "imshow(im_noisy)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0K84yIlscknO"
      },
      "source": [
        "#### Objective function\n",
        "* To train the network, we need an objective function (loss) that measures the difference between the network output and the clean reference image.\n",
        "* Lets use the Mean Squared Error (i.e. L2 loss).\n",
        "* You can check the [doc here](https://pytorch.org/docs/stable/generated/torch.nn.MSELoss.html?highlight=mse#torch.nn.MSELoss)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Yz10fJgycNjb"
      },
      "source": [
        "# The objectively is also fundamentally a neural network module\n",
        "objective = nn.MSELoss()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O3Qjw3CLdumI"
      },
      "source": [
        "# Lets test our objective\n",
        "# First we create a network designed for taking RGB images as input\n",
        "\n",
        "net = LinearConvNet(kernel_size=9, num_channels=3)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pdtQyx_ReOSd"
      },
      "source": [
        "# Now, lets compute the loss between the network output and the clean image\n",
        "\n",
        "# Make image 4d\n",
        "im_reference = im.view(1, *im.shape)\n",
        "\n",
        "# Add noise for the input\n",
        "im_input = add_noise(im_reference)\n",
        "\n",
        "# Apply network\n",
        "im_output = net(im_input)\n",
        "\n",
        "# Compute loss\n",
        "loss = objective(im_output, im_reference)\n",
        "\n",
        "print(loss)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AjMXGhIpffh1"
      },
      "source": [
        "* The value represents the Mean Squared Error over all pixels\n",
        "\n",
        "#### Computing Gradients\n",
        "* To train the network with Stochastic Gradient Descent (see lecture), we first need to compute the gradients of the objective function w.r.t. the weights we want to optimize.\n",
        "* Fortunately, this is extremely easy in PyTorch.\n",
        "* We just need to call the `backward()` function on the computed loss."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dFe2tbP4fVQx"
      },
      "source": [
        "# Compute gradients\n",
        "loss.backward()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JyisP0d2fZG8"
      },
      "source": [
        "# The gradients are stored in the .grad attribute of all weights in the network\n",
        "# Lets print some info:\n",
        "print('grad shape:', net.conv.weight.grad.shape)\n",
        "print('grad mean:', net.conv.weight.grad.mean())\n",
        "print('grad min:', net.conv.weight.grad.min())\n",
        "print('grad max:', net.conv.weight.grad.max())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5XbkEXoAkg81"
      },
      "source": [
        "* To perform a gradient descent step, we could essentially do the following\n",
        "```\n",
        "for w in net.parameters():\n",
        "      w += learning_rate * w.grad\n",
        "```\n",
        "* Here learning_rate is the size of the gradient step.\n",
        "* However, PyTorch already have convenient oprimizers implemented.\n",
        "* Lets check how it works next.\n",
        "\n",
        "#### The Optimizer\n",
        "* We first define an optimizer.\n",
        "* Lets choose the standard SGD ([see doc](https://pytorch.org/docs/stable/optim.html?highlight=sgd#torch.optim.SGD))."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9vinoDkIoQSL"
      },
      "source": [
        "# First define the network again\n",
        "net = LinearConvNet(kernel_size=9, num_channels=3)\n",
        "\n",
        "# To the optimizer, we specify the parameters we whish to optimize and the learning rate (lr)\n",
        "# Here, we also set the momentum parameter, which also averages gradients over time\n",
        "optimizer = torch.optim.SGD(net.parameters(), lr=0.01, momentum=0.9)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s2D6xDE7go-o"
      },
      "source": [
        "# Now, lets create a loop that optimizes the network on a single image\n",
        "\n",
        "im = imread('0705x4.png')\n",
        "\n",
        "num_steps = 2000\n",
        "\n",
        "for n in range(num_steps):\n",
        "    # Make image 4d\n",
        "    im_reference = im.view(1, *im.shape)\n",
        "\n",
        "    # Add noise for the input\n",
        "    im_input = add_noise(im_reference)\n",
        "\n",
        "    # We need to clear the old gradients first\n",
        "    optimizer.zero_grad()\n",
        "\n",
        "    # Apply network\n",
        "    im_output = net(im_input)\n",
        "\n",
        "    # Compute loss\n",
        "    loss = objective(im_output, im_reference)\n",
        "\n",
        "    # Compute the new gradients\n",
        "    loss.backward()\n",
        "\n",
        "    # Take an optimization step\n",
        "    optimizer.step()\n",
        "\n",
        "    # Print loss to show progress\n",
        "    if n % 50 == 0:\n",
        "         print('Iter {}/{}:  loss = {}'.format(n, num_steps, loss.item()))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_1YQQL82pOe8"
      },
      "source": [
        "# Compare the last input and output of the network\n",
        "imshow(im_input[...,:100,100:200])\n",
        "imshow(im_output[...,:100,100:200])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vxi912CzqPmD"
      },
      "source": [
        "* The network has learned to donoise the image to some extent!\n",
        "* We saw that the loss was decreasing quite rapidly over the iteration\n",
        "### 💡 **Exercise**\n",
        "* Apply the trained network on another image and check the results visually. What do you think?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U2ZlPuuGp4s6"
      },
      "source": [
        "# Implement your solution here\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BDx0zAFAsFdX"
      },
      "source": [
        "#### Creating a dataset loader\n",
        "\n",
        "* Note that we only trained the network above on a single image.\n",
        "* In general, this would not work well in practice, especially if using deeper networks.\n",
        "* The network would just overfit to that image and not generalize to other images.\n",
        "* Here, we will implement training on a full dataset using standard PyTorch tools.\n",
        "* First we will copy the data over to the colab instance."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "axsL-Qiqowax"
      },
      "source": [
        "# Make a directory called data\n",
        "!mkdir /content/data "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zSWwMMApovKA"
      },
      "source": [
        "# Copy the images for training and testing\n",
        "# This takes a while ... read and understand the next cells in the meantime!\n",
        "!rsync -avzh /content/gdrive/My\\ Drive/DIV2K_train_small/* /content/data/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tiJKIvT6qNLK"
      },
      "source": [
        "from pathlib import Path\n",
        "from torchvision import transforms     # Torchvision is a part of PyTorch that contains extra vision functionality\n",
        "\n",
        "class DenoisingDataset(torch.utils.data.Dataset):\n",
        "    def __init__(self, dir, noise_std=0.1, crop_size=200, pattern='*.png', start_ind=None, end_ind=None):\n",
        "        super().__init__()\n",
        "\n",
        "        self.noise_std = noise_std\n",
        "        self.crop_size = crop_size\n",
        "\n",
        "        # List all images in the folder\n",
        "        image_dir = Path(dir)\n",
        "        self.image_files = list(sorted(image_dir.glob(pattern)))\n",
        "        self.image_files = self.image_files[start_ind:end_ind]\n",
        "\n",
        "        # Create transform that crops the image and converts to tensor\n",
        "        self.transform = transforms.Compose([transforms.RandomCrop(size=self.crop_size, pad_if_needed=True),\n",
        "                                             transforms.ToTensor()])\n",
        "\n",
        "\n",
        "    # The dataset need to implement the getitem function. This should load and return an image pair.\n",
        "    # This function is called with a random index during training.\n",
        "    def __getitem__(self, index):\n",
        "        # Get image path\n",
        "        im_path = self.image_files[index]\n",
        "\n",
        "        # Load the PIL image\n",
        "        im_pil = Image.open(im_path)\n",
        "        \n",
        "        # Transform to tensor. This also do normalization\n",
        "        im_reference = self.transform(im_pil)\n",
        "\n",
        "        # Add noise\n",
        "        im_noisy = add_noise(im_reference, std=self.noise_std)\n",
        "\n",
        "        return im_reference, im_noisy\n",
        "\n",
        "    # The length function should return the number of images in the dataset\n",
        "    def __len__(self):\n",
        "        return len(self.image_files)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "04SpC3oexLkD"
      },
      "source": [
        "# Lets test the dataset loader.\n",
        "dataset_train = DenoisingDataset(dir='/content/data/')\n",
        "\n",
        "im_ref, im_input = dataset_train[4]\n",
        "\n",
        "imshow(im_ref)\n",
        "imshow(im_input)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TlCswdDtz3Aq"
      },
      "source": [
        "# Now we need to create a data loader. This is easy\n",
        "# batch_size is the number of images we sample in parallell.\n",
        "# The loader will automatically sample images from the dataset and concatenate them\n",
        "\n",
        "train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=32, shuffle=True, drop_last=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p2N_bCz1Cpet"
      },
      "source": [
        "# Lets test it!\n",
        "for im_ref, im_input in train_loader:\n",
        "    print('ref shape:', im_ref.shape)\n",
        "    print('input shape:', im_input.shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U2BumZDKqXOw"
      },
      "source": [
        "* Note that the tensors returned by the loader consists of 32 stacked color images of size 200x200.\n",
        "* 32 is the batch_size we chose in the previous cell.\n",
        "* The loop above goes through all images in the dataset once.\n",
        "\n",
        "#### The training loop\n",
        "\n",
        "* We saw earlier a simple example of a training loop for a single image.\n",
        "* Now we will write the training loop for the real case.\n",
        "* This will implement real Stochastic Gradient Descent, since also the data is sampled as random batches during training.\n",
        "* The training is usually split into **epochs**.\n",
        "* Each **epoch** cycles through the dataset once. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GpTdaVa_CrER"
      },
      "source": [
        "# We first write a function that runs a single epoch\n",
        "\n",
        "def run_epoch(net, objective, optimizer, data_loader, is_training=True, cuda=False):\n",
        "    # If we are training on GPU, make sure that the network is there\n",
        "    if cuda:\n",
        "        net.cuda()\n",
        "    else:\n",
        "        net.cpu()\n",
        "\n",
        "    # Set network to train/val mode\n",
        "    net.train(is_training)\n",
        "\n",
        "    # Turn of tracking of gradients if we are doing validation (saves memory)\n",
        "    torch.set_grad_enabled(is_training)\n",
        "\n",
        "    # For tracking the average loss\n",
        "    avg_loss = 0\n",
        "    num_batches = 0\n",
        "\n",
        "    # Loop over the dataset\n",
        "    for im_reference, im_input in data_loader:\n",
        "        # Move data to GPU if we use it\n",
        "        if cuda:\n",
        "            im_reference = im_reference.cuda()\n",
        "            im_input = im_input.cuda()\n",
        "        \n",
        "        # Reset gradients\n",
        "        if is_training:\n",
        "            optimizer.zero_grad()\n",
        "        \n",
        "        # Apply network\n",
        "        im_output = net(im_input)\n",
        "\n",
        "        # Compute loss\n",
        "        loss = objective(im_output, im_reference)\n",
        "\n",
        "        if is_training:\n",
        "            # Compute gradients\n",
        "            loss.backward()\n",
        "\n",
        "            # Update weights by taking one optimizer step\n",
        "            optimizer.step()\n",
        "        \n",
        "        # Update statistics\n",
        "        avg_loss += loss.item()\n",
        "        num_batches += 1\n",
        "\n",
        "    avg_loss /= num_batches\n",
        "\n",
        "    # Return average loss so that we can print that\n",
        "    return avg_loss\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "57nbZ86ryZqR"
      },
      "source": [
        "* In each epoch, we should train the network on the training set.\n",
        "* We should also evaluate the network at each epoch.\n",
        "* This is very important in order to monitor the training and making sure that the network does not overfit on the training data."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jT7t9zCkGYm5"
      },
      "source": [
        "# Now we write the full training loop\n",
        "# It should do both training and validation\n",
        "\n",
        "def train_network(net, objective, optimizer, train_loader, val_loader, num_epochs, cuda=False):\n",
        "    # Loop over all epochs\n",
        "    for ep in range(num_epochs):\n",
        "        # Run the training\n",
        "        train_loss = run_epoch(net, objective, optimizer, train_loader, is_training=True, cuda=cuda)\n",
        "\n",
        "        # Run the validation\n",
        "        val_loss = run_epoch(net, objective, optimizer, val_loader, is_training=False, cuda=cuda)\n",
        "\n",
        "        # Print stats\n",
        "        print('[Ep {}/{}] Train loss: {:.06f}  ,  Val loss: {:.06f}'.format(ep+1, num_epochs, train_loss, val_loss))\n",
        "\n",
        "    print('Training done!')\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RHOMwgfyy1TL"
      },
      "source": [
        "#### Running a first training\n",
        "* Now we have all the tools to run our first real training.\n",
        "* Lets set it up!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X0i7b_0eHaMW"
      },
      "source": [
        "# Lets first specify some general settings\n",
        "cuda = True         # Run with GPU on since that will be faster.\n",
        "batch_size = 32\n",
        "\n",
        "# Define the training dataset. The final 50 images (end_ind=-50) we can save them for validation\n",
        "dataset_train = DenoisingDataset(dir='/content/data/', end_ind=-50)\n",
        "\n",
        "# Similarily we define the validation dataset\n",
        "dataset_val = DenoisingDataset(dir='/content/data/', start_ind=-50)\n",
        "\n",
        "# Create the dataset loaders\n",
        "train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size, shuffle=True, drop_last=True)\n",
        "val_loader = torch.utils.data.DataLoader(dataset_val, batch_size=batch_size, shuffle=True, drop_last=True)\n",
        "\n",
        "# Create the network \n",
        "net_linear = LinearConvNet(kernel_size=9, num_channels=3)\n",
        "\n",
        "# Lets use the same objective\n",
        "objective = nn.MSELoss()\n",
        "\n",
        "# Also the same SGD optimizer as before\n",
        "optimizer = torch.optim.SGD(net_linear.parameters(), lr=0.05, momentum=0.9)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uYNCK_uZHx28"
      },
      "source": [
        "# Lets start the training!\n",
        "train_network(net_linear, objective, optimizer, train_loader, val_loader, num_epochs=50, cuda=cuda)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5m9IfkSCi-nr"
      },
      "source": [
        "* Now we have trained our first network. Lets test it on an image of the validation set."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b3NvU3CVH6cO"
      },
      "source": [
        "# We should move the network back to CPU\n",
        "net_linear.cpu()\n",
        "net_linear.eval()\n",
        "\n",
        "# Read an image from the validation set\n",
        "im_ref = imread(str(dataset_val.image_files[0]), base_path='')\n",
        "\n",
        "im_noisy = add_noise(im_ref)\n",
        "\n",
        "im_denoised = net_linear(im_noisy.unsqueeze(0))\n",
        "\n",
        "imshow(im_noisy[..., 100:200, :100])\n",
        "imshow(im_denoised[..., 100:200, :100])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YGvFxAU7ustp"
      },
      "source": [
        "### Creating a deep network\n",
        "* The network we have trained so far consists of just a single convolution.\n",
        "* Now we will build our first deep network.\n",
        "* The code below implements the CNN architecture in this figure.\n",
        "* Note that the second convolution downsamples the feature map by a factor 2 and that the third one up-samples it again.\n",
        "![could not be displayed](https://raw.githubusercontent.com/martin-danelljan/pytorch-tutorial/master/images/miniunet.png)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "poB5xHWcIWqi"
      },
      "source": [
        "\n",
        "class MiniUNet(nn.Module):\n",
        "    \"\"\"A small U-Net architecture with skip connections.\"\"\"\n",
        "\n",
        "    def __init__(self, bias=True, init_std=1e-2):\n",
        "        super().__init__()\n",
        "\n",
        "        # Initial conv that preserves the resolution\n",
        "        self.conv_init = nn.Conv2d(3, 16, kernel_size=7, stride=1, padding=3, bias=bias)\n",
        "\n",
        "        # A conv layer that also downsamples the data (stride=2)\n",
        "        self.conv1 = nn.Conv2d(16, 32, kernel_size=5, stride=2, padding=2, bias=bias)\n",
        "\n",
        "        # One conv layer that upsample the data again.\n",
        "        # These type of layers are called transposed convolutions.\n",
        "        self.convtp1 = nn.ConvTranspose2d(32, 16, kernel_size=6, stride=2, padding=2, bias=bias)\n",
        "\n",
        "        # Add a final conv layer that generates the output\n",
        "        self.conv_final = nn.Conv2d(16, 3, kernel_size=3, stride=1, padding=1, bias=bias)\n",
        "\n",
        "        # We use the rectified linear unit activation\n",
        "        self.activation = nn.ReLU(inplace=True)\n",
        "\n",
        "        # Initialization\n",
        "        for m in self.modules():\n",
        "            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):\n",
        "                # We randomly initialize the weights with some standard deviation\n",
        "                torch.nn.init.normal_(m.weight, std=init_std)\n",
        "                if m.bias is not None:\n",
        "                    torch.nn.init.zeros_(m.bias)\n",
        "\n",
        "\n",
        "    def forward(self, im):\n",
        "        # Layer 1: (same resolution)\n",
        "        x1 = self.activation(self.conv_init(im))\n",
        "\n",
        "        # Layer 2: (2x down sampling)\n",
        "        x2 = self.activation(self.conv1(x1))\n",
        "\n",
        "        # Layer 3: (2x up sampling)\n",
        "        x3 = self.activation(self.convtp1(x2))\n",
        "\n",
        "        # Layer 4: (same as initial resolution)\n",
        "        im_out = self.conv_final(x3)\n",
        "\n",
        "        return im_out"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VPQWY3bt2zNV"
      },
      "source": [
        "# Lets first specify some general settings\n",
        "cuda = True         # Run with GPU on since that will be faster.\n",
        "batch_size = 32\n",
        "\n",
        "# Define the training dataset. The final 50 images (end_ind=-50) we can save them for validation\n",
        "dataset_train = DenoisingDataset(dir='/content/data/', end_ind=-50)\n",
        "\n",
        "# Similarily we define the validation dataset\n",
        "dataset_val = DenoisingDataset(dir='/content/data/', start_ind=-50)\n",
        "\n",
        "# Create the dataset loaders\n",
        "train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size, shuffle=True, drop_last=True)\n",
        "val_loader = torch.utils.data.DataLoader(dataset_val, batch_size=batch_size, shuffle=True, drop_last=True)\n",
        "\n",
        "# Create the network \n",
        "miniunet = MiniUNet()\n",
        "\n",
        "# Lets use the same objective\n",
        "objective = nn.MSELoss()\n",
        "\n",
        "# Also the same SGD optimizer as before\n",
        "optimizer = torch.optim.SGD(miniunet.parameters(), lr=0.05, momentum=0.9)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ETvY51uX3BWs"
      },
      "source": [
        "# Lets start the training!\n",
        "train_network(miniunet, objective, optimizer, train_loader, val_loader, num_epochs=50, cuda=cuda)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4cZZCYtFPCGU"
      },
      "source": [
        "* We see that the network starts with a very high loss compared to that of the linear network.\n",
        "* Moreover, it converges very slowly.\n",
        "* Lets try to address that.\n",
        "### 💡 **Exercise**\n",
        "* Add residual connections (also known as skip connections) to the MiniUNet architecture according to the following figure.\n",
        "* \"+\" denotes standard pointwise addition.\n",
        "* Note that the long residual connection adds the noisy input image to the output.\n",
        "* This means that the network will predict the noise, which is then substracted from the input image.\n",
        "![could not be displayed](https://raw.githubusercontent.com/martin-danelljan/pytorch-tutorial/master/images/miniunet_residual.png)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9vcMvnQTQ1e7"
      },
      "source": [
        "# You can modify the code above or make a new copy of it in another cell\n",
        "# Then train you new network architecture here\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ga9IBleDTEAp"
      },
      "source": [
        "### 💡 **Exercise**\n",
        "* Although the loss starts at a reasonable value, the convergence is still slow.\n",
        "* In CNNs, this is often addressed using Batch Normalization layers.\n",
        "* Thease layers often makes deep networks significantly easier and faster to train.\n",
        "* Check the [PyTorch documentation for the BatchNorm2d layer](https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm2d.html?highlight=batchnorm#torch.nn.BatchNorm2d).\n",
        "* Then add Batch Normalization (BN) layers according to the figure below.\n",
        "* **Tip:** You can keep the optinal parameters to their default values.\n",
        "![could not be displayed](https://raw.githubusercontent.com/martin-danelljan/pytorch-tutorial/master/images/miniunet_bn.png)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MxO3O2MvUr-j"
      },
      "source": [
        "# Train your new architecture here\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZHJb_0BKU6ud"
      },
      "source": [
        "### 💡 **Exercise**\n",
        "* Visually check the denoising quality of your network using the code below.\n",
        "* Check for different images and different regions in the image.\n",
        "* What do you think about the quality? How does it compare with the linear network? "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FfGSPpGI3X2r"
      },
      "source": [
        "# We should move the network back to CPU\n",
        "miniunet.cpu()\n",
        "\n",
        "# The network should be in eval() mode. This is very important when using Batch Norm.\n",
        "miniunet.eval()\n",
        "\n",
        "# Read an image from the validation set\n",
        "im_ref = imread(str(dataset_val.image_files[0]), base_path='')\n",
        "\n",
        "# We need to select a crop that is divisible with 4\n",
        "im_ref = im_ref[:, :256, :256]\n",
        "\n",
        "im_noisy = add_noise(im_ref)\n",
        "\n",
        "im_denoised = miniunet(im_noisy.unsqueeze(0))\n",
        "\n",
        "imshow(im_noisy[..., 100:200, :100])\n",
        "imshow(im_denoised[..., 100:200, :100])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sgNhZVTbmlnB"
      },
      "source": [
        "Here are some extra exercises you can do if you get time. or after the lab if you are interested.\n",
        "\n",
        "### 💡 **Exercise (extra)**\n",
        "\n",
        "* Try applying the network to other images with other noise standard deviations.\n",
        "* How are the results?\n",
        "\n",
        "### 💡 **Exercise (extra)**\n",
        "\n",
        "* To improve the performance of the network for other noise distributions, you can try to randomly sample the noise standard deviation during training.\n",
        "* Add an option in the `DenoisingDataset` for this.\n",
        "* Note that the `add_noise` function has an optinal parameter for the standard deviation.\n",
        "\n",
        "### 💡 **Exercise (extra)**\n",
        "\n",
        "* Try making the network deeper by adding additional layers in the middle.\n",
        "\n",
        "### 💡 **Exercise (extra)**\n",
        "\n",
        "* Replace SGD with the more recent [Adam optimizer](https://pytorch.org/docs/stable/optim.html?highlight=adam#torch.optim.Adam).\n",
        "\n",
        "### 💡 **Exercise (extra)**\n",
        "\n",
        "* Try changing to a [Leaky Relu activation](https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html?highlight=leaky%20relu#torch.nn.LeakyReLU).\n",
        "\n",
        "### 💡 **Exercise (extra)**\n",
        "\n",
        "* Integrate a learning rate scheduler that decreases the learning rate for each epoch.\n",
        "* Use, for example, the [StepLR](https://pytorch.org/docs/stable/optim.html?highlight=schedule#torch.optim.lr_scheduler.StepLR). \n"
      ]
    }
  ]
}