Architecture.ipynb

{
 "metadata": {
  "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.8.5-final"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#System Libs\n",
    "import os\n",
    "import glob\n",
    "import json\n",
    "\n",
    "#ImageProc Libs\n",
    "import cv2\n",
    "from albumentations import (Compose,Resize)\n",
    "from albumentations.pytorch import ToTensorV2\n",
    "\n",
    "#Viz\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "#DL Libs\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset,DataLoader\n",
    "from torchsummary import  summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "files = open('configs.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "ff = json.load(files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "images_files = glob.glob(\"/home/lustbeast/PaperWork/Veena Ma'am/Sketch to Face/IIITD_SketchDatabase/Semi-forensic database/IIIT-D student and staff/photo/*.jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class sketchDataset(Dataset):\n",
    "    def __init__(self,paths,transforms=None):\n",
    "        self.paths = paths\n",
    "        self.transforms = transforms\n",
    "    def __len__(self):\n",
    "        return len(self.paths)\n",
    "    def __getitem__(self,idx):\n",
    "        img = cv2.imread(self.paths[idx])\n",
    "        if self.transforms:\n",
    "            img = self.transforms(image=img)['image']\n",
    "\n",
    "        return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def augs():\n",
    "    return Compose([\n",
    "        Resize(384,384),\n",
    "        ToTensorV2()\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = sketchDataset(images_files,augs())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = DataLoader(dataset,num_workers = 4,batch_size=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = next(iter(dataset)).half()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "writer = SummaryWriter(\"runs/graph_viz/encoder/redefined\")"
   ]
  },
  {
   "source": [
    "## Architecture\n",
    "\n",
    "### Modules:\n",
    "    1. MultiScale-A\n",
    "    2. Reduction\n",
    "    3. MultiScale-B\n",
    "\n",
    "<img src=\"RES/arch.png\"></img>\n",
    "\n",
    "\n",
    "Things to be considered:\n",
    "\n",
    "    * In MultiScale-A module, The stream ID =1 given in the paper has been neglected due to average pooling. So, this module just contains three parallel streams.\n",
    "    * In MultiScale-B module, The stream ID =4 given in the paper has been neglected due to the same reason. So, this module also contains just three parallel streams"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiScaleA(nn.Module):\n",
    "    def __init__(self,in_channels,stream_1_1x1_filters,stream_2_1x1_filters,stream_2_3x3_filters,stream_3_1x1_filters,stream_3_3x3_filters):\n",
    "        super(MultiScaleA,self).__init__()\n",
    "        \n",
    "        self.stream1 = conv_mod(in_channels,stream_1_1x1_filters,kernel_size=(1,1)) #Stream ID=1\n",
    "        self.stream2 = nn.Sequential(\n",
    "            conv_mod(in_channels,stream_2_1x1_filters,kernel_size=(1,1)),\n",
    "            conv_mod(stream_2_1x1_filters,stream_2_3x3_filters,kernel_size=(3,3),padding=1),     #Stream ID=2\n",
    "        )\n",
    "        self.stream3 = nn.Sequential(\n",
    "            conv_mod(in_channels,stream_3_1x1_filters,kernel_size=(1,1)),\n",
    "            conv_mod(stream_3_1x1_filters,stream_3_3x3_filters[0],kernel_size=(3,3),padding=1),\n",
    "            conv_mod(stream_3_3x3_filters[0],stream_3_3x3_filters[1],kernel_size=(3,3),padding=1)\n",
    "        )\n",
    "\n",
    "    def forward(self,x):\n",
    "\n",
    "        stream1 = self.stream1(x)\n",
    "        stream2 = self.stream2(x)\n",
    "        stream3 = self.stream3(x)\n",
    "\n",
    "        concat = torch.cat([stream1,stream2,stream3],axis=1)\n",
    "\n",
    "        return concat\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Reduction(nn.Module):\n",
    "    def __init__(self,in_channels,red_stream_2_3x3_filters,red_stream_3_1x1_filters,red_stream_3_3x3_filters):\n",
    "\n",
    "        super(Reduction,self).__init__()\n",
    "        self.stream1_MF = nn.MaxPool2d(kernel_size=(3,3),stride=2)\n",
    "        \n",
    "        self.stream2_CF = conv_mod(in_channels,red_stream_2_3x3_filters,kernel_size=(3,3),stride=(2,2))\n",
    "        \n",
    "        self.stream3_CF = nn.Sequential(\n",
    "            conv_mod(in_channels,red_stream_3_1x1_filters,kernel_size=(1,1)),\n",
    "            conv_mod(red_stream_3_1x1_filters,red_stream_3_3x3_filters[0],kernel_size=(3,3)),\n",
    "            conv_mod(red_stream_3_3x3_filters[0],red_stream_3_3x3_filters[1],kernel_size=(3,3),stride=(2,2),padding=1)\n",
    "        )\n",
    "        \n",
    "\n",
    "    def forward(self,x):\n",
    "\n",
    "        stream1_MF = self.stream1_MF(x)\n",
    "        stream2_CF = self.stream2_CF(x)\n",
    "        stream3_CF = self.stream3_CF(x)\n",
    "\n",
    "        print(stream1_MF.shape,stream2_CF.shape,stream3_CF.shape)\n",
    "\n",
    "        return torch.cat([stream1_MF,stream2_CF,stream3_CF],axis=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiScaleB(nn.Module):\n",
    "    def __init__(self,in_channels,Bstream_1_1x1,Bstream_2_1x1,Bstream_2_3x3,Bstream_3_1x1,Bstream_3_3x3):\n",
    "        super(MultiScaleB,self).__init__()\n",
    "        self.st1 = conv_mod(in_channels,Bstream_1_1x1,kernel_size=(1,1))\n",
    "\n",
    "        self.st2 = nn.Sequential(\n",
    "            conv_mod(in_channels,Bstream_2_1x1,kernel_size=(1,1)),\n",
    "            conv_mod(Bstream_2_1x1,Bstream_2_3x3[0],kernel_size=(1,3)),\n",
    "            conv_mod(Bstream_2_3x3[0],Bstream_2_3x3[1],kernel_size=(3,1),padding=(1,1))\n",
    "            )\n",
    "\n",
    "        self.st3 = nn.Sequential(\n",
    "            conv_mod(in_channels,Bstream_3_1x1,kernel_size=(1,1)),\n",
    "            conv_mod(Bstream_3_1x1,Bstream_3_3x3[0],kernel_size=(1,3)),\n",
    "            conv_mod(Bstream_3_3x3[0],Bstream_3_3x3[1],kernel_size=(3,1)),\n",
    "            conv_mod(Bstream_3_3x3[1],Bstream_3_3x3[2],kernel_size=(1,3)),\n",
    "            conv_mod(Bstream_3_3x3[2],Bstream_3_3x3[3],kernel_size=(3,1),padding=(2,2))\n",
    "        )\n",
    "\n",
    "\n",
    "    def forward(self,x):\n",
    "\n",
    "        st1 = self.st1(x)\n",
    "        st2 = self.st2(x)\n",
    "        st3 = self.st3(x)\n",
    "\n",
    "        print(st1.shape,st2.shape,st3.shape)\n",
    "\n",
    "        return torch.cat([st1,st2,st3],axis=1)\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class conv_mod(nn.Module):\n",
    "    def __init__(self,in_channels,out_channels,kernel_size=(1,1),stride=1,padding=0,activation='relu'):\n",
    "        super(conv_mod,self).__init__()\n",
    "        self.mod = nn.Sequential(\n",
    "            nn.BatchNorm2d(in_channels),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(in_channels,out_channels,kernel_size=kernel_size,stride=stride,padding=padding)\n",
    "\n",
    "        )\n",
    "\n",
    "    def forward(self,x):\n",
    "        x = self.mod(x)\n",
    "        return x"
   ]
  },
  {
   "source": [
    "## Architecture - Encoder"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Multi_Block(nn.Module):\n",
    "    def __init__(self,blocks):\n",
    "        super(Multi_Block,self).__init__()\n",
    "        self.BlockA_params = [i for i in blocks[0].values()]\n",
    "        self.BlockB_params = [j for j in blocks[1].values()]\n",
    "        self.BlockC_params = [k for k in blocks[2].values()]\n",
    "\n",
    "        self.mulA = MultiScaleA(self.BlockA_params[0],self.BlockA_params[1],self.BlockA_params[2],self.BlockA_params[3],self.BlockA_params[4],self.BlockA_params[5])\n",
    "        self.red = Reduction(self.BlockB_params[0],self.BlockB_params[1],self.BlockB_params[2],self.BlockB_params[3])\n",
    "\n",
    "        self.mulB = MultiScaleB(self.BlockC_params[0],self.BlockC_params[1],self.BlockC_params[2],self.BlockC_params[3],self.BlockC_params[4],self.BlockC_params[5])\n",
    "\n",
    "    def forward(self,x):\n",
    "\n",
    "        x = self.mulA(x)\n",
    "        x = self.red(x)\n",
    "        x = self.mulB(x)\n",
    "\n",
    "        return x\n",
    "class self_attention(nn.Module):\n",
    "    def __init__(self,feature):\n",
    "        super(self_attention,self).__init__()\n",
    "        self.f = nn.Conv2d(feature.shape[1],feature.shape[1],kernel_size=(1,1))\n",
    "        self.g = nn.Conv2d(feature.shape[1],feature.shape[1],kernel_size=(1,1))\n",
    "        self.h = nn.Conv2d(feature.shape[1],feature.shape[1],kernel_size=(1,1))\n",
    "        self.soft = nn.Softmax()\n",
    "    def forward(self,x):\n",
    "        f = self.f(x)\n",
    "        g = self.g(x)\n",
    "        h = self.h(x)\n",
    "        fg = torch.matmul(f,g)\n",
    "        fg_soft = self.soft(fg)\n",
    "        return torch.matmul(fg_soft,h)\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self,params):\n",
    "        super(Encoder,self).__init__()\n",
    "        self.A = Multi_Block(params['BlockA'])\n",
    "        self.B = Multi_Block(params['BlockB'])\n",
    "        self.C = Multi_Block(params['BlockC'])\n",
    "        self.sing_1 = conv_mod(416,256,kernel_size=(5,5))\n",
    "        self.sing_2 = conv_mod(256,128,kernel_size=(5,5))\n",
    "        self.sing_3 = conv_mod(128,64,kernel_size=(5,5))\n",
    "        self.sing_4 = conv_mod(64,32,kernel_size=(5,5))\n",
    "       \n",
    "\n",
    "    def forward(self,x):\n",
    "        multi_x_1 = self.A(x)\n",
    "        multi_x_2 = self.B(multi_x_1)\n",
    "        multi_x_3 = self.C(multi_x_2)\n",
    "        single_x_1 = self.sing_1(multi_x_3)\n",
    "        single_x_2 = self.sing_2(single_x_1)\n",
    "        single_x_3 = self.sing_3(single_x_2)\n",
    "        single_x_4 = self.sing_4(single_x_3)\n",
    "        \n",
    "        return multi_x_1,multi_x_2,multi_x_3,single_x_1,single_x_2,single_x_3,single_x_4"
   ]
  },
  {
   "source": [
    "## Architecture - Decoder"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Decoder,self).__init__()\n",
    "        self.up1 = nn.ConvTranspose2d(32,64,kernel_size=(5,5))\n",
    "        self.up2 = nn.ConvTranspose2d(64,128,kernel_size=(5,5))\n",
    "        self.up3 = nn.ConvTranspose2d(128,256,kernel_size=(5,5))\n",
    "        self.up4 = nn.ConvTranspose2d(256,416,kernel_size=(5,5))\n",
    "    def forward(self,x):\n",
    "        x = self.up1(x)\n",
    "        x = self.up2(x)\n",
    "        x = self.up3(x)\n",
    "        x = self.up4(x)\n",
    "\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "enc = Encoder(params=ff).half().cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = torch.randn(1,3,384,384)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "torch.Size([1, 224, 191, 191]) torch.Size([1, 32, 191, 191]) torch.Size([1, 256, 191, 191])\ntorch.Size([1, 32, 191, 191]) torch.Size([1, 128, 191, 191]) torch.Size([1, 256, 191, 191])\ntorch.Size([1, 256, 95, 95]) torch.Size([1, 32, 95, 95]) torch.Size([1, 256, 95, 95])\ntorch.Size([1, 32, 95, 95]) torch.Size([1, 128, 95, 95]) torch.Size([1, 256, 95, 95])\ntorch.Size([1, 256, 47, 47]) torch.Size([1, 32, 47, 47]) torch.Size([1, 256, 47, 47])\ntorch.Size([1, 32, 47, 47]) torch.Size([1, 128, 47, 47]) torch.Size([1, 256, 47, 47])\n"
     ]
    }
   ],
   "source": [
    "enc_out = enc(img.half().to(\"cuda:0\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "dec  = Decoder().to('cuda:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "dec_out = dec(enc_out[6].to(\"cuda:0\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "torch.Size([1, 416, 47, 47])"
      ]
     },
     "metadata": {},
     "execution_count": 34
    }
   ],
   "source": [
    "dec_out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decb = DecoderB(416,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_layers(model, layer_type_old, layer_type_new, convert_weights=False, num_groups=None):\n",
    "    for name, module in reversed(model._modules.items()):\n",
    "        if len(list(module.children())) > 0:\n",
    "            # recurse\n",
    "            model._modules[name] = convert_layers(module, layer_type_old, layer_type_new, convert_weights)\n",
    "\n",
    "        if type(module) == layer_type_old:\n",
    "            layer_old = module\n",
    "            layer_new = layer_type_new(module.num_features if num_groups is None else num_groups, module.num_features, module.eps, module.affine) \n",
    "\n",
    "            if convert_weights:\n",
    "                layer_new.weight = layer_old.weight\n",
    "                layer_new.bias = layer_old.bias\n",
    "\n",
    "            model._modules[name] = layer_new\n",
    "\n",
    "    return model\n",
    "\n",
    "\n",
    "# Replace BatchNorm with GroupNorm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "output_type": "error",
     "ename": "ModuleAttributeError",
     "evalue": "'conv_mod' object has no attribute 'num_features'",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleAttributeError\u001b[0m                      Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-46-1b0926d5bc9e>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mconvert_layers\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0menc\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mconv_mod\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mConvTranspose2d\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-44-980bed5aefcb>\u001b[0m in \u001b[0;36mconvert_layers\u001b[0;34m(model, layer_type_old, layer_type_new, convert_weights, num_groups)\u001b[0m\n\u001b[1;32m      7\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodule\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mlayer_type_old\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m             \u001b[0mlayer_old\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 9\u001b[0;31m             \u001b[0mlayer_new\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlayer_type_new\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum_features\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mnum_groups\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mnum_groups\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum_features\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meps\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maffine\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     10\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mconvert_weights\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__getattr__\u001b[0;34m(self, name)\u001b[0m\n\u001b[1;32m    776\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mname\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mmodules\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    777\u001b[0m                 \u001b[0;32mreturn\u001b[0m \u001b[0mmodules\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 778\u001b[0;31m         raise ModuleAttributeError(\"'{}' object has no attribute '{}'\".format(\n\u001b[0m\u001b[1;32m    779\u001b[0m             type(self).__name__, name))\n\u001b[1;32m    780\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mModuleAttributeError\u001b[0m: 'conv_mod' object has no attribute 'num_features'"
     ]
    }
   ],
   "source": [
    "convert_layers(enc,conv_mod,torch.nn.ConvTranspose2d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ]
}