{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a29a94a4",
   "metadata": {},
   "source": [
    "# TP3: Image Segmentation With a U-Net - From CPU to GPU"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29af8073",
   "metadata": {},
   "source": [
    "The goal of this notebook is to introduce the notion of **hardware acceleration** in the context of deep learning. We will use the application of image segmentation with a U-Net to illustrate the difference in performance between a CPU and a GPU.\n",
    "\n",
    "**Image segmentation** is an important task in  computer vision. The goal is to find multiple areas in an image and to assign labels to these area.  It provides a different kind of information than: \n",
    "- **image classification** caracterizes images with global labels;\n",
    "- **object detection** usually relies on finding bounding-boxes around  detected objects\n",
    "\n",
    "Segmentation is useful and can be used in real-world applications such as medical imaging, clothes segmentation, flooding maps, self-driving cars, etc. There are two types of image segmentation:\n",
    "- Semantic segmentation: classify each pixel with a label.\n",
    "- Instance segmentation: classify each pixel and differentiate each object instance.\n",
    "\n",
    "U-Net is a semantic segmentation technique [originally proposed for medical imaging segmentation](https://arxiv.org/abs/1505.04597). It’s one of the earlier deep learning segmentation models. This architecture is still widely used in more advanced models like Generative Adversarial or Diffusion Network. \n",
    "\n",
    "The model architecture is fairly simple: an encoder (for downsampling) and a decoder (for upsampling) with skip connections. U-Net is only based on convolutions. More specifically, the output classification is done at pixel level with a *(1,1)* convolution. It has therefore the following advantages: \n",
    "- parameter and data efficiency, \n",
    "- independent of the input size. \n",
    "\n",
    "The following image is taken from the original paper:\n",
    "\n",
    "<img src=\"https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/u-net-architecture.png\"  width=\"60%\" height=\"30%\">\n",
    "\n",
    "\n",
    "In this session, we will consider a binary segmentation task.  \n",
    "\n",
    "# 1.  Cell nuclei segmentation: the dataset\n",
    "\n",
    "Cell nuclei segmentation is an essential step in the biological analysis of microscopy images. \n",
    "This segmentation can be manually achieved with dedicated software, however it is very costly. \n",
    "In this lab session, the starting point is this [nature paper](https://www.nature.com/articles/s41597-020-00608-w). To quote some part of the paper: \n",
    "\n",
    "Fully-automated nuclear image segmentation is the prerequisite to ensure statistically significant, quantitative analyses of tissue preparations,applied in digital pathology or quantitative microscopy. The design of segmentation methods that work independently of the tissue type or preparation is complex, due to variations in nuclear morphology, staining intensity, cell density and nuclei aggregations. Machine learning-based segmentation methods can overcome these challenges, however high quality expert-annotated images are required for training. Currently, the limited number of annotated fluorescence image datasets publicly available do not cover a broad range of tissues and preparations. We present a comprehensive, annotated dataset including tightly aggregated nuclei of multiple tissues for the training of machine learning-based nuclear segmentation algorithms. The proposed dataset covers sample preparation methods frequently used in quantitative immunofluorescence microscopy. \n",
    "\n",
    "To spare some preprocessing time, this lab session starts with this pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66bcbb87",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip3 install tqdm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision.transforms as transforms\n",
    "import pickle\n",
    "import time\n",
    "from tqdm.notebook import trange, tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f368198f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -O nuclei_cells_segmentations.pck \"https://drive.usercontent.google.com/download?id=1afhnCuKP3A6rtZe2Ze4dAgQ4kZGR2_uD&export=download&confirm=t\"\n",
    "!curl -L -o nuclei_cells_segmentations.pck \"https://drive.usercontent.google.com/download?id=1afhnCuKP3A6rtZe2Ze4dAgQ4kZGR2_uD&export=download&confirm=t\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b34f50ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "fn = \"nuclei_cells_segmentations.pck\"\n",
    "with open(fn, \"rb\") as f:\n",
    "    X, Y = pickle.load(f)\n",
    "print(f\"Nuclei images: {X.shape}, Cell images: {Y.shape}\")\n",
    "N = X.shape[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2426b838",
   "metadata": {},
   "source": [
    "Let us visualize some images from the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b24c73a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15, 3))\n",
    "for i in range(0, 10):\n",
    "    plt.subplot(2, 10, i + 1)\n",
    "    plt.imshow(X[i].squeeze())\n",
    "    plt.axis(\"off\")\n",
    "    if i == 0:\n",
    "        plt.ylabel(\"Nuclei\")\n",
    "    plt.subplot(2, 10, i + 11)\n",
    "    plt.imshow(Y[i].squeeze())\n",
    "    plt.axis(\"off\")\n",
    "    if i == 0:\n",
    "        plt.ylabel(\"Cells\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd8b4881",
   "metadata": {},
   "source": [
    "This pickle contains a modified version of the dataset: \n",
    "- the same amount of images\n",
    "- all the images are resized to 128,128\n",
    "- the segmentation task is converted in a binary pixel classification: nuclei or not. \n",
    "\n",
    "The goal is now to train a U-Net on this dataset (69 images for training and 10 for \"test\"). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c25c8465",
   "metadata": {},
   "outputs": [],
   "source": [
    "M = 69\n",
    "trainset = X[:M], Y[:M]\n",
    "testset = X[M:], Y[M:]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0928ce5c",
   "metadata": {},
   "source": [
    "# 2 U-Net Architecture\n",
    "\n",
    "Following the previous picture of U-Net, the network is composed of 3 parts: encoder, bottleneck, decoder.  \n",
    "These three steps rely on a convolutional block (convolution, relu, convolution, relu) \n",
    "\n",
    "The first step is the **encoder**. The goal is to compress the  \"geometrical\" information with local features. The encoder first applies a convolution of kernel size (3,3) to extract $F=64$ features. Then the information is compressed using max-pooling (factor 2). The next step does the same:  extract $2\\times F=128$ features from the $F=64$, then compression with max-pooling. This operation is repeated 4 times in total to get at the end $F\\times 8 = 512$ channels that represent global features extracted from the input image. \n",
    "\n",
    "The **bottleneck** layer is a convolutional layer which doubles the number of channels. The idea is to create a \"dense\" representation of the image to gather both global and local features. \n",
    "\n",
    "The **decoder** part is similar to the encoder part but reversed. While we used max-pooling for downsampling in the encoder, the upsampling operation consists in **transposed convolution**. The goal is to increase (so upsample) the spatial dimensions of intermediate feature maps. \n",
    "\n",
    "The last peculiarity is the output layer for classification at the pixel level. In U-Net this last layer is (once again) a convolutional layer. This means that with the last hidden layer, we recover the same spatial dimension as the input with $F$ feature maps. The classification is carried out for each pixel independently, but the decision is based on $F$ features that encode global information. \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f64fe5f",
   "metadata": {},
   "source": [
    "## 2. U-Net: step by step\n",
    "\n",
    "Now the goal is to implement U-Net. As a proposed roadmap we propose the following step: \n",
    "- a function to create a convolutional block\n",
    "- a module for the encoder\n",
    "- a module for the decoder\n",
    "- and a U-Net module to wrap everything\n",
    "\n",
    "The number of feature map ($F=64$ in the original work) must be a variable of the UNet. For the first round of experiment, we can use $F=8$.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c56ee22",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_conv_block(in_c, out_c, kernel_size=3, stride=1, padding=1):\n",
    "    \"\"\"\n",
    "    Creates a convolutional block with two convolutional layers, each followed by a ReLU activation and batch normalization.\n",
    "\n",
    "    Args:\n",
    "        in_c (int): Number of input channels.\n",
    "        out_c (int): Number of output channels.\n",
    "        kernel_size (int, optional): Size of the convolving kernel. Default is 3.\n",
    "        stride (int, optional): Stride of the convolution. Default is 1.\n",
    "        padding (int, optional): Zero-padding added to both sides of the input. Default is 1.\n",
    "\n",
    "    Returns:\n",
    "        nn.Sequential: A sequential container of the convolutional block.\n",
    "    \"\"\"\n",
    "    mod = nn.Sequential(\n",
    "        nn.Conv2d(\n",
    "            in_channels=in_c,\n",
    "            out_channels=out_c,\n",
    "            kernel_size=kernel_size,\n",
    "            stride=stride,\n",
    "            padding=padding,\n",
    "        ),\n",
    "        nn.ReLU(),\n",
    "        nn.BatchNorm2d(num_features=out_c),\n",
    "        nn.Conv2d(\n",
    "            in_channels=out_c,\n",
    "            out_channels=out_c,\n",
    "            kernel_size=kernel_size,\n",
    "            stride=stride,\n",
    "            padding=padding,\n",
    "        ),\n",
    "        nn.BatchNorm2d(num_features=out_c),\n",
    "        nn.ReLU(),\n",
    "    )\n",
    "    return mod\n",
    "\n",
    "\n",
    "class encoder_block(nn.Module):\n",
    "    \"\"\"Encoder block:\n",
    "    - a conv block\n",
    "    - followed by a MaxPooling: the height and width is divided by 2\n",
    "    We need both the output before and after the max-pooling.\n",
    "    - the output before the maxpooling is used for the \"residual\" connection\n",
    "    - the output after is given to the next encoder block or the bottleneck block.\n",
    "\n",
    "    Remark: we could use a Sequential module but for the forward we need\n",
    "    to get x and p\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, in_c, out_c):\n",
    "        super().__init__()\n",
    "        self.conv = make_conv_block(in_c, out_c)\n",
    "        self.pool = nn.MaxPool2d(kernel_size=(2, 2))\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        x = self.conv(inputs)\n",
    "        p = self.pool(x)\n",
    "        return x, p\n",
    "\n",
    "\n",
    "class decoder_block(nn.Module):\n",
    "    \"\"\"Decoder block:\n",
    "    - First upsampling is done with a ConvTranspose2D\n",
    "    - The \"residual\" connection is implemented with the cat function\n",
    "    - Then a conv_block is applied.\n",
    "    The ConvTranspose doubles the height and width.\n",
    "    The forward call takes two args: the inputs from the previous block and the residual input.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, in_c, out_c):\n",
    "        super().__init__()\n",
    "        self.up = nn.ConvTranspose2d(in_c, out_c, kernel_size=2, stride=2, padding=0)\n",
    "        self.conv = make_conv_block(out_c + out_c, out_c)\n",
    "\n",
    "    def forward(self, inputs, skip):\n",
    "        x = self.up(inputs)\n",
    "        x = torch.cat([x, skip], axis=1)\n",
    "        x = self.conv(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class UNet(nn.Module):\n",
    "    \"\"\"\n",
    "    UNet model for image segmentation.\n",
    "\n",
    "    Args:\n",
    "        factor (int): Base number of filters for the convolutional layers. Default is 64.\n",
    "\n",
    "    Attributes:\n",
    "        f (int): Base number of filters for the convolutional layers.\n",
    "        e1 (nn.Module): First encoder block.\n",
    "        e2 (nn.Module): Second encoder block.\n",
    "        e3 (nn.Module): Third encoder block.\n",
    "        e4 (nn.Module): Fourth encoder block.\n",
    "        b (nn.Module): Bottleneck convolutional block.\n",
    "        d1 (nn.Module): First decoder block.\n",
    "        d2 (nn.Module): Second decoder block.\n",
    "        d3 (nn.Module): Third decoder block.\n",
    "        d4 (nn.Module): Fourth decoder block.\n",
    "        outputs (nn.Conv2d): Final convolutional layer for generating the output.\n",
    "\n",
    "    Methods:\n",
    "        forward(inputs):\n",
    "            Forward pass of the UNet model.\n",
    "\n",
    "            Args:\n",
    "                inputs (torch.Tensor): Input tensor of shape (batch_size, channels, height, width).\n",
    "\n",
    "            Returns:\n",
    "                torch.Tensor: Output tensor of shape (batch_size, 1, height, width).\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, factor=64):\n",
    "        super().__init__()\n",
    "        self.f = factor\n",
    "        \"\"\" Encoder \"\"\"\n",
    "        self.e1 = encoder_block(1, self.f)\n",
    "        self.e2 = encoder_block(self.f, self.f * 2)\n",
    "        self.e3 = encoder_block(self.f * 2, self.f * 4)\n",
    "        self.e4 = encoder_block(self.f * 4, self.f * 8)\n",
    "\n",
    "        \"\"\" Bottleneck \"\"\"\n",
    "        self.b = make_conv_block(self.f * 8, self.f * 16)\n",
    "\n",
    "        \"\"\" Decoder \"\"\"\n",
    "        self.d1 = decoder_block(self.f * 16, self.f * 8)\n",
    "        self.d2 = decoder_block(self.f * 8, self.f * 4)\n",
    "        self.d3 = decoder_block(self.f * 4, self.f * 2)\n",
    "        self.d4 = decoder_block(self.f * 2, self.f)\n",
    "\n",
    "        \"\"\" Classifier \"\"\"\n",
    "        self.outputs = nn.Conv2d(self.f, 1, kernel_size=1, padding=0)\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        \"\"\"Encoder\"\"\"\n",
    "        s1, p1 = self.e1(inputs)\n",
    "        s2, p2 = self.e2(p1)\n",
    "        s3, p3 = self.e3(p2)\n",
    "        s4, p4 = self.e4(p3)\n",
    "\n",
    "        \"\"\" Bottleneck \"\"\"\n",
    "        b = self.b(p4)\n",
    "\n",
    "        \"\"\" Decoder \"\"\"\n",
    "        d1 = self.d1(b, s4)\n",
    "        d2 = self.d2(d1, s3)\n",
    "        d3 = self.d3(d2, s2)\n",
    "        d4 = self.d4(d3, s1)\n",
    "\n",
    "        \"\"\" Classifier \"\"\"\n",
    "        outputs = self.outputs(d4)\n",
    "        return outputs\n",
    "\n",
    "\n",
    "class LightUNet(nn.Module):\n",
    "    \"\"\"\n",
    "    LightUNet is a lightweight implementation of the U-Net architecture for image segmentation tasks.\n",
    "    Attributes:\n",
    "        f (int): Factor to control the number of filters in the convolutional layers.\n",
    "        e1 (nn.Module): Encoder block for the first stage.\n",
    "        b (nn.Module): Bottleneck convolutional block.\n",
    "        d1 (nn.Module): Decoder block for the first stage.\n",
    "        outputs (nn.Conv2d): Convolutional layer to produce the final output.\n",
    "    Methods:\n",
    "        forward(inputs):\n",
    "            Forward pass through the network.\n",
    "            Args:\n",
    "                inputs (torch.Tensor): Input tensor with shape (batch_size, channels, height, width).\n",
    "            Returns:\n",
    "                torch.Tensor: Output tensor with shape (batch_size, 1, height, width).\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, factor=8):\n",
    "        super().__init__()\n",
    "        self.f = factor\n",
    "        \"\"\" Encoder \"\"\"\n",
    "        self.e1 = encoder_block(1, self.f)\n",
    "\n",
    "        \"\"\" Bottleneck \"\"\"\n",
    "        self.b = make_conv_block(self.f, self.f * 2)\n",
    "\n",
    "        \"\"\" Decoder \"\"\"\n",
    "        self.d1 = decoder_block(self.f * 2, self.f)\n",
    "\n",
    "        \"\"\" Classifier \"\"\"\n",
    "        self.outputs = nn.Conv2d(self.f, 1, kernel_size=1, padding=0)\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        \"\"\"Encoder\"\"\"\n",
    "        s1, p1 = self.e1(inputs)\n",
    "\n",
    "        \"\"\" Bottleneck \"\"\"\n",
    "        b = self.b(p1)\n",
    "\n",
    "        \"\"\" Decoder \"\"\"\n",
    "        d1 = self.d1(b, s1)\n",
    "\n",
    "        \"\"\" Classifier \"\"\"\n",
    "        outputs = self.outputs(d1)\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19c33ddc",
   "metadata": {},
   "source": [
    "We can now define two different U-Net with different number of features. Let us compare the speed of inference between the two models.\n",
    "By default, the model is loaded on CPU. We will then compare the speed of inference on CPU and GPU.\n",
    "\n",
    "\n",
    "First, let us check if your computer has a GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b795d84f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c67434b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"GPU available using Cuda: {torch.cuda.is_available()}\")\n",
    "print(\n",
    "    f\"GPU available using MPS: {torch.backends.mps.is_available() and torch.backends.mps.is_built()}\"\n",
    ")\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "    gpu_available = True\n",
    "elif torch.backends.mps.is_available() and torch.backends.mps.is_built():\n",
    "    device = torch.device(\"mps\")\n",
    "    gpu_available = True\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "    gpu_available = False\n",
    "\n",
    "\n",
    "print(f\"Device: {device}\")\n",
    "print(f\"Number of GPUs: {torch.cuda.device_count()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65e35ec0",
   "metadata": {},
   "source": [
    "<font color='blue'>TODO:</font> \n",
    "- Run the following cell to compare the speed of inference on CPU and GPU.\n",
    "- Comment the results.\n",
    "- Consider running your code on Colab if your own computer does not have a GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8aab7f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_cpu = UNet(factor=64)\n",
    "print(f\"Number of parameters: {sum(p.numel() for p in model_cpu.parameters())}\")\n",
    "light_model_cpu = LightUNet(factor=8)\n",
    "print(f\"Number of parameters: {sum(p.numel() for p in light_model_cpu.parameters())}\")\n",
    "\n",
    "times_model_cpu = []\n",
    "times_light_model_cpu = []\n",
    "model_cpu.eval()\n",
    "light_model_cpu.eval()\n",
    "if torch.cuda.is_available():\n",
    "    times_model_gpu = []\n",
    "    times_light_model_gpu = []\n",
    "    model_gpu = UNet(factor=64).cuda()\n",
    "    light_model_gpu = LightUNet(factor=8).cuda()\n",
    "    model_gpu.eval()\n",
    "    light_model_gpu.eval()\n",
    "\n",
    "for _ in trange(10):\n",
    "    t0 = time.time()\n",
    "    model_cpu(trainset[0])\n",
    "    t1 = time.time()\n",
    "    times_model_cpu.append(t1 - t0)\n",
    "    t0 = time.time()\n",
    "    light_model_cpu(trainset[0])\n",
    "    t1 = time.time()\n",
    "    times_light_model_cpu.append(t1 - t0)\n",
    "    if torch.cuda.is_available():\n",
    "        t0 = time.time()\n",
    "        model_gpu(trainset[0].cuda())\n",
    "        t1 = time.time()\n",
    "        times_model_gpu.append(t1 - t0)\n",
    "        t0 = time.time()\n",
    "        light_model_gpu(trainset[0].cuda())\n",
    "        t1 = time.time()\n",
    "        times_light_model_gpu.append(t1 - t0)\n",
    "print(f\"UNet CPU: {np.mean(times_model_cpu):.3f} s\")\n",
    "if torch.cuda.is_available():\n",
    "    print(f\"UNet GPU: {np.mean(times_model_gpu):.3f} s\")\n",
    "    print(f\"Inference speedup: {np.mean(times_model_cpu)/np.mean(times_model_gpu):.1f}\")\n",
    "print(f\"LightUNet CPU: {np.mean(times_light_model_cpu):.3f} s\")\n",
    "if torch.cuda.is_available():\n",
    "    print(f\"LightUNet GPU: {np.mean(times_light_model_gpu):.3f} s\")\n",
    "    print(\n",
    "        f\"Inference speedup: {np.mean(times_light_model_cpu)/np.mean(times_light_model_gpu):.1f}\"\n",
    "    )\n",
    "    del model_gpu, light_model_gpu\n",
    "del model_cpu, light_model_cpu"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8749c614",
   "metadata": {},
   "source": [
    "We can also compare the speed of inference of GPU with different batch sizes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3d68023",
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    time_inference_cpu_mean = []\n",
    "    time_inference_gpu_mean = []\n",
    "    time_inference_cpu_std = []\n",
    "    time_inference_gpu_std = []\n",
    "    light_model_cpu = LightUNet(factor=8)\n",
    "    light_model_cpu.eval()\n",
    "    light_model_gpu = LightUNet(factor=8).cuda()\n",
    "    light_model_gpu.eval()\n",
    "    with torch.no_grad():\n",
    "        for i in trange(0, 7):\n",
    "            time_inference_cpu = []\n",
    "            time_inference_gpu = []\n",
    "            batch_size = 2**i\n",
    "            for i in range(20):\n",
    "                t0 = time.time()\n",
    "                light_model_cpu(trainset[0][:batch_size])\n",
    "                t1 = time.time()\n",
    "                time_inference_cpu.append(t1 - t0)\n",
    "            time_inference_cpu_mean.append(np.mean(time_inference_cpu))\n",
    "            time_inference_cpu_std.append(np.std(time_inference_cpu))\n",
    "            for i in range(20):\n",
    "                t0 = time.time()\n",
    "                light_model_gpu(trainset[0][:batch_size].cuda())\n",
    "                t1 = time.time()\n",
    "                time_inference_gpu.append(t1 - t0)\n",
    "            time_inference_gpu_mean.append(np.mean(time_inference_gpu))\n",
    "            time_inference_gpu_std.append(np.std(time_inference_gpu))\n",
    "    del light_model_cpu, light_model_gpu\n",
    "\n",
    "    x = [2**i for i in range(7)]\n",
    "    plt.figure(figsize=(15, 5))\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.plot(x, time_inference_cpu_mean, label=\"CPU\")\n",
    "    plt.fill_between(\n",
    "        x,\n",
    "        np.array(time_inference_cpu_mean) - np.array(time_inference_cpu_std),\n",
    "        np.array(time_inference_cpu_mean) + np.array(time_inference_cpu_std),\n",
    "        alpha=0.3,\n",
    "    )\n",
    "    affine_approx = np.polyfit(x[3:], time_inference_cpu_mean[3:], 1)\n",
    "    plt.plot(\n",
    "        x,\n",
    "        np.polyval(affine_approx, x),\n",
    "        linestyle=\"--\",\n",
    "        color=\"black\",\n",
    "        label=f\"Affine approx {affine_approx[0]:.3e}x\",\n",
    "    )\n",
    "\n",
    "    plt.xlabel(\"Batch size\")\n",
    "    plt.ylabel(\"Inference time (s)\")\n",
    "    plt.legend()\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.plot(x, time_inference_gpu_mean, label=\"GPU\")\n",
    "    plt.fill_between(\n",
    "        x,\n",
    "        np.array(time_inference_gpu_mean) - np.array(time_inference_gpu_std),\n",
    "        np.array(time_inference_gpu_mean) + np.array(time_inference_gpu_std),\n",
    "        alpha=0.3,\n",
    "    )\n",
    "    affine_approx = np.polyfit(x[3:], time_inference_gpu_mean[3:], 1)\n",
    "    plt.plot(\n",
    "        x,\n",
    "        np.polyval(affine_approx, x),\n",
    "        linestyle=\"--\",\n",
    "        color=\"black\",\n",
    "        label=f\"Affine approx {affine_approx[0]:.3e}x\",\n",
    "    )\n",
    "    plt.xlabel(\"Batch size\")\n",
    "    plt.ylabel(\"Inference time (s)\")\n",
    "    plt.title(\"Inference time for LightUNet\")\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c2a875f",
   "metadata": {},
   "source": [
    "We can do the same thing with the training time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cc41b46",
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    time_training_cpu_mean = []\n",
    "    time_training_gpu_mean = []\n",
    "    time_training_cpu_std = []\n",
    "    time_training_gpu_std = []\n",
    "    light_model_cpu = LightUNet(factor=8)\n",
    "    light_model_cpu.train()\n",
    "    light_model_gpu = LightUNet(factor=8).cuda()\n",
    "    light_model_gpu.train()\n",
    "    criterion = nn.MSELoss()\n",
    "    optimizer_cpu = torch.optim.Adam(light_model_cpu.parameters(), lr=0.001)\n",
    "    optimizer_gpu = torch.optim.Adam(light_model_gpu.parameters(), lr=0.001)\n",
    "    for i in trange(0, 7):\n",
    "        time_inference_cpu = []\n",
    "        time_inference_gpu = []\n",
    "        batch_size = 2**i\n",
    "        for i in range(2):\n",
    "            t0 = time.time()\n",
    "            optimizer_cpu.zero_grad()\n",
    "            outputs = light_model_cpu(trainset[0][:batch_size])\n",
    "            loss = criterion(outputs, trainset[1][:batch_size])\n",
    "            loss.backward()\n",
    "            optimizer_cpu.step()\n",
    "            t1 = time.time()\n",
    "            time_inference_cpu.append(t1 - t0)\n",
    "        time_training_cpu_mean.append(np.mean(time_inference_cpu))\n",
    "        time_training_cpu_std.append(np.std(time_inference_cpu))\n",
    "        for i in range(20):\n",
    "            t0 = time.time()\n",
    "            optimizer_gpu.zero_grad()\n",
    "            outputs = light_model_gpu(trainset[0][:batch_size].cuda())\n",
    "            loss = criterion(outputs, trainset[1][:batch_size].cuda())\n",
    "            loss.backward()\n",
    "            optimizer_gpu.step()\n",
    "            t1 = time.time()\n",
    "            time_inference_gpu.append(t1 - t0)\n",
    "        time_training_gpu_mean.append(np.mean(time_inference_gpu))\n",
    "        time_training_gpu_std.append(np.std(time_inference_gpu))\n",
    "    del light_model_cpu, light_model_gpu\n",
    "\n",
    "    x = [2**i for i in range(7)]\n",
    "    plt.figure(figsize=(15, 5))\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.plot(x, time_training_cpu_mean, label=\"CPU\")\n",
    "    plt.fill_between(\n",
    "        x,\n",
    "        np.array(time_training_cpu_mean) - np.array(time_training_cpu_std),\n",
    "        np.array(time_training_cpu_mean) + np.array(time_training_cpu_std),\n",
    "        alpha=0.3,\n",
    "    )\n",
    "    affine_approx = np.polyfit(x[3:], time_training_cpu_mean[3:], 1)\n",
    "    plt.plot(\n",
    "        x,\n",
    "        np.polyval(affine_approx, x),\n",
    "        linestyle=\"--\",\n",
    "        color=\"black\",\n",
    "        label=f\"Affine approx {affine_approx[0]:.3e}x\",\n",
    "    )\n",
    "    plt.xlabel(\"Batch size\")\n",
    "    plt.ylabel(\"Training time (s)\")\n",
    "    plt.legend()\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.plot(x, time_training_gpu_mean, label=\"GPU\")\n",
    "    plt.fill_between(\n",
    "        x,\n",
    "        np.array(time_training_gpu_mean) - np.array(time_training_gpu_std),\n",
    "        np.array(time_training_gpu_mean) + np.array(time_training_gpu_std),\n",
    "        alpha=0.3,\n",
    "    )\n",
    "    affine_approx = np.polyfit(x[3:], time_training_gpu_mean[3:], 1)\n",
    "    plt.plot(\n",
    "        x,\n",
    "        np.polyval(affine_approx, x),\n",
    "        linestyle=\"--\",\n",
    "        color=\"black\",\n",
    "        label=f\"Affine approx {affine_approx[0]:.3e}x\",\n",
    "    )\n",
    "    plt.xlabel(\"Batch size\")\n",
    "    plt.ylabel(\"Training time (s)\")\n",
    "    plt.title(\"Training time for LightUNet\")\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cfc335e",
   "metadata": {},
   "source": [
    "<font color='blue'>TODO:</font>\n",
    "- Compare the training/inference time on CPU and GPU.\n",
    "- Comment the results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c6a760c",
   "metadata": {},
   "source": [
    "From now on, the code will be GPU or CPU agnostic. The code will automatically detect if a GPU is available and use it. If not, the code will run on CPU. To do so, we will use the `torch.device` object. \n",
    "\n",
    "```python\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = model.to(device)\n",
    "input = input.to(device)\n",
    "```\n",
    "\n",
    "And every time, we need to manipulate results, we will use the following code: \n",
    "\n",
    "```python\n",
    "output = model(input)\n",
    "output = output.cpu().detach().numpy()\n",
    "```\n",
    "The `detach()` method is used to detach the output from the computational graph (no Grad attached). The `cpu()` method is used to move the output from the GPU to the CPU. The `numpy()` method is used to convert the output to a numpy array.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "636ac03f",
   "metadata": {},
   "source": [
    "# 3. Training the U-Net\n",
    "\n",
    "We will now build functions to train the U-Net. The next cell contains the following functions:\n",
    "- a function to train the model and evaluate the model\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b504c7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(\n",
    "    model, trainset, devset=None, nepochs=10, lr0=1e-2, threshold=0, device=\"cpu\"\n",
    "):\n",
    "    \"\"\"Train a model\n",
    "    Args:\n",
    "    - model: the model to train\n",
    "    - trainset: a tuple (X,Y) of the training data\n",
    "    - devset: a tuple (X,Y) of the development data\n",
    "    - nepochs: number of epochs\n",
    "    - lr0: initial learning rate\n",
    "    - threshold: threshold for the accuracy\n",
    "    - device: the device to use\n",
    "    Output:\n",
    "    - the trained model\n",
    "    \"\"\"\n",
    "    # Prepare the data, loss, and optimizer\n",
    "    criterion = nn.BCEWithLogitsLoss()\n",
    "    optim = torch.optim.Adam(model.parameters(), lr0)\n",
    "    losses = torch.zeros(nepochs)\n",
    "    dlosses = torch.zeros(nepochs)\n",
    "    accs = torch.zeros(nepochs)\n",
    "    train, labels = trainset\n",
    "    dev, dlabels = devset\n",
    "\n",
    "    # Training loop\n",
    "    for e in range(nepochs):\n",
    "        ### train part\n",
    "        optim.zero_grad()\n",
    "        preds = model(train)\n",
    "        l = criterion(preds, labels)\n",
    "        l.backward()\n",
    "        optim.step()\n",
    "        losses[e] = l.item()\n",
    "        accs[e] = (\n",
    "            ((preds > threshold) * 1 == labels).sum()\n",
    "            * 100\n",
    "            / (preds.shape[0] * preds.shape[1] * preds.shape[2] * preds.shape[3])\n",
    "        )\n",
    "        ### dev eval\n",
    "        model.eval()\n",
    "        preds = model(dev)\n",
    "        l = criterion(preds, dlabels)\n",
    "        dlosses[e] = l.item()\n",
    "        model.train()\n",
    "        if e % 10 == 0:\n",
    "            print(\n",
    "                f\"Epoch {e} | Loss {losses[e]:.3f} | Acc {accs[e]:.2f} | Dev Loss {dlosses[e]:.3f}\"\n",
    "            )\n",
    "    print(\n",
    "        f\"Epoch {e} | Loss {losses[e]:.3f} | Acc {accs[e]:.2f} | Dev Loss {dlosses[e]:.3f}\"\n",
    "    )\n",
    "    plt.figure(figsize=(10, 5))\n",
    "    plt.plot(losses, label=\"Train\")\n",
    "    plt.plot(dlosses, label=\"Dev\")\n",
    "    plt.xlabel(\"Epoch\")\n",
    "    plt.ylabel(\"Loss\")\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0e2043f",
   "metadata": {},
   "source": [
    "<font color='blue'>TODO:</font>\n",
    "- Run the following cell to train the U-Net for 10 epochs (on CPU).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "869ae11c",
   "metadata": {},
   "outputs": [],
   "source": [
    "m8 = LightUNet(8)\n",
    "m8 = train(m8, trainset, testset, 10, lr0=2e-2, threshold=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='blue'>TODO:</font>\n",
    "- Modify the code to use the GPU.\n",
    "- Train the U-Net for 100 epochs.\n",
    "- Compare the training time between CPU and GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "293c681f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(\n",
    "    model, trainset, devset=None, nepochs=10, lr0=1e-2, threshold=0, device=\"cpu\"\n",
    "):\n",
    "    \"\"\"Train a model\n",
    "    Args:\n",
    "    - model: the model to train\n",
    "    - trainset: a tuple (X,Y) of the training data\n",
    "    - devset: a tuple (X,Y) of the development data\n",
    "    - nepochs: number of epochs\n",
    "    - lr0: initial learning rate\n",
    "    - threshold: threshold for the accuracy\n",
    "    - device: the device to use\n",
    "    Output:\n",
    "    - the trained model\n",
    "    \"\"\"\n",
    "    ### Code to complete here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "818b7861",
   "metadata": {},
   "outputs": [],
   "source": [
    "m8 = LightUNet(8).to(device)\n",
    "m8 = train(m8, trainset, testset, 100, lr0=2e-2, threshold=0, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d71bb350",
   "metadata": {},
   "source": [
    "<font color='blue'>TODO:</font>\n",
    "- Modify the code to return the model with the best validation loss.\n",
    "- Ideally, show on the plot the best model.\n",
    "\n",
    "One way to do so, is to save the model for the best validation loss and load it at the end of the training.\n",
    "To do so, we will use the following code to save the model at a given epoch. \n",
    "\n",
    "```python\n",
    "state_dict = model.state_dict()\n",
    "torch.save(state_dict, 'model.pth')\n",
    "```\n",
    "\n",
    "And the following code to load the model at the end of the training. \n",
    "\n",
    "```python\n",
    "model.load_state_dict(torch.load('model.pth'))\n",
    "```\n",
    "When the model is light (like in this case), you can simply track the weight directly using the following code: \n",
    "\n",
    "```python\n",
    "weight = model.state_dict()\n",
    "model.load_state_dict(weight)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e998851d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(\n",
    "    model, trainset, devset=None, nepochs=10, lr0=1e-2, threshold=0, device=\"cpu\"\n",
    "):\n",
    "    \"\"\"Train a model\n",
    "    Args:\n",
    "    - model: the model to train\n",
    "    - trainset: a tuple (X,Y) of the training data\n",
    "    - devset: a tuple (X,Y) of the development data\n",
    "    - nepochs: number of epochs\n",
    "    - lr0: initial learning rate\n",
    "    - threshold: threshold for the accuracy\n",
    "    - device: the device to use\n",
    "    Output:\n",
    "    - the best trained model\n",
    "    \"\"\"\n",
    "    ### Code to complete here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "129e75b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "m8 = LightUNet(8).to(device)\n",
    "m8 = train(m8, trainset, testset, 100, lr0=2e-2, threshold=0, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c37cb2b6",
   "metadata": {},
   "source": [
    "The following cell contains the code to assess the performance of the model on the test set. Include the following functions:\n",
    "- a function to plot the results\n",
    "- a function to plot the PR curve\n",
    "\n",
    "The PR curve is a good way to assess the performance of a binary classifier. It is a plot of the precision against the recall for the different thresholds. \n",
    "\n",
    "The precision is the number of true positive divided by the number of true positive and false positive. Intuitively, it is the ability of the classifier not to label as positive a sample that is negative.\n",
    "\n",
    "The recall is the number of true positive divided by the number of true positive and false negative. Intuitively, it is the ability of the classifier to find all the positive samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3fcea2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plotres(m, testset, device):\n",
    "    \"\"\"Plot the results of a segmentation model\n",
    "    Args:\n",
    "    - m: the model\n",
    "    - testset: the test set\n",
    "    \"\"\"\n",
    "    m.eval()\n",
    "    x, y = testset\n",
    "    if device == torch.device(\"cuda\"):\n",
    "        x, y = x.cuda(), y.cuda()\n",
    "    testid = range(x.shape[0])\n",
    "    plt.figure(figsize=(5, 10))\n",
    "    for i in testid:\n",
    "        out = m(x[i].unsqueeze(0)).squeeze().detach()\n",
    "        plt.subplot(10, 4, i * 4 + 1)\n",
    "        if i == 0:\n",
    "            plt.title(\"Input\")\n",
    "        plt.imshow(x[i].cpu().squeeze())\n",
    "        plt.axis(\"off\")\n",
    "        plt.subplot(10, 4, i * 4 + 2)\n",
    "        if i == 0:\n",
    "            plt.title(\"Output\")\n",
    "        out = m(x[i].unsqueeze(0)).squeeze().detach()\n",
    "        plt.imshow(out.cpu())\n",
    "        plt.axis(\"off\")\n",
    "        plt.subplot(10, 4, i * 4 + 3)\n",
    "        if i == 0:\n",
    "            plt.title(\"Output>0\")\n",
    "        plt.imshow((out > 0).cpu())\n",
    "        plt.axis(\"off\")\n",
    "        plt.subplot(10, 4, i * 4 + 4)\n",
    "        if i == 0:\n",
    "            plt.title(\"Ground Truth\")\n",
    "        plt.imshow(y[i].cpu().squeeze() == 1)\n",
    "        plt.axis(\"off\")\n",
    "    plt.show()\n",
    "    m.train()\n",
    "\n",
    "\n",
    "def evalseg(preds, labels, thr=0):\n",
    "    \"\"\"Evaluate a segmentation model\n",
    "    Args:\n",
    "    - preds: the predictions\n",
    "    - labels: the ground truth\n",
    "    - thr: the threshold\n",
    "    Output:\n",
    "    - accuracy, precision, recall\n",
    "    \"\"\"\n",
    "    Tot = preds.numel()\n",
    "    if preds.ndim == 3:\n",
    "        preds = preds.unsqueeze(0)\n",
    "    out = preds.squeeze().detach()\n",
    "    selx = out > thr\n",
    "    sely = (labels == 1).squeeze()\n",
    "    x1 = (selx * 1).sum()\n",
    "    y1 = (sely * 1).sum()\n",
    "    inters = selx == sely\n",
    "    acc = inters.sum() * 100 / Tot\n",
    "    prec = selx[inters].sum() * 100 / x1\n",
    "    rapp = sely[inters].sum() * 100 / y1\n",
    "    return acc.cpu().item(), prec.cpu().item(), rapp.cpu().item()\n",
    "\n",
    "\n",
    "def plotroc(m, testset, start=-2, end=2, nsteps=50, device=\"cpu\"):\n",
    "    \"\"\"Plot the ROC curve of a segmentation model\n",
    "    Args:\n",
    "    - m: the model or a list of models\n",
    "    - testset: the test set\n",
    "    - start: the start of the threshold\n",
    "    - end: the end of the threshold\n",
    "    - nsteps: the number of steps\n",
    "    - device: the device\n",
    "    \"\"\"\n",
    "    plt.figure(figsize=(5, 5))\n",
    "    lm = None\n",
    "    if isinstance(m, list):\n",
    "        lm = m\n",
    "    else:\n",
    "        lm = [m]\n",
    "    idx = 0\n",
    "    x, y = testset\n",
    "    if device == torch.device(\"cuda\"):\n",
    "        x, y = x.cuda(), y.cuda()\n",
    "    for m in lm:\n",
    "        preds = m(x)\n",
    "        roc = []\n",
    "        for thr in np.linspace(start, end, nsteps):\n",
    "            acc, prec, rap = evalseg(preds, y, thr)\n",
    "            roc.append([prec, rap])\n",
    "            # print(thr, prec, rap)\n",
    "        roc = np.array(roc)\n",
    "        plt.scatter(roc[:, 0], roc[:, 1], label=str(idx), alpha=0.5)\n",
    "        idx += 1\n",
    "    plt.xlabel(\"Precision\")\n",
    "    plt.ylabel(\"Recall\")\n",
    "    plt.legend()\n",
    "    plt.xlim(0, 100)\n",
    "    plt.ylim(0, 100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fd7072f",
   "metadata": {},
   "source": [
    "<font color='blue'>TODO:</font>\n",
    "- Run the following cell to assess the performance of the model on the test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "981fbf1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "plotres(m8, testset, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c77e298",
   "metadata": {},
   "outputs": [],
   "source": [
    "plotroc(m8, testset, -5, 5, 50, device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d63aacdd",
   "metadata": {},
   "source": [
    "<font color='blue'>TODO:</font>\n",
    "- Train different U-Net with different number of features (4 and 16).\n",
    "- Make sure to have the best learning rate for each model by observing the training loss.\n",
    "- Compare the performance of the different models on the test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99fc8fd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "m4 = LightUNet(4).to(device)\n",
    "m4 = train(m4, trainset, testset, 100, lr0=..., device=device)\n",
    "plotres(m4, testset, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1233514",
   "metadata": {},
   "outputs": [],
   "source": [
    "m16 = LightUNet(16).to(device)\n",
    "m16 = train(m16, trainset, testset, 100, lr0=..., device=device)\n",
    "plotres(m16, testset, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8c1f24f",
   "metadata": {},
   "outputs": [],
   "source": [
    "plotroc([m4, m8, m16], testset, -5, 5, 50, device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5248ef1",
   "metadata": {},
   "source": [
    "<font color='blue'>TODO:</font>\n",
    "- Train a larger U-Net with 8 and 16 features for 200 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c04bb1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "bm8 = UNet(8).to(device)\n",
    "bm8 = train(bm8, trainset, testset, 200, lr0=..., device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4792833a",
   "metadata": {},
   "outputs": [],
   "source": [
    "plotroc([m8, bm8], testset, -5, 5, 50, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2458d5b",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "bm16 = UNet(16).to(device)\n",
    "bm16 = train(bm16, trainset, testset, 200, lr0=..., device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f626b1d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "plotroc([m16, bm16], testset, -5, 5, 50, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4266a775",
   "metadata": {},
   "outputs": [],
   "source": [
    "bm32 = UNet(32).to(device)\n",
    "bm32 = train(bm32, trainset, testset, 200, lr0=..., device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17ba79c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "plotroc([bm8, bm16, bm32], testset, -5, 5, 50, device)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}