UNet_model_diff_channels.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e6cefc7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from base import BaseModel\n",
    "from torch.optim import *\n",
    "from torchvision import models\n",
    "import numpy as np\n",
    "from abc import abstractmethod\n",
    "import argparse\n",
    "import collections\n",
    "import torch\n",
    "import numpy as np\n",
    "import data_loader.data_loaders as module_data\n",
    "import model.loss as module_loss\n",
    "import model.metric as module_metric\n",
    "import model.model as module_arch\n",
    "from parse_config import ConfigParser\n",
    "from trainer import Trainer\n",
    "from utils import prepare_device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fb8ca19d",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LOG: Using pretrained convolutional layer conv_1 Conv2d(11, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "LOG: Using pretrained convolutional layer conv_1 Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "LOG: Using pretrained convolutional layer conv_1 Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "LOG: Using pretrained convolutional layer conv_2 Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "LOG: Using pretrained convolutional layer conv_1 Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "LOG: Using pretrained convolutional layer conv_2 Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "\n",
      "\n",
      "####################################################################################################\n",
      "(LOG): The following Model Topology will be Utilized: ENC_4_DEC_4\n",
      "####################################################################################################\n",
      "\n",
      "\n",
      "in_tensor.shape is:  torch.Size([4, 11, 64, 64]) out_tensor.shape is:  torch.Size([4, 2, 64, 64])\n"
     ]
    }
   ],
   "source": [
    "class UNet_down_block(BaseModel):\n",
    "    \"\"\"\n",
    "        Encoder class\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, input_channel, output_channel, conv_1=None, conv_2=None):\n",
    "        super(UNet_down_block, self).__init__()\n",
    "        if conv_1:\n",
    "            print('LOG: Using pretrained convolutional layer conv_1', conv_1)\n",
    "        if conv_2:\n",
    "            print('LOG: Using pretrained convolutional layer conv_2', conv_2)\n",
    "        self.input_channels = input_channel\n",
    "        self.output_channels = output_channel\n",
    "        self.conv1 = conv_1 if conv_1 else nn.Conv2d(input_channel, output_channel, kernel_size=3, padding=1)\n",
    "        self.conv2 = conv_2 if conv_2 else nn.Conv2d(output_channel, output_channel, kernel_size=3, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(num_features=output_channel)\n",
    "        self.bn2 = nn.BatchNorm2d(num_features=output_channel)\n",
    "        self.activate = nn.ReLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.activate(self.bn1(self.conv1(x)))\n",
    "        x = self.activate(self.bn2(self.conv2(x)))\n",
    "        return x\n",
    "\n",
    "\n",
    "class UNet_up_block(BaseModel):\n",
    "    \"\"\"\n",
    "        Decoder class\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, prev_channel, input_channel, output_channel):\n",
    "        super(UNet_up_block, self).__init__()\n",
    "        self.output_channels = output_channel\n",
    "        self.tr_conv_1 = nn.ConvTranspose2d(\n",
    "            input_channel, input_channel, kernel_size=2, stride=2)\n",
    "        self.conv_1 = nn.Conv2d(prev_channel+input_channel, output_channel, kernel_size=3, stride=1, padding=1)\n",
    "        self.conv_2 = nn.Conv2d(output_channel, output_channel, kernel_size=3, stride=1, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(num_features=output_channel)\n",
    "        self.bn2 = nn.BatchNorm2d(num_features=output_channel)\n",
    "        self.activate = nn.ReLU()\n",
    "\n",
    "    def forward(self, prev_feature_map, x):\n",
    "        x = self.tr_conv_1(x)\n",
    "        x = self.activate(x)\n",
    "        x = torch.cat((x, prev_feature_map), dim=1)\n",
    "        x = self.activate(self.bn1(self.conv_1(x)))\n",
    "        x = self.activate(self.bn2(self.conv_2(x)))\n",
    "        return x\n",
    "\n",
    "\n",
    "class UNet(BaseModel):\n",
    "    def __init__(self, topology, input_channels, num_classes):\n",
    "        super(UNet, self).__init__()\n",
    "        # these topologies are possible right now\n",
    "        self.topologies = {\n",
    "            \"ENC_1_DEC_1\": self.ENC_1_DEC_1,\n",
    "            \"ENC_2_DEC_2\": self.ENC_2_DEC_2,\n",
    "            \"ENC_3_DEC_3\": self.ENC_3_DEC_3,\n",
    "            \"ENC_4_DEC_4\": self.ENC_4_DEC_4,\n",
    "        }\n",
    "        assert topology in self.topologies\n",
    "        vgg_trained = models.vgg11(pretrained=True)\n",
    "        pretrained_layers = list(vgg_trained.features)\n",
    "        self.max_pool = nn.MaxPool2d(2, 2)\n",
    "        self.dropout = nn.Dropout2d(0.6)\n",
    "        self.activate = nn.ReLU()\n",
    "        #self.encoder_1 = UNet_down_block(input_channels, 64) # this is the original\n",
    "        self.encoder_1 = UNet_down_block(11, 64, conv_1=nn.Conv2d(11, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)))\n",
    "        self.encoder_2 = UNet_down_block(64, 128, conv_1=pretrained_layers[3])\n",
    "        self.encoder_3 = UNet_down_block(128, 256, conv_1=pretrained_layers[6], conv_2=pretrained_layers[8])\n",
    "        self.encoder_4 = UNet_down_block(256, 512, conv_1=pretrained_layers[11], conv_2=pretrained_layers[13])\n",
    "        self.mid_conv_64_64_a = nn.Conv2d(64, 64, 3, padding=1)\n",
    "        self.mid_conv_64_64_b = nn.Conv2d(64, 64, 3, padding=1)\n",
    "        self.mid_conv_128_128_a = nn.Conv2d(128, 128, 3, padding=1)\n",
    "        self.mid_conv_128_128_b = nn.Conv2d(128, 128, 3, padding=1)\n",
    "        self.mid_conv_256_256_a = nn.Conv2d(256, 256, 3, padding=1)\n",
    "        self.mid_conv_256_256_b = nn.Conv2d(256, 256, 3, padding=1)\n",
    "        self.mid_conv_512_1024 = nn.Conv2d(512, 1024, 3, padding=1)\n",
    "        self.mid_conv_1024_1024 = nn.Conv2d(1024, 1024, 3, padding=1)\n",
    "        self.decoder_4 = UNet_up_block(prev_channel=self.encoder_4.output_channels,\n",
    "                                       input_channel=self.mid_conv_1024_1024.out_channels, output_channel=256)\n",
    "        self.decoder_3 = UNet_up_block(prev_channel=self.encoder_3.output_channels,\n",
    "                                       input_channel=self.decoder_4.output_channels, output_channel=128)\n",
    "        self.decoder_2 = UNet_up_block(prev_channel=self.encoder_2.output_channels,\n",
    "                                       input_channel=self.decoder_3.output_channels, output_channel=64)\n",
    "        self.decoder_1 = UNet_up_block(prev_channel=self.encoder_1.output_channels,\n",
    "                                       input_channel=self.decoder_2.output_channels, output_channel=64)\n",
    "        self.binary_last_conv = nn.Conv2d(64, num_classes, kernel_size=1)\n",
    "        self.softmax = nn.Softmax(dim=1)\n",
    "        self.forward = self.topologies[topology]\n",
    "        print('\\n\\n' + \"#\" * 100)\n",
    "        print(\"(LOG): The following Model Topology will be Utilized: {}\".format(self.forward.__name__))\n",
    "        print(\"#\" * 100 + '\\n\\n')\n",
    "\n",
    "    def ENC_1_DEC_1(self, x_in):\n",
    "        x1_cat = self.encoder_1(x_in)\n",
    "        x1_cat_1 = self.dropout(x1_cat)\n",
    "        x1 = self.max_pool(x1_cat_1)\n",
    "        x_mid = self.mid_conv_64_64_a(x1)\n",
    "        x_mid = self.activate(x_mid)\n",
    "        x_mid = self.mid_conv_64_64_b(x_mid)\n",
    "        x_mid = self.activate(x_mid)\n",
    "        x_mid = self.dropout(x_mid)\n",
    "        x = self.decoder_1(x1_cat, x_mid)\n",
    "        x = self.binary_last_conv(x)\n",
    "        # return the final vector and the corresponding softmax-ed prediction\n",
    "        return x, self.softmax(x)\n",
    "\n",
    "    def ENC_2_DEC_2(self, x_in):\n",
    "        x1_cat = self.encoder_1(x_in)\n",
    "        x1 = self.max_pool(x1_cat)\n",
    "        x2_cat = self.encoder_2(x1)\n",
    "        x2_cat_1 = self.dropout(x2_cat)\n",
    "        x2 = self.max_pool(x2_cat_1)\n",
    "        x_mid = self.mid_conv_128_128_a(x2)\n",
    "        x_mid = self.activate(x_mid)\n",
    "        x_mid = self.mid_conv_128_128_b(x_mid)\n",
    "        x_mid = self.activate(x_mid)\n",
    "        x_mid = self.dropout(x_mid)\n",
    "        x = self.decoder_2(x2_cat, x_mid)\n",
    "        x = self.decoder_1(x1_cat, x)\n",
    "        x = self.binary_last_conv(x)\n",
    "        # return the final vector and the corresponding softmax-ed prediction\n",
    "        return x, self.softmax(x)\n",
    "\n",
    "    def ENC_3_DEC_3(self, x_in):\n",
    "        x1_cat = self.encoder_1(x_in)\n",
    "        x1 = self.max_pool(x1_cat)\n",
    "        x2_cat = self.encoder_2(x1)\n",
    "        x2_cat_1 = self.dropout(x2_cat)\n",
    "        x2 = self.max_pool(x2_cat_1)\n",
    "        x3_cat = self.encoder_3(x2)\n",
    "        x3 = self.max_pool(x3_cat)\n",
    "        x_mid = self.mid_conv_256_256_a(x3)\n",
    "        x_mid = self.activate(x_mid)\n",
    "        x_mid = self.mid_conv_256_256_b(x_mid)\n",
    "        x_mid = self.activate(x_mid)\n",
    "        x_mid = self.dropout(x_mid)\n",
    "        x = self.decoder_3(x3_cat, x_mid)\n",
    "        x = self.decoder_2(x2_cat, x)\n",
    "        x = self.decoder_1(x1_cat, x)\n",
    "        x = self.binary_last_conv(x)\n",
    "        # return the final vector and the corresponding softmax-ed prediction\n",
    "        return x, self.softmax(x)\n",
    "\n",
    "    def ENC_4_DEC_4(self, x_in):\n",
    "        x1_cat = self.encoder_1(x_in)\n",
    "        x1 = self.max_pool(x1_cat)\n",
    "        x2_cat = self.encoder_2(x1)\n",
    "        x2_cat_1 = self.dropout(x2_cat)\n",
    "        x2 = self.max_pool(x2_cat_1)\n",
    "        x3_cat = self.encoder_3(x2)\n",
    "        x3 = self.max_pool(x3_cat)\n",
    "        x4_cat = self.encoder_4(x3)\n",
    "        x4_cat_1 = self.dropout(x4_cat)\n",
    "        x4 = self.max_pool(x4_cat_1)\n",
    "        x_mid = self.mid_conv_512_1024(x4)\n",
    "        x_mid = self.activate(x_mid)\n",
    "        x_mid = self.mid_conv_1024_1024(x_mid)\n",
    "        x_mid = self.activate(x_mid)\n",
    "        x_mid = self.dropout(x_mid)\n",
    "        x = self.decoder_4(x4_cat, x_mid)\n",
    "        x = self.decoder_3(x3_cat, x)\n",
    "        x = self.decoder_2(x2_cat, x)\n",
    "        x = self.decoder_1(x1_cat, x)\n",
    "        x = self.binary_last_conv(x)\n",
    "        # return the final vector and the corresponding softmax-ed prediction\n",
    "        return x, self.softmax(x)\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def check_model(topology, input_channels, num_classes, input_shape):\n",
    "    model = UNet(topology=topology, input_channels=input_channels,num_classes=num_classes)\n",
    "    model.eval()\n",
    "    in_tensor = torch.Tensor(*input_shape)\n",
    "    with torch.no_grad():\n",
    "        out_tensor, softmaxed = model(in_tensor)\n",
    "        print(\"in_tensor.shape is: \", in_tensor.shape, \"out_tensor.shape is: \", out_tensor.shape)\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    \n",
    "    check_model(topology=\"ENC_4_DEC_4\", input_channels=11, # input_channels was 7..i am changing it here to see how it works\n",
    "                num_classes=2, input_shape=[4, 11, 64, 64]) # input_shape was 4,7,64,64 so changing to 18 here also"
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}