UNET.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "UNET.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyPYXE4jvtyDTlCVtGaWUrqY",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sakethbachu/UNET/blob/master/UNET.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UA6XHQiu72bU",
        "colab_type": "code",
        "outputId": "bb2928b2-8120-432d-9e5c-4d319c709658",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 218
        }
      },
      "source": [
        "!pip3 install http://download.pytorch.org/whl/cu80/torch-0.3.0.post4-cp36-cp36m-linux_x86_64.whl "
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting torch==0.3.0.post4\n",
            "\u001b[?25l  Downloading http://download.pytorch.org/whl/cu80/torch-0.3.0.post4-cp36-cp36m-linux_x86_64.whl (592.3MB)\n",
            "\u001b[K     |████████████████████████████████| 592.3MB 74.6MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torch==0.3.0.post4) (1.17.5)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.6/dist-packages (from torch==0.3.0.post4) (3.13)\n",
            "\u001b[31mERROR: torchvision 0.4.2 has requirement torch==1.3.1, but you'll have torch 0.3.0.post4 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: fastai 1.0.60 has requirement torch>=1.0.0, but you'll have torch 0.3.0.post4 which is incompatible.\u001b[0m\n",
            "Installing collected packages: torch\n",
            "  Found existing installation: torch 1.3.1\n",
            "    Uninstalling torch-1.3.1:\n",
            "      Successfully uninstalled torch-1.3.1\n",
            "Successfully installed torch-0.3.0.post4\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K22Q10XLkVTx",
        "colab_type": "code",
        "outputId": "f5dd5c45-2517-4b2f-8ce0-56492dd7c5bf",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 218
        }
      },
      "source": [
        "!pip3 install torchvision"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.6/dist-packages (0.4.2)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from torchvision) (1.12.0)\n",
            "Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.6/dist-packages (from torchvision) (6.2.2)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torchvision) (1.17.5)\n",
            "Collecting torch==1.3.1\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/88/95/90e8c4c31cfc67248bf944ba42029295b77159982f532c5689bcfe4e9108/torch-1.3.1-cp36-cp36m-manylinux1_x86_64.whl (734.6MB)\n",
            "\u001b[K     |████████████████████████████████| 734.6MB 12kB/s \n",
            "\u001b[?25hInstalling collected packages: torch\n",
            "  Found existing installation: torch 0.3.0.post4\n",
            "    Uninstalling torch-0.3.0.post4:\n",
            "      Successfully uninstalled torch-0.3.0.post4\n",
            "Successfully installed torch-1.3.1\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Fg_Td_-nlVgp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.utils.data as data\n",
        "import torchvision\n",
        "from torchvision import transforms\n",
        "import torch.optim as optim\n",
        "import matplotlib.pyplot as plt\n",
        "import cv2\n",
        "import os"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FronK_QvlfEp",
        "colab_type": "code",
        "outputId": "fbb12eb7-5615-4620-fb6e-ebabf6021780",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 121
        }
      },
      "source": [
        "\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive');"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Go to this URL in a browser: https://accounts.google.com/o/oauth2/auth?client_id=947318989803-6bn6qk8qdgf4n4g3pfee6491hc0brc4i.apps.googleusercontent.com&redirect_uri=urn%3aietf%3awg%3aoauth%3a2.0%3aoob&response_type=code&scope=email%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdocs.test%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive.photos.readonly%20https%3a%2f%2fwww.googleapis.com%2fauth%2fpeopleapi.readonly\n",
            "\n",
            "Enter your authorization code:\n",
            "··········\n",
            "Mounted at /content/gdrive\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DTnAanQ6JIGu",
        "colab_type": "text"
      },
      "source": [
        "Defining two functions, one loading the raw input images and the other loading the labels. Meanwhile resizing the images to 128*128. Converting the labels to gray scale images as CE loss requires (N).N:mini batch size. The softmax is already applied for the CEloss."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JBuC7vQkmVpC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "def load_img(folder):\n",
        "  #an iter variable\n",
        "    images = []\n",
        "    for filename in os.listdir(folder): \n",
        "        img = cv2.imread(os.path.join(folder, filename)) \n",
        "        img = cv2.resize(img, (128, 128), interpolation = cv2.INTER_AREA)\n",
        "        images.append(img)\n",
        "    return images\n",
        "\n",
        "def load_label_img(folder):\n",
        "  img1 = [] #an empty list for the files\n",
        "  for filename in os.listdir(folder): \n",
        "      img = cv2.imread(os.path.join(folder, filename))\n",
        "      img = cv2.resize(img, (128, 128), interpolation = cv2.INTER_AREA)\n",
        "      img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n",
        "      img1.append(img)\n",
        "  return img1"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tfUqkuNgtIuh",
        "colab_type": "text"
      },
      "source": [
        "Converting the colors to grayscale and appending them in a list. This helps in giving back the index which will be the class number and this labels the image. The new labels will be in the form of numpy array of the form (128, 128, 1)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "powz9MVqqYOR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def colortogray(cn):\n",
        "    cn = np.reshape(cn, (1, 1, 3));\n",
        "    cn = cv2.cvtColor(cn, cv2.COLOR_BGR2GRAY);\n",
        "    return cn\n",
        "#these are the colors that are used for making the boundaries(ie classfication colours)    \n",
        "colors = [];\n",
        "colors.append(colortogray(np.array([64, 128, 64], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([128, 0, 192], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([192, 128, 0], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([64, 128, 0], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([0, 0, 128], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([128, 0, 64], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([192, 0, 64], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([64, 128, 192], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([128, 192, 192], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([128, 64, 64], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([192, 0, 128], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([64, 0, 192], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([64, 128, 128], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([192, 0, 192], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([64, 64, 128], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([128, 192, 64], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([0, 64, 64], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([128, 64, 128], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([192, 128, 128], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([192, 0, 0], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([128, 128, 192], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([128, 128, 128], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([192, 128, 64], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([64, 0, 0], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([64, 64, 0], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([128, 64, 192], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([0, 128, 128], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([192, 128, 192], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([64, 0, 64], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([0, 192, 192], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([0, 0, 0], dtype = 'uint8')))\n",
        "colors.append(colortogray(np.array([0, 192, 64], dtype = 'uint8')))\n",
        "\n",
        "def class_pixel(label_img):\n",
        "     #      if label_img[0].any() == c[0] and label_img[1].any() == c[1] and label_img[2].any() == c[2]:\n",
        "      #for i in range(128):\n",
        "          #for j in range(128):\n",
        "            #for k in range(3):\n",
        "             # if label_img[k, i, j].any() == c.any():\n",
        "                  #class_pix = index\n",
        "    #return class_pix\n",
        "\n",
        "    class_pix = np.ones([128, 128, 1], dtype = int);\n",
        "    for index, c in enumerate(colors):\n",
        "        class_pix[label_img == c] = index\n",
        "    return class_pix\n",
        "\n",
        "# Convert all segmented images into labeled images with appropriate class number.\n",
        "\n",
        "def label_img_list(img_list):\n",
        "    images = []\n",
        "    for image in img_list:\n",
        "        images.append(class_pixel(image))\n",
        "    return images\n",
        "           "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_T0SqN_ysy-9",
        "colab_type": "text"
      },
      "source": [
        "We are implementing the advanced dataloader to implement mini batch size. We are implementing magic methods such as getitem() and len(). The former returns  an item of the given index and the latter gives the length of the dataset. Also we are defining the transforms that are applied for the data that is loaded."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wJMCbG2FV9Wy",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "transform_img = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])\n",
        "transform_img_label = transforms.Compose([transforms.ToTensor()])\n",
        "\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "\n",
        "#Advance Dataloaders\n",
        "\n",
        "class trainset(data.Dataset):\n",
        "    def __init__(self, transform = None, root_train = None, root_train_label = None, transform_label = None):\n",
        "        self.train_img = load_img(root_train)\n",
        "        self.transform = transform\n",
        "        self.transform_label = transform_label\n",
        "        self.train_label_img = label_img_list(load_label_img(root_train_label))\n",
        "        \n",
        "    def __len__(self):\n",
        "        return len(self.train_img)\n",
        "    \n",
        "    def __getitem__(self, index):\n",
        "        img = self.transform(self.train_img[index])\n",
        "        label = self.transform_label(self.train_label_img[index])\n",
        "        return img, label\n",
        "\n",
        "class valset(data.Dataset):\n",
        "    def __init__(self, transform = None, root_val = None, root_val_label = None, transform_label = None):\n",
        "        self.val_img = load_img(root_val)\n",
        "        self.transform = transform\n",
        "        self.transform_label = transform_label\n",
        "        self.val_label_img = label_img_list(load_label_img(root_val_label))\n",
        "        \n",
        "    def __len__(self):\n",
        "        return len(self.val_img);\n",
        "    \n",
        "    def __getitem__(self, index):\n",
        "        img = self.transform(self.val_img[index])\n",
        "        label = self.transform_label(self.val_label_img[index])\n",
        "        return img, label\n",
        "\n",
        "    \n",
        "class testset(data.Dataset):\n",
        "    def __init__(self, transform = None, root_test = None, root_test_label = None, transform_label = None):\n",
        "        self.test_img = load_img(root_test)\n",
        "        self.transform = transform\n",
        "        self.transform_label = transform_label\n",
        "        self.test_label_img = label_img_list(load_label_img(root_test_label))\n",
        "        \n",
        "    def __len__(self):\n",
        "        return len(self.test_img)\n",
        "    \n",
        "    def __getitem__(self, index):\n",
        "        img = self.transform(self.test_img[index])\n",
        "        label = self.transform_label(self.test_label_img[index])\n",
        "        return img, label\n",
        "traindataset = trainset(transform_img, '/content/gdrive/My Drive/Camvid dataset/train/', '/content/gdrive/My Drive/Camvid dataset/train_labels/', transform_img_label)\n",
        "# valdataset = valset(transform_img, '/content/gdrive/My Drive/Camvid dataset/val/', '/content/gdrive/My Drive/Camvid dataset/val_labels/', transform_img_label)\n",
        "testdataset = testset(transform_img, '/content/gdrive/My Drive/Camvid dataset/test/', '/content/gdrive/My Drive/Camvid dataset/test_labels/', transform_img_label)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kIRRdolLvAFd",
        "colab_type": "text"
      },
      "source": [
        "Defining the train and test loader with batch size 1, for not filling up the gpu space."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kL4wo80Avuyf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train_loader = data.DataLoader(traindataset, batch_size = 1, shuffle=True,  num_workers=4)\n",
        "test_loader = data.DataLoader(testdataset, batch_size = 1, shuffle=True,  num_workers=4)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5l_nRm-pvzZT",
        "colab_type": "text"
      },
      "source": [
        "Defining the Unet model, the first part of the U is made by conv2d class and the second part is made by convTranspose2d."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BxfcZ6wlhhfi",
        "colab_type": "code",
        "outputId": "1908871b-2e22-45a2-e42d-103eaabe636b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 437
        }
      },
      "source": [
        "class u_net(nn.Module):\n",
        "    def __init__(self):\n",
        "      super().__init__()\n",
        "      self.conv1 = nn.Conv2d(3, 64, 3)\n",
        "      self.conv2 = nn.Conv2d(64, 128, 3)\n",
        "      self.conv3 = nn.Conv2d(128, 256, 3)\n",
        "      self.conv4 = nn.Conv2d(256, 512, 3)\n",
        "      self.conv5 = nn.Conv2d(512, 1024, 3)\n",
        "      self.conv6 = nn.Conv2d(1024, 512, 3)\n",
        "      self.conv7 = nn.Conv2d(512, 512, 3)\n",
        "      self.conv8 = nn.Conv2d(512, 256, 3)\n",
        "      self.conv9 = nn.Conv2d(256, 256, 3)\n",
        "      self.conv10 = nn.Conv2d(256, 128, 3)\n",
        "      self.conv11 = nn.Conv2d(128, 128, 3)\n",
        "      self.conv12 = nn.Conv2d(64, 64, 3)\n",
        "      self.b1 = nn.BatchNorm2d(64)\n",
        "      self.b2 = nn.BatchNorm2d(128)\n",
        "      self.b3 = nn.BatchNorm2d(256)\n",
        "      self.b4 = nn.BatchNorm2d(512)\n",
        "      self.b5 = nn.BatchNorm2d(1024)\n",
        "      self.convT1 = nn.ConvTranspose2d(1024, 512, 2, 2)\n",
        "      self.convT2 = nn.ConvTranspose2d(512, 256, 2, 2)\n",
        "      self.convT3 = nn.ConvTranspose2d(256, 256, 2, 2)\n",
        "      self.convT4 = nn.ConvTranspose2d(128, 64, 2, 2)\n",
        "      self.convT5 = nn.ConvTranspose2d(64, 32, 2, 2)\n",
        "      self.pool1 = nn.MaxPool2d(2, 2)\n",
        "\n",
        "    def forward(self, x):\n",
        "      x = F.relu(self.b1(self.conv1(x)))\n",
        "      x = F.relu(self.b1(self.conv12(x)))\n",
        "      x = F.relu(self.b2(self.conv2(x)))\n",
        "      x = self.pool1(x)\n",
        "      x = F.relu(self.b2(self.conv11(x)))\n",
        "      x = F.relu(self.b3(self.conv3(x)))\n",
        "      x1 = x\n",
        "      x1 = x1[:, :, int((58 - 24)/2) : int((58 + 24)/2), int((58 - 24)/2) : int((58 + 24)/2)];\n",
        "      x = self.pool1(x)\n",
        "      x = F.relu(self.b3(self.conv9(x)))\n",
        "      x = F.relu(self.b4(self.conv4(x)))\n",
        "      x2 = x\n",
        "      x2 = x2[:, :, int((25 - 16)/2) : int((25 + 16)/2), int((25 - 16)/2) : int((25 + 16)/2)];\n",
        "      x = self.pool1(x)\n",
        "      x = F.relu(self.b4(self.conv7(x)))\n",
        "      x = F.relu(self.b5(self.conv5(x)))\n",
        "      x = self.b4(self.convT1(x))\n",
        "      x = torch.cat((x2, x), dim = 1)\n",
        "      x = F.relu(self.b4(self.conv6(x)))\n",
        "      x = F.relu(self.b4(self.conv7(x)))\n",
        "      x = self.b3(self.convT2(x))\n",
        "      x = torch.cat((x1, x), dim = 1)\n",
        "      x = F.relu(self.b3(self.conv8(x)))\n",
        "      x = F.relu(self.b3(self.conv9(x)))\n",
        "      x = self.b3(self.convT3(x))\n",
        "      x = F.relu(self.b2(self.conv10(x)))\n",
        "      x = F.relu(self.b2(self.conv11(x)))\n",
        "      x = F.relu(self.b2(self.conv11(x)))\n",
        "      x = F.relu(self.b2(self.conv11(x)))\n",
        "      x = self.b1(self.convT4(x))\n",
        "      x = self.convT5(x)\n",
        "      del x1\n",
        "      del x2\n",
        "      return x\n",
        "net = u_net()\n",
        "print(net)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "u_net(\n",
            "  (conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))\n",
            "  (conv2): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1))\n",
            "  (conv3): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1))\n",
            "  (conv4): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1))\n",
            "  (conv5): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
            "  (conv6): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1))\n",
            "  (conv7): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1))\n",
            "  (conv8): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1))\n",
            "  (conv9): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
            "  (conv10): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1))\n",
            "  (conv11): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))\n",
            "  (conv12): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))\n",
            "  (b1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "  (b2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "  (b3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "  (b4): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "  (b5): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "  (convT1): ConvTranspose2d(1024, 512, kernel_size=(2, 2), stride=(2, 2))\n",
            "  (convT2): ConvTranspose2d(512, 256, kernel_size=(2, 2), stride=(2, 2))\n",
            "  (convT3): ConvTranspose2d(256, 256, kernel_size=(2, 2), stride=(2, 2))\n",
            "  (convT4): ConvTranspose2d(128, 64, kernel_size=(2, 2), stride=(2, 2))\n",
            "  (convT5): ConvTranspose2d(64, 32, kernel_size=(2, 2), stride=(2, 2))\n",
            "  (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
            ")\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p7YgrmSo00eI",
        "colab_type": "code",
        "outputId": "a90ce2c1-762a-4076-9009-237b359fa137",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 454
        }
      },
      "source": [
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "print(device)\n",
        "UNET = u_net()\n",
        "UNET.to(device)\n"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "cuda:0\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "u_net(\n",
              "  (conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))\n",
              "  (conv2): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1))\n",
              "  (conv3): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1))\n",
              "  (conv4): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1))\n",
              "  (conv5): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
              "  (conv6): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1))\n",
              "  (conv7): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1))\n",
              "  (conv8): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1))\n",
              "  (conv9): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
              "  (conv10): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1))\n",
              "  (conv11): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))\n",
              "  (conv12): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))\n",
              "  (b1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "  (b2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "  (b3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "  (b4): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "  (b5): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "  (convT1): ConvTranspose2d(1024, 512, kernel_size=(2, 2), stride=(2, 2))\n",
              "  (convT2): ConvTranspose2d(512, 256, kernel_size=(2, 2), stride=(2, 2))\n",
              "  (convT3): ConvTranspose2d(256, 256, kernel_size=(2, 2), stride=(2, 2))\n",
              "  (convT4): ConvTranspose2d(128, 64, kernel_size=(2, 2), stride=(2, 2))\n",
              "  (convT5): ConvTranspose2d(64, 32, kernel_size=(2, 2), stride=(2, 2))\n",
              "  (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
              ")"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eNEjeffq19yD",
        "colab_type": "text"
      },
      "source": [
        "Defining the loss as crossentropyloss and the other parameters such as learning rate etc. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8M0Lft7i4EQs",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "loss1 = nn.CrossEntropyLoss()\n",
        "optimizer = optim.Adam(UNET.parameters(), lr = 0.0001, betas = (0.9, 0.999), eps = 1e-08, weight_decay=0, amsgrad=False)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KcV429Ux4Mhq",
        "colab_type": "text"
      },
      "source": [
        "Defining the loop for training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3NlexAr26rSD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "for epoch in range(200):\n",
        "  for i, data in enumerate(train_loader):  \n",
        "    inputs, labels = data;\n",
        "    if labels.size() == torch.Size([1, 1, 128, 128]):\n",
        "      labels = labels.reshape(1, 128, 128)\n",
        "    inputs, labels = inputs.to(device), labels.to(device)\n",
        "    optimizer.zero_grad()\n",
        "    outputs = UNET(inputs)\n",
        "    loss = loss1(outputs, labels)\n",
        "    loss.backward()\n",
        "    optimizer.step()\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fEcuFQdXBcST",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "Saving the model to the gdrive."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SDgaHv0mzNst",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "PATH = '/content/gdrive/My Drive/camvidmodel/saved1.pth'\n",
        "torch.save(UNET.state_dict(),PATH); # A state_dict is simply a Python dictionary \n",
        "# object that maps each layer to its parameter tensor.\n"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}