Discriminator (1).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 192,
      "metadata": {
        "id": "nnZLVqF31p4Z"
      },
      "outputs": [],
      "source": [
        "from collections import OrderedDict\n",
        "import torch\n",
        "import torch.nn as nn"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class DisInputBlock(nn.Module):\n",
        "  def __init__(self, in_channels = None, out_channels = None, is_batch_norm=False):\n",
        "    super(DisInputBlock, self).__init__()\n",
        "\n",
        "    self.in_channels = in_channels\n",
        "    self.out_channels = out_channels\n",
        "    self.is_batch_norm = is_batch_norm\n",
        "\n",
        "    self.kernel_size = 3\n",
        "    self.stride = 1\n",
        "    self.padding = 1\n",
        "\n",
        "    self.model = self.input_block()\n",
        "\n",
        "  def input_block(self):\n",
        "    layers = OrderedDict()\n",
        "\n",
        "    if self.is_batch_norm:\n",
        "      layers[\"conv\"] = nn.Conv2d(\n",
        "          self.in_channels, self.out_channels, self.kernel_size, self.stride*2, self.padding)\n",
        "\n",
        "      layers[\"batch_norm\"] = nn.BatchNorm2d(self.out_channels)\n",
        "    else:\n",
        "      layers[\"conv\"] = nn.Conv2d(\n",
        "          self.in_channels, self.out_channels, self.kernel_size, self.stride, self.padding)\n",
        "\n",
        "    layers[\"lReLU\"] = nn.LeakyReLU(negative_slope=0.2)\n",
        "\n",
        "    return nn.Sequential(layers)\n",
        "\n",
        "  def forward(self, x):\n",
        "    return self.model(x)"
      ],
      "metadata": {
        "id": "jv7Vb1hn19UK"
      },
      "execution_count": 193,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class DisFeatureBlock(nn.Module):\n",
        "  def __init__(self, in_channels = None, out_channels = None):\n",
        "    super(DisFeatureBlock, self).__init__()\n",
        "\n",
        "    self.in_channels = in_channels\n",
        "    self.out_channels = out_channels\n",
        "\n",
        "    self.kernel_size = 3\n",
        "    self.stride = 1\n",
        "    self.padding = 1\n",
        "\n",
        "    self.model = self.feature_block()\n",
        "\n",
        "  def feature_block(self):\n",
        "    layers = []\n",
        "    layers.append(nn.Conv2d(\n",
        "          self.in_channels, self.out_channels, self.kernel_size, self.stride*2, self.padding))\n",
        "\n",
        "    layers.append(nn.BatchNorm2d(self.out_channels))\n",
        "\n",
        "    layers.append(nn.LeakyReLU(negative_slope=0.2))\n",
        "\n",
        "    return nn.Sequential(*layers)\n",
        "\n",
        "\n",
        "  def forward(self, x):\n",
        "    return self.model(x)"
      ],
      "metadata": {
        "id": "AVohAhZZ3zcN"
      },
      "execution_count": 194,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class DisOutBlock(nn.Module):\n",
        "  def __init__(self, in_channels=None, out_channels=None):\n",
        "    super(DisOutBlock, self).__init__()\n",
        "\n",
        "    self.in_channels = in_channels\n",
        "    self.out_channels = out_channels\n",
        "    self.kernel_size = 1\n",
        "\n",
        "    self.model = self.out_block()\n",
        "\n",
        "  def out_block(self):\n",
        "    layers = OrderedDict()\n",
        "\n",
        "    layers[\"out_conv\"] = nn.Conv2d(self.in_channels, self.out_channels, self.kernel_size)\n",
        "    layers[\"lReLU\"] = nn.LeakyReLU(negative_slope=0.2)\n",
        "\n",
        "    return nn.Sequential(layers)\n",
        "\n",
        "  def forward(self, x):\n",
        "    return self.model(x)"
      ],
      "metadata": {
        "id": "RmvhdFUj8eE6"
      },
      "execution_count": 195,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Discriminator(nn.Module):\n",
        "  def __init__(self):\n",
        "    super(Discriminator, self).__init__()\n",
        "\n",
        "    self.in_channels = 64\n",
        "    self.total_repetitive = 6\n",
        "\n",
        "    self.input = nn.Sequential(*[DisInputBlock(\n",
        "        in_channels=in_channels, out_channels=out_channels, is_batch_norm=batch_norm) for batch_norm, in_channels, out_channels in [[False, 3, 64], [True, 64, 64]]])\n",
        "\n",
        "    layers = []\n",
        "\n",
        "    for index in range(self.total_repetitive):\n",
        "      layers.append(DisFeatureBlock(self.in_channels, self.in_channels*2 if index%2==0 else self.in_channels))\n",
        "\n",
        "      if index%2==0:\n",
        "        self.in_channels=self.in_channels*2\n",
        "      else:\n",
        "        self.in_channels = self.in_channels\n",
        "\n",
        "    self.features_block = nn.Sequential(*layers)\n",
        "\n",
        "    self.avg_pool = nn.AdaptiveAvgPool2d(output_size=1)\n",
        "\n",
        "    self.out = nn.Sequential(*[DisOutBlock(in_channels, out_channels) for in_channels, out_channels in [[512, 1024], [1024, 1]]])\n",
        "\n",
        "\n",
        "  def forward(self, x):\n",
        "    x = self.input(x)\n",
        "\n",
        "    x = self.features_block(x)\n",
        "\n",
        "    x = self.avg_pool(x)\n",
        "\n",
        "    return self.out(x)"
      ],
      "metadata": {
        "id": "wiRZK10F1v7V"
      },
      "execution_count": 196,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "netD = Discriminator()\n",
        "\n",
        "netD(torch.randn(64, 3, 256, 256)).size()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LqhK0Lmx5hK0",
        "outputId": "d8a68b94-a22e-4bd9-80f4-a9c8674adc00"
      },
      "execution_count": 197,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([64, 1, 1, 1])"
            ]
          },
          "metadata": {},
          "execution_count": 197
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "sum(params.numel() for params in netD.parameters())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "smORCrnU5oAN",
        "outputId": "f82d47df-4959-4186-a659-ef07dc550a74"
      },
      "execution_count": 198,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "5215425"
            ]
          },
          "metadata": {},
          "execution_count": 198
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "netD"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Mf_-A_m8Ar6E",
        "outputId": "041a37b3-3c3a-4a6c-c811-71d8a8a62541"
      },
      "execution_count": 199,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Discriminator(\n",
              "  (input): Sequential(\n",
              "    (0): DisInputBlock(\n",
              "      (model): Sequential(\n",
              "        (conv): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "        (lReLU): LeakyReLU(negative_slope=0.2)\n",
              "      )\n",
              "    )\n",
              "    (1): DisInputBlock(\n",
              "      (model): Sequential(\n",
              "        (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
              "        (batch_norm): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (lReLU): LeakyReLU(negative_slope=0.2)\n",
              "      )\n",
              "    )\n",
              "  )\n",
              "  (features_block): Sequential(\n",
              "    (0): DisFeatureBlock(\n",
              "      (model): Sequential(\n",
              "        (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
              "        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (2): LeakyReLU(negative_slope=0.2)\n",
              "      )\n",
              "    )\n",
              "    (1): DisFeatureBlock(\n",
              "      (model): Sequential(\n",
              "        (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
              "        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (2): LeakyReLU(negative_slope=0.2)\n",
              "      )\n",
              "    )\n",
              "    (2): DisFeatureBlock(\n",
              "      (model): Sequential(\n",
              "        (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
              "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (2): LeakyReLU(negative_slope=0.2)\n",
              "      )\n",
              "    )\n",
              "    (3): DisFeatureBlock(\n",
              "      (model): Sequential(\n",
              "        (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
              "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (2): LeakyReLU(negative_slope=0.2)\n",
              "      )\n",
              "    )\n",
              "    (4): DisFeatureBlock(\n",
              "      (model): Sequential(\n",
              "        (0): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
              "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (2): LeakyReLU(negative_slope=0.2)\n",
              "      )\n",
              "    )\n",
              "    (5): DisFeatureBlock(\n",
              "      (model): Sequential(\n",
              "        (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
              "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (2): LeakyReLU(negative_slope=0.2)\n",
              "      )\n",
              "    )\n",
              "  )\n",
              "  (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
              "  (out): Sequential(\n",
              "    (0): DisOutBlock(\n",
              "      (model): Sequential(\n",
              "        (out_conv): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
              "        (lReLU): LeakyReLU(negative_slope=0.2)\n",
              "      )\n",
              "    )\n",
              "    (1): DisOutBlock(\n",
              "      (model): Sequential(\n",
              "        (out_conv): Conv2d(1024, 1, kernel_size=(1, 1), stride=(1, 1))\n",
              "        (lReLU): LeakyReLU(negative_slope=0.2)\n",
              "      )\n",
              "    )\n",
              "  )\n",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 199
        }
      ]
    }
  ]
}