cnn_architectures.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "cnn_architectures.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyOrPkUMY3CNnjbWJIMPaPLV",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Cranjis-McB/CNN_ARCHITECTURES/blob/main/cnn_architectures.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# CNN ARCHITECTURES\n",
        "\n",
        "This notebook contains the **Popular CNN Architecures** implemented from the Scratch in **PyTorch** framework.\n",
        "\n",
        "### **Implemented Architectures**:\n",
        "\n",
        "\n",
        "\n",
        "*   **ResNet** (Version - 18, 34, 50, 101, 152)\n",
        "*   **EfficientNet** (Version - b0-b7)\n",
        "*   **EfficientNet-V2** (Version - S, M, L)\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "o4PqNuf6XC8o"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "je24Dkk_W7GP"
      },
      "outputs": [],
      "source": [
        "# Import useful Modules.\n",
        "import torch\n",
        "from torch import nn\n",
        "from math import ceil"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Resnet"
      ],
      "metadata": {
        "id": "jScjA4bWX88O"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Tier-2 ResBlock class. (Resnet-18,Resnet-34)\n",
        "class ResBlock_Tier2(nn.Module):\n",
        "    \n",
        "    def __init__(self, in_channels, intermediate_channels):\n",
        "        super(ResBlock_Tier2, self).__init__()\n",
        "        \n",
        "        # Layers\n",
        "        # To make the residual input equal size as output channel.\n",
        "        \n",
        "        downsample = False\n",
        "        self.skip_connection = nn.Sequential() # Default\n",
        "        \n",
        "        if intermediate_channels == 2*in_channels:\n",
        "            self.skip_connection = nn.Sequential(nn.Conv2d(in_channels, intermediate_channels, kernel_size=1, stride=2, bias = False),\n",
        "                                             nn.BatchNorm2d(intermediate_channels)\n",
        "                                            )\n",
        "            downsample = True\n",
        "            \n",
        "\n",
        "        \n",
        "        # Downsampling the output shape.\n",
        "        if downsample:\n",
        "            self.conv1 = nn.Conv2d(in_channels, intermediate_channels, kernel_size=3, stride=2, padding = 1, bias = False)\n",
        "            \n",
        "        else:\n",
        "            self.conv1 = nn.Conv2d(in_channels, intermediate_channels, kernel_size=3, stride=1, padding = 1, bias = False)\n",
        "            \n",
        "        self.bn1 = nn.BatchNorm2d(intermediate_channels)\n",
        "        \n",
        "        self.conv2 = nn.Conv2d(intermediate_channels, intermediate_channels, kernel_size=3, stride=1, padding = 1, bias = False)\n",
        "        self.bn2 = nn.BatchNorm2d(intermediate_channels)\n",
        "        \n",
        "        self.relu = nn.ReLU()\n",
        "        \n",
        "    def forward(self, x):\n",
        "        \n",
        "        # Residual to be added later.\n",
        "        identity = self.skip_connection(x)\n",
        "        \n",
        "        #---------------------------\n",
        "        # Layer-1\n",
        "        x = self.conv1(x)\n",
        "        x = self.bn1(x)\n",
        "        x = self.relu(x)\n",
        "        #----------------------------\n",
        "        \n",
        "        #---------------------------\n",
        "        # Layer-2\n",
        "        x = self.conv2(x)\n",
        "        x = self.bn2(x)\n",
        "        #print(x.size())\n",
        "        #print(identity.size())\n",
        "        x += identity\n",
        "        x = self.relu(x)\n",
        "        #----------------------------\n",
        "        \n",
        "        return x\n",
        "    \n",
        "#---------------------------------------------------------------------------------------------------\n",
        "'''\n",
        "  This is Tier-2 Resnet class (Resnet-18, 34)\n",
        "  It takes input as img_channels:\n",
        "  for RGB = 3\n",
        "  for Gray = 1),\n",
        "  \n",
        "  num_layers: \n",
        "  for Resnet-18 = [2, 2, 2, 2]\n",
        "  for Resnet-34 = [3, 4, 6, 3]\n",
        "  \n",
        "  , and number of classes. \n",
        "  for ex : Imagenet = 1000, MNIST = 10 etc\n",
        "  \n",
        "  Output : Resnet Model\n",
        "  \n",
        "'''\n",
        "class ResNet_Tier2(nn.Module):\n",
        "    \n",
        "    def __init__(self, img_channels, num_layers, num_classes):\n",
        "        super(ResNet_Tier2, self).__init__()\n",
        "        \n",
        "        # Layers \n",
        "        # Layer-0 Output shape : 64 X 56 X 56\n",
        "        self.layer0 = nn.Sequential(\n",
        "            nn.Conv2d(img_channels, 64, kernel_size=7, stride=2, padding=3),\n",
        "            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),\n",
        "            nn.BatchNorm2d(64),\n",
        "            nn.ReLU()\n",
        "        )\n",
        "        \n",
        "        # Residual Blocks layers\n",
        "        self.layer1 = self._make_layer(ResBlock_Tier2, num_layers[0], 64, 64)\n",
        "        self.layer2 = self._make_layer(ResBlock_Tier2, num_layers[1], 64, 128)\n",
        "        self.layer3 = self._make_layer(ResBlock_Tier2, num_layers[2], 128, 256)\n",
        "        self.layer4 = self._make_layer(ResBlock_Tier2, num_layers[3], 256, 512)\n",
        "        \n",
        "        # FC Layers\n",
        "        self.avgpool = nn.AdaptiveAvgPool2d((1,1))\n",
        "        self.fc  = nn.Linear(512, num_classes)\n",
        "        \n",
        "    def forward(self, x):\n",
        "        x = self.layer0(x)\n",
        "        x = self.layer1(x)\n",
        "        x = self.layer2(x)\n",
        "        x = self.layer3(x)\n",
        "        x = self.layer4(x)\n",
        "        x = self.avgpool(x)\n",
        "        x = x.reshape(x.shape[0], -1)\n",
        "        x = self.fc(x)\n",
        "        \n",
        "        return x\n",
        "        \n",
        "    def _make_layer(self, ResBlock_Tier2, num_residual_blocks, in_channels, intermediate_channels):\n",
        "        \n",
        "        layers = []\n",
        "        only_once = True\n",
        "        \n",
        "        for i in range(num_residual_blocks):\n",
        "            layers.append(ResBlock_Tier2(in_channels, intermediate_channels))\n",
        "            if only_once:\n",
        "                in_channels = intermediate_channels\n",
        "                only_once = False\n",
        "            \n",
        "        return nn.Sequential(*layers)"
      ],
      "metadata": {
        "id": "Xou9BV_sX6iG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Tier-1 ResBlock class. (Resnet-50, Resnet-101, Resnet-152)\n",
        "class ResBlock_Tier1(nn.Module):\n",
        "    \n",
        "    def __init__(self, in_channels, intermediate_channels):\n",
        "        super(ResBlock_Tier1, self).__init__()\n",
        "        \n",
        "        # Layers\n",
        "        # To make the residual input equal size as output channel.\n",
        "        \n",
        "        stride = 1 # Default value\n",
        "        downsample = False\n",
        "        \n",
        "        if intermediate_channels == in_channels/2:\n",
        "            stride = 2\n",
        "            downsample = True\n",
        "            \n",
        "        self.skip_connection = nn.Sequential(nn.Conv2d(in_channels, intermediate_channels*4, kernel_size=1, stride=stride, bias = False),\n",
        "                                             nn.BatchNorm2d(intermediate_channels*4)\n",
        "                                            )\n",
        "        \n",
        "        # Downsampling the output shape.\n",
        "        if downsample:\n",
        "            self.conv1 = nn.Conv2d(in_channels, intermediate_channels, kernel_size=1, stride=2, bias = False)\n",
        "            \n",
        "        else:\n",
        "            self.conv1 = nn.Conv2d(in_channels, intermediate_channels, kernel_size=1, stride=1, bias = False)\n",
        "            \n",
        "        self.bn1 = nn.BatchNorm2d(intermediate_channels)\n",
        "        \n",
        "        self.conv2 = nn.Conv2d(intermediate_channels, intermediate_channels, kernel_size=3, stride=1, padding = 1, bias = False)\n",
        "        self.bn2 = nn.BatchNorm2d(intermediate_channels)\n",
        "        \n",
        "        self.conv3 = nn.Conv2d(intermediate_channels, intermediate_channels*4, kernel_size=1, stride=1, bias = False)\n",
        "        self.bn3 = nn.BatchNorm2d(intermediate_channels*4)\n",
        "        \n",
        "        self.relu = nn.ReLU()\n",
        "        \n",
        "    def forward(self, x):\n",
        "        \n",
        "        # Residual to be added later.\n",
        "        identity = self.skip_connection(x)\n",
        "        \n",
        "        #---------------------------\n",
        "        # Layer-1\n",
        "        x = self.conv1(x)\n",
        "        x = self.bn1(x)\n",
        "        x = self.relu(x)\n",
        "        #----------------------------\n",
        "        \n",
        "        #---------------------------\n",
        "        # Layer-2\n",
        "        x = self.conv2(x)\n",
        "        x = self.bn2(x)\n",
        "        x = self.relu(x)\n",
        "        #----------------------------\n",
        "        \n",
        "        #---------------------------\n",
        "        # Layer-3\n",
        "        x = self.conv3(x)\n",
        "        x = self.bn3(x)\n",
        "        #print(\"Size identity\", identity.size())\n",
        "        #print(\"Size X\", x.size())\n",
        "        x += identity\n",
        "        x = self.relu(x)\n",
        "        #----------------------------\n",
        "        \n",
        "        return x\n",
        "    \n",
        "#---------------------------------------------------------------------------------------------------\n",
        "'''\n",
        "  This is Tier-1 Resnet class (Resnet-50, 101, 152)\n",
        "  It takes input as img_channels:\n",
        "  for RGB = 3\n",
        "  for Gray = 1),\n",
        "  \n",
        "  num_layers: \n",
        "  for Resnet-50 = [3, 4, 6, 3]\n",
        "  for Resnet-101 = [3, 4, 23, 3]\n",
        "  for Resnet-152 = [3, 8, 36, 3]\n",
        "  \n",
        "  , and number of classes. \n",
        "  for ex : Imagenet = 1000, MNIST = 10 etc\n",
        "  \n",
        "  Output : Resnet Model\n",
        "  \n",
        "'''\n",
        "class ResNet_Tier1(nn.Module):\n",
        "    \n",
        "    def __init__(self, img_channels, num_layers, num_classes):\n",
        "        super(ResNet_Tier1, self).__init__()\n",
        "        \n",
        "        # Layers \n",
        "        # Layer-0 Output shape : 64 X 56 X 56\n",
        "        self.layer0 = nn.Sequential(\n",
        "            nn.Conv2d(img_channels, 64, kernel_size=7, stride=2, padding=3),\n",
        "            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),\n",
        "            nn.BatchNorm2d(64),\n",
        "            nn.ReLU()\n",
        "        )\n",
        "        \n",
        "        # Residual Blocks layers\n",
        "        self.layer1 = self._make_layer(ResBlock_Tier1, num_layers[0], 64, 64)\n",
        "        self.layer2 = self._make_layer(ResBlock_Tier1, num_layers[1], 256, 128)\n",
        "        self.layer3 = self._make_layer(ResBlock_Tier1, num_layers[2], 512, 256)\n",
        "        self.layer4 = self._make_layer(ResBlock_Tier1, num_layers[3], 1024, 512)\n",
        "        \n",
        "        # FC Layers\n",
        "        self.avgpool = nn.AdaptiveAvgPool2d((1,1))\n",
        "        self.fc  = nn.Linear(512 * 4, num_classes)\n",
        "        \n",
        "    def forward(self, x):\n",
        "        x = self.layer0(x)\n",
        "        x = self.layer1(x)\n",
        "        x = self.layer2(x)\n",
        "        x = self.layer3(x)\n",
        "        x = self.layer4(x)\n",
        "        x = self.avgpool(x)\n",
        "        x = x.reshape(x.shape[0], -1)\n",
        "        x = self.fc(x)\n",
        "        \n",
        "        return x\n",
        "        \n",
        "    def _make_layer(self, ResBlock_Tier1, num_residual_blocks, in_channels, intermediate_channels):\n",
        "        \n",
        "        layers = []\n",
        "        only_once = True\n",
        "        \n",
        "        for i in range(num_residual_blocks):\n",
        "            layers.append(ResBlock_Tier1(in_channels, intermediate_channels))\n",
        "            if only_once:\n",
        "                in_channels = intermediate_channels*4\n",
        "                only_once = False\n",
        "            \n",
        "        return nn.Sequential(*layers)"
      ],
      "metadata": {
        "id": "JPckAYCuX4xo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "'''Resnets'''\n",
        "\n",
        "# Tier-2 Resnets.\n",
        "def ResNet18(img_channel=3, num_classes=1000):\n",
        "    return ResNet_Tier2(img_channel, [2, 2, 2, 2], num_classes)\n",
        "\n",
        "\n",
        "def ResNet34(img_channel=3, num_classes=1000):\n",
        "    return ResNet_Tier2(img_channel, [3, 4, 6, 3], num_classes)\n",
        "\n",
        "#--------------------------------------------------------------------\n",
        "\n",
        "# Tier-1 Resnets.\n",
        "def ResNet50(img_channel=3, num_classes=1000):\n",
        "    return ResNet_Tier1(img_channel, [3, 4, 6, 3], num_classes)\n",
        "\n",
        "\n",
        "def ResNet101(img_channel=3, num_classes=1000):\n",
        "    return ResNet_Tier1(img_channel, [3, 4, 23, 3], num_classes)\n",
        "\n",
        "\n",
        "def ResNet152(img_channel=3, num_classes=1000):\n",
        "    return ResNet_Tier1(img_channel, [3, 8, 36, 3], num_classes)"
      ],
      "metadata": {
        "id": "ghntw19YYLYl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## EfficientNet"
      ],
      "metadata": {
        "id": "OuYwku5UYtvv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "''' A simple Convolution, Batch Normalization, and Activation Class'''\n",
        "\n",
        "class ConvBnAct(nn.Module):\n",
        "    \n",
        "    def __init__(self, n_in, n_out, kernel_size = 3, stride = 1, \n",
        "                 padding = 0, groups = 1, bn = True, act = True,\n",
        "                 bias = False\n",
        "                ):\n",
        "        \n",
        "        super(ConvBnAct, self).__init__()\n",
        "        \n",
        "        self.conv = nn.Conv2d(n_in, n_out, kernel_size = kernel_size,\n",
        "                              stride = stride, padding = padding,\n",
        "                              groups = groups, bias = bias\n",
        "                             )\n",
        "        self.batch_norm = nn.BatchNorm2d(n_out) if bn else nn.Identity()\n",
        "        self.activation = nn.SiLU() if act else nn.Identity()\n",
        "        \n",
        "    def forward(self, x):\n",
        "        \n",
        "        x = self.conv(x)\n",
        "        x = self.batch_norm(x)\n",
        "        x = self.activation(x)\n",
        "        \n",
        "        return x\n",
        "    \n",
        "#------------------------------------------------------------------------------\n",
        "\n",
        "''' Squeeze and Excitation Block '''\n",
        "\n",
        "class SqueezeExcitation(nn.Module):\n",
        "    \n",
        "    def __init__(self, n_in, reduced_dim):\n",
        "        super(SqueezeExcitation, self).__init__()\n",
        "        \n",
        "        \n",
        "        self.se = nn.Sequential(\n",
        "            nn.AdaptiveAvgPool2d(1),\n",
        "            nn.Conv2d(n_in, reduced_dim, kernel_size=1),\n",
        "            nn.SiLU(),\n",
        "            nn.Conv2d(reduced_dim, n_in, kernel_size=1),\n",
        "            nn.Sigmoid()\n",
        "        )\n",
        "       \n",
        "    def forward(self, x):\n",
        "        \n",
        "        y = self.se(x)\n",
        "        \n",
        "        return x * y\n",
        "                                    \n",
        "#------------------------------------------------------------------------------\n",
        "\n",
        "''' Stochastic Depth Module'''\n",
        "\n",
        "class StochasticDepth(nn.Module):\n",
        "    \n",
        "    def __init__(self, survival_prob = 0.8):\n",
        "        super(StochasticDepth, self).__init__()\n",
        "        \n",
        "        self.p =  survival_prob\n",
        "        \n",
        "    def forward(self, x):\n",
        "        \n",
        "        if not self.training:\n",
        "            return x\n",
        "        \n",
        "        binary_tensor = torch.rand(x.shape[0], 1, 1, 1, device=x.device) < self.p\n",
        "        \n",
        "        return torch.div(x, self.p) * binary_tensor\n",
        "        \n",
        "#-------------------------------------------------------------------------------\n",
        "\n",
        "''' Residual Bottleneck Block with Expansion Factor = N as defined in Mobilenet-V2 paper\n",
        "    with Squeeze and Excitation Block and Stochastic Depth. \n",
        "'''\n",
        "\n",
        "class MBConvN(nn.Module):\n",
        "    \n",
        "    def __init__(self, n_in, n_out, kernel_size = 3, \n",
        "                 stride = 1, expansion_factor = 6,\n",
        "                 reduction = 4, # Squeeze and Excitation Block\n",
        "                 survival_prob = 0.8 # Stochastic Depth\n",
        "                ):\n",
        "        \n",
        "        super(MBConvN, self).__init__()\n",
        "        \n",
        "        self.skip_connection = (stride == 1 and n_in == n_out) \n",
        "        intermediate_channels = int(n_in * expansion_factor)\n",
        "        padding = (kernel_size - 1)//2\n",
        "        reduced_dim = int(n_in//reduction)\n",
        "        \n",
        "        self.expand = nn.Identity() if (expansion_factor == 1) else ConvBnAct(n_in, intermediate_channels, kernel_size = 1)\n",
        "        self.depthwise_conv = ConvBnAct(intermediate_channels, intermediate_channels,\n",
        "                                        kernel_size = kernel_size, stride = stride, \n",
        "                                        padding = padding, groups = intermediate_channels\n",
        "                                       )\n",
        "        self.se = SqueezeExcitation(intermediate_channels, reduced_dim = reduced_dim)\n",
        "        self.pointwise_conv = ConvBnAct(intermediate_channels, n_out, \n",
        "                                        kernel_size = 1, act = False\n",
        "                                       )\n",
        "        self.drop_layers = StochasticDepth(survival_prob = survival_prob)\n",
        "        \n",
        "    def forward(self, x):\n",
        "        \n",
        "        residual = x\n",
        "        \n",
        "        x = self.expand(x)\n",
        "        x = self.depthwise_conv(x)\n",
        "        x = self.se(x)\n",
        "        x = self.pointwise_conv(x)\n",
        "        \n",
        "        if self.skip_connection:\n",
        "            x = self.drop_layers(x)\n",
        "            x += residual\n",
        "        \n",
        "        return x\n",
        "    \n",
        "\n",
        "#----------------------------------------------------------------------------------------------\n",
        "\n",
        "'''Efficient-net Class'''\n",
        "\n",
        "class EfficientNet(nn.Module):\n",
        "    \n",
        "    '''Generic Efficient net class which takes width multiplier, Depth multiplier, and Survival Prob.'''\n",
        "    \n",
        "    def __init__(self, width_mult = 1, depth_mult = 1, \n",
        "                 dropout_rate = 0.2, num_classes = 1000):\n",
        "        super(EfficientNet, self).__init__()\n",
        "        \n",
        "        last_channel = ceil(1280 * width_mult)\n",
        "        self.features = self._feature_extractor(width_mult, depth_mult, last_channel)\n",
        "        self.avgpool = nn.AdaptiveAvgPool2d(1)\n",
        "        self.classifier = nn.Sequential(\n",
        "            nn.Dropout(dropout_rate),\n",
        "            nn.Linear(last_channel, num_classes)\n",
        "        )\n",
        "        \n",
        "    def forward(self, x):\n",
        "        \n",
        "        x = self.features(x)\n",
        "        x = self.avgpool(x)\n",
        "        x = self.classifier(x.view(x.shape[0], -1))\n",
        "        \n",
        "        return x\n",
        "    \n",
        "        \n",
        "    def _feature_extractor(self, width_mult, depth_mult, last_channel):\n",
        "        \n",
        "        channels = 4*ceil(int(32*width_mult) / 4)\n",
        "        layers = [ConvBnAct(3, channels, kernel_size = 3, stride = 2, padding = 1)]\n",
        "        in_channels = channels\n",
        "        \n",
        "        kernels = [3, 3, 5, 3, 5, 5, 3]\n",
        "        expansions = [1, 6, 6, 6, 6, 6, 6]\n",
        "        num_channels = [16, 24, 40, 80, 112, 192, 320]\n",
        "        num_layers = [1, 2, 2, 3, 3, 4, 1]\n",
        "        strides =[1, 2, 2, 2, 1, 2, 1]\n",
        "        \n",
        "        # Scale channels and num_layers according to width and depth multipliers.\n",
        "        scaled_num_channels = [4*ceil(int(c*width_mult) / 4) for c in num_channels]\n",
        "        scaled_num_layers = [int(d * depth_mult) for d in num_layers]\n",
        "\n",
        "        \n",
        "        for i in range(len(scaled_num_channels)):\n",
        "             \n",
        "            layers += [MBConvN(in_channels if repeat==0 else scaled_num_channels[i], \n",
        "                               scaled_num_channels[i],\n",
        "                               kernel_size = kernels[i],\n",
        "                               stride = strides[i] if repeat==0 else 1, \n",
        "                               expansion_factor = expansions[i]\n",
        "                              )\n",
        "                       for repeat in range(scaled_num_layers[i])\n",
        "                      ]\n",
        "            in_channels = scaled_num_channels[i]\n",
        "        \n",
        "        layers.append(ConvBnAct(in_channels, last_channel, kernel_size = 1, stride = 1, padding = 0))\n",
        "    \n",
        "        return nn.Sequential(*layers)"
      ],
      "metadata": {
        "id": "iVDw_nobYo-B"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Compound scaling factors for efficient-net family.\n",
        "EFFICIENT_NET_CONFIG = {\n",
        "    # tuple of width multiplier, depth multiplier, resolution, and Survival Prob\n",
        "    \"b0\" : (1.0, 1.0, 224, 0.2),\n",
        "    \"b1\" : (1.0, 1.1, 240, 0.2),\n",
        "    \"b2\" : (1.1, 1.2, 260, 0.3),\n",
        "    \"b3\" : (1.2, 1.4, 300, 0.3),\n",
        "    \"b4\" : (1.4, 1.8, 380, 0.4),\n",
        "    \"b5\" : (1.6, 2.2, 456, 0.4),\n",
        "    \"b6\" : (1.8, 2.6, 528, 0.5),\n",
        "    \"b7\" : (2.0, 3.1, 600, 0.5)\n",
        "}"
      ],
      "metadata": {
        "id": "Nw7tNxReZf6K"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Efficient-nets.\n",
        "def EffNet(version, num_classes=1000):\n",
        "  width_mult, depth_mult, res, dropout_rate = EFFICIENT_NET_CONFIG[version]\n",
        "  return EfficientNet(width_mult, depth_mult, dropout_rate, num_classes)"
      ],
      "metadata": {
        "id": "NZxelPdoZqkM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "##  EfficientNet-V2"
      ],
      "metadata": {
        "id": "wzNI4Q3x0Q2s"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "Eff_V2_SETTINGS = {\n",
        "    # expansion factor, k, stride, n_in, n_out, num_layers, use_fusedMBCONV\n",
        "    's' : [\n",
        "        [1, 3, 1, 24, 24, 2, True],\n",
        "        [4, 3, 2, 24, 48, 4, True],\n",
        "        [4, 3, 2, 48, 64, 4, True],\n",
        "        [4, 3, 2, 64, 128, 6, False],\n",
        "        [6, 3, 1, 128, 160, 9, False],\n",
        "        [6, 3, 2, 160, 256, 15, False]\n",
        "    ],\n",
        "    \n",
        "    'm' : [\n",
        "        [1, 3, 1, 24, 24, 3, True],\n",
        "        [4, 3, 2, 24, 48, 5, True],\n",
        "        [4, 3, 2, 48, 80, 5, True],\n",
        "        [4, 3, 2, 80, 160, 7, False],\n",
        "        [6, 3, 1, 160, 176, 14, False],\n",
        "        [6, 3, 2, 176, 304, 18, False],\n",
        "        [6, 3, 1, 304, 512, 5, False]\n",
        "    ],\n",
        "    \n",
        "    'l' : [\n",
        "        [1, 3, 1, 32, 32, 4, True],\n",
        "        [4, 3, 2, 32, 64, 7, True],\n",
        "        [4, 3, 2, 64, 96, 7, True],\n",
        "        [4, 3, 2, 96, 192, 10, False],\n",
        "        [6, 3, 1, 192, 224, 19, False],\n",
        "        [6, 3, 2, 224, 384, 25, False],\n",
        "        [6, 3, 1, 384, 640, 7, False]\n",
        "    ]\n",
        "}"
      ],
      "metadata": {
        "id": "NMZk_hoo0YZe"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "'''A simple Convolution + Batch Normalization + Activation Class'''\n",
        "\n",
        "class ConvBnAct(nn.Module):\n",
        "    \n",
        "    def __init__(\n",
        "        self,\n",
        "        n_in, # in_channels\n",
        "        n_out, # out_channels\n",
        "        k_size = 3, # Kernel Size\n",
        "        stride = 1, \n",
        "        padding = 0,\n",
        "        groups = 1, \n",
        "        act = True, \n",
        "        bn = True, \n",
        "        bias = False\n",
        "    ):\n",
        "        super(ConvBnAct, self).__init__()\n",
        "        \n",
        "        self.conv = nn.Conv2d(n_in, n_out, kernel_size = k_size, stride = stride,\n",
        "                              padding = padding, groups = groups,bias = bias\n",
        "                             )\n",
        "        self.batch_norm = nn.BatchNorm2d(n_out) if bn else nn.Identity()\n",
        "        self.activation = nn.SiLU() if act else nn.Identity()\n",
        "        \n",
        "    def forward(self, x):\n",
        "        x = self.conv(x)\n",
        "        x = self.batch_norm(x)\n",
        "        x = self.activation(x)\n",
        "        \n",
        "        return x\n",
        "    \n",
        "#--------------------------------------------------------------------------------------------\n",
        "\n",
        "'''Squeeze and Excitation Class'''\n",
        "\n",
        "class SqueezeExcitation(nn.Module):\n",
        "    \n",
        "    def __init__(\n",
        "        self,\n",
        "        n_in, # In_channels\n",
        "        reduced_dim\n",
        "    ):\n",
        "        super(SqueezeExcitation, self).__init__()\n",
        "      \n",
        "        self.squeeze = nn.AdaptiveAvgPool2d(1)\n",
        "        self.excite = nn.Sequential(nn.Conv2d(n_in, reduced_dim, kernel_size=1),\n",
        "                                   nn.SiLU(),\n",
        "                                   nn.Conv2d(reduced_dim, n_in, kernel_size=1),\n",
        "                                   nn.Sigmoid()\n",
        "                                   )\n",
        "        \n",
        "    def forward(self, x):\n",
        "        y = self.squeeze(x)\n",
        "        y = self.excite(y)\n",
        "            \n",
        "        return x * y\n",
        "        \n",
        "#--------------------------------------------------------------------------------------\n",
        "\n",
        "''' Stochastic Depth Class'''\n",
        "\n",
        "class StochasticDepth(nn.Module):\n",
        "    \n",
        "    def __init__(\n",
        "        self,\n",
        "        survival_prob = 0.8\n",
        "    ):\n",
        "        super(StochasticDepth, self).__init__()\n",
        "        \n",
        "        self.p =  survival_prob\n",
        "        \n",
        "    def forward(self, x):\n",
        "        \n",
        "        if not self.training:\n",
        "            return x\n",
        "        \n",
        "        binary_tensor = torch.rand(x.shape[0], 1, 1, 1, device=x.device) < self.p\n",
        "        \n",
        "        return torch.div(x, self.p) * binary_tensor\n",
        "        \n",
        "#-------------------------------------------------------------------------------\n",
        "\n",
        "'''MBCONV Class'''\n",
        "\n",
        "class MBConvN(nn.Module):\n",
        "    \n",
        "    def __init__(\n",
        "        self,\n",
        "        n_in, # In_channels\n",
        "        n_out, # out_channels\n",
        "        k_size = 3, # kernel_size\n",
        "        stride = 1,\n",
        "        expansion_factor = 4,\n",
        "        reduction_factor = 4, # SqueezeExcitation Block\n",
        "        survival_prob = 0.8 # StochasticDepth Block\n",
        "    ):\n",
        "        super(MBConvN, self).__init__()\n",
        "        reduced_dim = int(n_in//4)\n",
        "        expanded_dim = int(expansion_factor * n_in)\n",
        "        padding = (k_size - 1)//2\n",
        "        \n",
        "        self.use_residual = (n_in == n_out) and (stride == 1)\n",
        "        self.expand = nn.Identity() if (expansion_factor == 1) else ConvBnAct(n_in, expanded_dim, k_size = 1)\n",
        "        self.depthwise_conv = ConvBnAct(expanded_dim, expanded_dim,\n",
        "                                        k_size, stride = stride,\n",
        "                                        padding = padding, groups = expanded_dim\n",
        "                                       )\n",
        "        self.se = SqueezeExcitation(expanded_dim, reduced_dim)\n",
        "        self.drop_layers = StochasticDepth(survival_prob)\n",
        "        self.pointwise_conv = ConvBnAct(expanded_dim, n_out, k_size = 1, act = False)\n",
        "        \n",
        "    def forward(self, x):\n",
        "        \n",
        "        residual = x.clone()\n",
        "        x = self.expand(x)\n",
        "        x = self.depthwise_conv(x)\n",
        "        x = self.se(x)\n",
        "        x = self.pointwise_conv(x)\n",
        "        \n",
        "        if self.use_residual:\n",
        "            x = self.drop_layers(x)\n",
        "            x += residual\n",
        "            \n",
        "        return x\n",
        "    \n",
        "#--------------------------------------------------------------------------------------\n",
        "\n",
        "'''Fused-MBCONV Class'''\n",
        "\n",
        "class FusedMBConvN(nn.Module):\n",
        "    \n",
        "    def __init__(\n",
        "        self,\n",
        "        n_in, # In_channels\n",
        "        n_out, # out_channels\n",
        "        k_size = 3, # kernel_size\n",
        "        stride = 1,\n",
        "        expansion_factor = 4,\n",
        "        reduction_factor = 4, # SqueezeExcitation Block\n",
        "        survival_prob = 0.8 # StochasticDepth Block\n",
        "    ):\n",
        "        super(FusedMBConvN, self).__init__()\n",
        "        \n",
        "        reduced_dim = int(n_in//4)\n",
        "        expanded_dim = int(expansion_factor * n_in)\n",
        "        padding = (k_size - 1)//2\n",
        "        \n",
        "        self.use_residual = (n_in == n_out) and (stride == 1)\n",
        "        #self.expand = nn.Identity() if (expansion_factor == 1) else ConvBnAct(n_in, expanded_dim, k_size = 1)\n",
        "        self.conv = ConvBnAct(n_in, expanded_dim,\n",
        "                              k_size, stride = stride,\n",
        "                              padding = padding, groups = 1\n",
        "                             )\n",
        "        #self.se = SqueezeExcitation(expanded_dim, reduced_dim)\n",
        "        self.drop_layers = StochasticDepth(survival_prob)\n",
        "        self.pointwise_conv = nn.Identity() if (expansion_factor == 1) else ConvBnAct(expanded_dim, n_out, k_size = 1, act = False)\n",
        "        \n",
        "    def forward(self, x):\n",
        "        \n",
        "        residual = x.clone()\n",
        "        #x = self.conv(x)\n",
        "        x = self.conv(x)\n",
        "        #x = self.se(x)\n",
        "        x = self.pointwise_conv(x)\n",
        "        \n",
        "        if self.use_residual:\n",
        "            x = self.drop_layers(x)\n",
        "            x += residual\n",
        "            \n",
        "        return x\n",
        "    \n",
        "#-----------------------------------------------------------------------------------------------\n",
        "\n",
        "class EfficientNetV2(nn.Module):\n",
        "    \n",
        "    def __init__(\n",
        "    self,\n",
        "    version = 's',\n",
        "    dropout_rate = 0.2,\n",
        "    num_classes = 1000\n",
        "    ):\n",
        "        super(EfficientNetV2, self).__init__()\n",
        "        last_channel = 1280\n",
        "        self.features = self._feature_extractor(version, last_channel)\n",
        "        self.classifier = nn.Sequential(\n",
        "            nn.AdaptiveAvgPool2d((1,1)),\n",
        "            nn.Flatten(),\n",
        "            nn.Dropout(dropout_rate, inplace = True),\n",
        "            nn.Linear(last_channel, num_classes)\n",
        "        )\n",
        "        \n",
        "    def forward(self, x):\n",
        "        x = self.features(x)\n",
        "        x = self.classifier(x)\n",
        "        \n",
        "        return x\n",
        "        \n",
        "    def _feature_extractor(self, version, last_channel):\n",
        "        \n",
        "        # Extract the Config\n",
        "        config = Eff_V2_SETTINGS[version]\n",
        "        \n",
        "        layers = []\n",
        "        layers.append(ConvBnAct(3, config[0][3], k_size = 3, stride = 2, padding = 1))\n",
        "        #in_channel = config[0][3]\n",
        "        \n",
        "        for (expansion_factor, k, stride, n_in, n_out, num_layers, use_fused) in config:\n",
        "            \n",
        "            if use_fused:\n",
        "                layers += [FusedMBConvN(n_in if repeat==0 else n_out, \n",
        "                                        n_out,\n",
        "                                        k_size=k,\n",
        "                                        stride = stride if repeat==0 else 1,\n",
        "                                        expansion_factor=expansion_factor\n",
        "                                       ) for repeat in range(num_layers)\n",
        "                          ]\n",
        "            else:\n",
        "                \n",
        "                layers += [MBConvN(n_in if repeat==0 else n_out, \n",
        "                                   n_out,\n",
        "                                   k_size=k,\n",
        "                                   stride = stride if repeat==0 else 1,\n",
        "                                   expansion_factor=expansion_factor\n",
        "                                   ) for repeat in range(num_layers)\n",
        "                      ]\n",
        "                \n",
        "        layers.append(ConvBnAct(config[-1][4], last_channel, k_size = 1))   \n",
        "            \n",
        "        return nn.Sequential(*layers)"
      ],
      "metadata": {
        "id": "izERnCuU0YCP"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Efficient-nets.\n",
        "def EffNetV2(version, num_classes=1000):\n",
        "  return EfficientNetV2(version = version, num_classes = num_classes)"
      ],
      "metadata": {
        "id": "Jk8bFHF50X6m"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "--------------------------------------------------------------------------"
      ],
      "metadata": {
        "id": "_5WHGAGWbMeb"
      }
    }
  ]
}