BYOL_Pytorch (1).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "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.7.2"
    },
    "colab": {
      "name": "BYOL-Pytorch.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DEsLktBKRtLF",
        "colab_type": "text"
      },
      "source": [
        "# BYOL-Pytorch  \n",
        "Pytorch Implementation of BYOL: Bootstrap Your Own Latent: A New Approach to Self-Supervised Learning (https://arxiv.org/abs/2006.07733).   \n",
        "Major part of Code is inspired from https://github.com/sthalles/PyTorch-BYOL.  \n",
        "The Code has more appropriate Naming Convention. \n",
        "# Default Training\n",
        "* Running the Python File without any changes trains BYOL with **CIFAR10** Dataset.\n",
        "* All the Parameters are contained in ___Params Object___ in the script.\n",
        "# Custom Training\n",
        "* Change the __Dataset Object__.\n",
        "* Update the Required Parameters in the ___Params Object___.   \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Iguz4yzHRuOI",
        "colab_type": "text"
      },
      "source": [
        "## Import Statement"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JXwnvO2WFr5N",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.utils.data.dataloader\n",
        "from torchvision import datasets\n",
        "from torchvision.transforms import transforms\n",
        "import torchvision.models\n",
        "import cv2\n",
        "import numpy as np\n",
        "from tqdm import tqdm"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gcTb54y7Fr5U",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "3708d645-f203-47a1-cf03-c48b989f87f6"
      },
      "source": [
        "np.random.seed(0)\n",
        "torch.manual_seed(42)"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<torch._C.Generator at 0x7f42bbec33b0>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OjicKwA5IS4a",
        "colab_type": "text"
      },
      "source": [
        "# Augmentation Functions    \n",
        "Given an image, augmentation is applied to create two different views. Augmentation used over here is very similar to that of [SimCLR](https://arxiv.org/abs/2002.05709). \n",
        "### GaussianBlur(Class)   \n",
        "![alt text](https://wikimedia.org/api/rest_v1/media/math/render/svg/dd16b16869269dba008d19c0969515a1d50b3ae2)\n",
        "*   **Parameter** = Kernel Size\n",
        "*   **Output**     = Gaussian Blur Transformed Image   \n",
        "### Transforms(Function)\n",
        "*   **Parameter** = Input Dimension of the Image.\n",
        "*   **Output**     = Composes a torchvision.transforms Object with all the Transformation functions intact.\n",
        "### MultiViewDataInjector(Class)\n",
        "*   **Parameter** = Input Image.\n",
        "*   **Output**    = Applies **Transforms** Function to result two different augmented image.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O1ws4HuZFr5d",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "class GaussianBlur(object):\n",
        "    \"\"\"blur a single image on CPU\"\"\"\n",
        "\n",
        "    def __init__(self, kernel_size):\n",
        "        radias = kernel_size // 2\n",
        "        kernel_size = radias * 2 + 1\n",
        "        self.blur_h = nn.Conv2d(3, 3, kernel_size=(kernel_size, 1),\n",
        "                                stride=1, padding=0, bias=False, groups=3)\n",
        "        self.blur_v = nn.Conv2d(3, 3, kernel_size=(1, kernel_size),\n",
        "                                stride=1, padding=0, bias=False, groups=3)\n",
        "        self.k = kernel_size\n",
        "        self.r = radias\n",
        "\n",
        "        self.blur = nn.Sequential(\n",
        "            nn.ReflectionPad2d(radias),\n",
        "            self.blur_h,\n",
        "            self.blur_v\n",
        "        )\n",
        "\n",
        "        self.pil_to_tensor = transforms.ToTensor()\n",
        "        self.tensor_to_pil = transforms.ToPILImage()\n",
        "\n",
        "    def __call__(self, img):\n",
        "        img = self.pil_to_tensor(img).unsqueeze(0)\n",
        "\n",
        "        sigma = np.random.uniform(0.1, 2.0)\n",
        "        x = np.arange(-self.r, self.r + 1)\n",
        "        x = np.exp(-np.power(x, 2) / (2 * sigma * sigma))\n",
        "        x = x / x.sum()\n",
        "        x = torch.from_numpy(x).view(1, -1).repeat(3, 1)\n",
        "\n",
        "        self.blur_h.weight.data.copy_(x.view(3, 1, self.k, 1))\n",
        "        self.blur_v.weight.data.copy_(x.view(3, 1, 1, self.k))\n",
        "\n",
        "        with torch.no_grad():\n",
        "            img = self.blur(img)\n",
        "            img = img.squeeze()\n",
        "\n",
        "        img = self.tensor_to_pil(img)\n",
        "\n",
        "        return img"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zJYmcyhVFr5i",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def Transforms(Input_Dim,S=1):\n",
        "    Color_Jitter = transforms.ColorJitter(0.8*S,0.8*S,0.8*S,0.2*S)\n",
        "    Data_Transforms = transforms.Compose([transforms.RandomResizedCrop(size=Input_Dim[0]),\n",
        "                                         transforms.RandomHorizontalFlip(),\n",
        "                                         transforms.RandomApply([Color_Jitter],p=0.75),\n",
        "                                         transforms.RandomGrayscale(p=0.2),\n",
        "                                         GaussianBlur(int(0.1*Input_Dim[0])),\n",
        "                                         transforms.ToTensor(),\n",
        "                                        ])\n",
        "    return Data_Transforms"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "30KwuU7oFr5m",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class MultiViewDataInjector(object):\n",
        "    def __init__(self,Transforms):\n",
        "        self.transforms = Transforms\n",
        "    def __call__(self,Sample,*Consistent_Flip):\n",
        "        if Consistent_Flip:\n",
        "            Sample  =  torchvision.transforms.RandomHorizontalFlip()\n",
        "        Output = [transforms(Sample) for transforms in self.transforms]\n",
        "        return Output"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BDHuM4ssFr5q",
        "colab_type": "text"
      },
      "source": [
        "# Model   \n",
        "Contains two basic Neural Networks \n",
        "*  **MLP_Base** - Creates the **Latent Space** from the Encoder.\n",
        "*  **Skeleton Net** - Encompases MLP_BASE for **Latent Space** creation and uses **ResNet18** to learn Feature Representations.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LSCHcyc6Fr5r",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class MLP_Base(nn.Module):\n",
        "    def __init__(self,Inp,Hidden,Projection):\n",
        "        super(MLP_Base,self).__init__()\n",
        "        self.Linear1 = nn.Linear(Inp,Hidden)\n",
        "        self.BatchNorm = nn.BatchNorm1d(Hidden)\n",
        "        self.Linear2 = nn.Linear(Hidden,Projection)\n",
        "    def forward(self,Input):\n",
        "        Linear_Inp = torch.relu(self.BatchNorm(self.Linear1(Input)))\n",
        "        Linear_Out = self.Linear2(Linear_Inp)\n",
        "        return Linear_Out"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3R2korLYFr5x",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class SkeletonNet(nn.Module):\n",
        "    def __init__(self,Hid,Proj):\n",
        "        super(SkeletonNet,self).__init__()\n",
        "        Resnet = torchvision.models.resnet18(pretrained=False)\n",
        "        self.Encoder = torch.nn.Sequential(*list(Resnet.children())[:-1])\n",
        "        self.Proj = MLP_Base(Resnet.fc.in_features,Hid,Proj)\n",
        "    def forward(self,Input):\n",
        "        Enc_Out = self.Encoder(Input)\n",
        "        Enc_Out = Enc_Out.view(Enc_Out.size(0),Enc_Out.size(1))\n",
        "        Final = self.Proj(Enc_Out)\n",
        "        return Final"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aCEF9jV8Fr53",
        "colab_type": "text"
      },
      "source": [
        "# Training Class"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yANwkv3lFr53",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class BYOL:\n",
        "    def __init__(self,Online_Net,Target_Net,Predictor,Optim,Params):\n",
        "        self.Online_Net = Online_Net\n",
        "        self.Target_Net = Target_Net\n",
        "        self.Predictor  = Predictor\n",
        "        self.Optim      = Optim\n",
        "        self.Device     = Params['Device']\n",
        "        self.Epochs     = Params['Epochs']\n",
        "        self.Moment        = Params['M']\n",
        "        self.Batch_Size = Params['Batch_Size']\n",
        "        self.Save_Path = 'Models/BYOL.pth'\n",
        "    @torch.no_grad()\n",
        "    def Update_Target_Params(self):\n",
        "        for Param_Online,Param_Target in zip(self.Online_Net.parameters(),self.Target_Net.parameters()):\n",
        "            Param_Target = Param_Target.data *self.Moment + Param_Online.data*(1-self.Moment)\n",
        "    @staticmethod          \n",
        "    def Loss(Rep1,Rep2):\n",
        "        Norm_Rep1 = F.normalize(Rep1,dim=-1,p=2) #L2-Normalized Rep One\n",
        "        Norm_Rep2 = F.normalize(Rep2,dim=-1,p=2) #L2 Normalized Rep Two\n",
        "        Loss = -2 * (Norm_Rep1*Norm_Rep2).sum(dim=-1)\n",
        "        return Loss \n",
        "    def Init_Target_Network(self):\n",
        "        for Param_Online,Param_Target in zip(self.Online_Net.parameters(),self.Target_Net.parameters()):\n",
        "            Param_Target.data.copy_(Param_Online.data) #Init Target with Param_Online\n",
        "            Param_Target.requires_grad = False\n",
        "    def TrainLoop(self,View1,View2):\n",
        "        self.Optim.zero_grad()\n",
        "        Pred1 = self.Predictor(self.Online_Net(View1))\n",
        "        Pred2 = self.Predictor(self.Online_Net(View2))\n",
        "        with torch.no_grad():\n",
        "            Target2 = self.Target_Net(View1)\n",
        "            Target1 = self.Target_Net(View2)\n",
        "        Loss_Calc = self.Loss(Pred1,Target1) + self.Loss(Pred2,Target2)\n",
        "        return Loss_Calc.mean()\n",
        "    def Train(self,Trainset):\n",
        "        TrainLoader = torch.utils.data.DataLoader(Trainset,batch_size=self.Batch_Size,drop_last=False,shuffle=True)\n",
        "        self.Init_Target_Network()\n",
        "        for Epoch in range(self.Epochs):\n",
        "          Loss_Count = 0.0\n",
        "          print(\"Epoch {}\".format(Epoch))\n",
        "          for (View_1,View_2),_ in tqdm(TrainLoader):\n",
        "              View_1 = View_1.to(self.Device)\n",
        "              View_2 = View_2.to(self.Device)\n",
        "              Loss = self.TrainLoop(View_1,View_2)\n",
        "              Loss_Count += Loss.item()\n",
        "              Loss.backward()\n",
        "              self.Optim.step()\n",
        "              self.Update_Target_Params()\n",
        "          Epoch_Loss = Loss_Count/len(TrainLoader)\n",
        "          print(\"\\n Epoch {} Loss:{} : \".format(Epoch,Epoch_Loss))\n",
        "        self.Save(self.Save_Path)\n",
        "    def Save(self,Save):\n",
        "        torch.save({'Online_Net':self.Online_Net.state_dict(),\n",
        "                    'Enc_Net':self.Online_Net.Encoder.state_dict(),\n",
        "                    'Target_Net':self.Target_Net.state_dict(),\n",
        "                    'Optim':self.Optim.state_dict()},Save)\n"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SMEfYm3hFr58",
        "colab_type": "text"
      },
      "source": [
        "# Main Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yV-8hgamFr58",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "Parameters = {'Epochs':50,'M':0.99,'Batch_Size':64,'Device':'cuda','Hidden':512,'Proj':128,'LR':0.03}"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GHGyBW_QFr6A",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "3e6dfb91-ef0b-4fdb-cd59-534eb044cb11"
      },
      "source": [
        "Data_Transforms = Transforms((3,32,32))\n",
        "Dataset = datasets.CIFAR10('./data',download=True,transform=MultiViewDataInjector([Data_Transforms,Data_Transforms]))"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Files already downloaded and verified\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wd4ia_ilFr6E",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "Online_Network = SkeletonNet(Parameters['Hidden'],Parameters['Proj'])\n",
        "Predictor = MLP_Base(Online_Network.Proj.Linear2.out_features,Parameters['Hidden'],Parameters['Proj'])\n",
        "Target_Network = SkeletonNet(Parameters['Hidden'],Parameters['Proj'])"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "otHq9L2EFr6K",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "2ab41072-5d94-4958-a5da-1b5bb1c9b48d"
      },
      "source": [
        "Online_Network.to(Parameters['Device'])\n",
        "Predictor.to(Parameters['Device'])\n",
        "Target_Network.to(Parameters['Device'])\n",
        "print(\"Models Made.\")"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Models Made.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yrNO-ksrFr6P",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "Optimizer = torch.optim.SGD(list(Online_Network.parameters())+list(Predictor.parameters()),lr=0.03)"
      ],
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yTCxCJZjFr6T",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "Trainer = BYOL(Online_Network,Target_Network,Predictor,Optimizer,Parameters)"
      ],
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fFaoH14dFr6a",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 118
        },
        "outputId": "a98b7a02-8214-4463-f13f-ba1d5823e7eb"
      },
      "source": [
        "Trainer.Train(Dataset)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\r  0%|          | 0/782 [00:00<?, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch 0\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 782/782 [02:38<00:00,  4.93it/s]\n",
            "  0%|          | 0/782 [00:00<?, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            " Epoch 0 Loss:-2.237157367731985 : \n",
            "Epoch 1\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            " 73%|███████▎  | 568/782 [01:54<00:42,  5.04it/s]"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ioq4CAvRHjgY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}