NIU_discriminator.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model and trial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NLayerDiscriminator(nn.Module):\n",
    "    \"\"\"\n",
    "    'PatchGAN' classifier described in the original pix2pix paper.\n",
    "    It can classify whether 70×70 overlapping patches are real or fake.\n",
    "    Such a patch-level discriminator architecture has fewer parameters\n",
    "    than a full-image discriminator and can work on arbitrarily-sized images\n",
    "    in a fully convolutional fashion.\n",
    "    \"\"\"\n",
    "    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.InstanceNorm2d):\n",
    "        \"\"\"\n",
    "        Construct a PatchGAN discriminator\n",
    "        \n",
    "        Parameters:\n",
    "            1. input_nc (int): the number of channels in input images\n",
    "            2. ndf (int): the number filters\n",
    "            3. n_layers (int): the number of conv layers in the discriminator\n",
    "            4. norm_layer (nn.Module): normalization layer \n",
    "        \"\"\"\n",
    "        super(NLayerDiscriminator, self).__init__()\n",
    "        use_bias = (norm_layer == nn.InstanceNorm2d)\n",
    "    \n",
    "        kernel_size, padding = 4, 1\n",
    "        \n",
    "        # Leading Conv\n",
    "        sequence = [\n",
    "            nn.Conv2d(input_nc, ndf, kernel_size=kernel_size, stride=2, padding=padding),\n",
    "            nn.LeakyReLU(.2, True)\n",
    "        ]\n",
    "        \n",
    "        # gradually increase the number of filters and decrease the dimension of the image\n",
    "        mult_i = 1\n",
    "        for n in range(n_layers):\n",
    "            mult_o = min(mult_i * 2, 8)\n",
    "            ic, oc = ndf * mult_i, ndf * mult_o\n",
    "            sequence += [\n",
    "                nn.Conv2d(\n",
    "                    ic, oc, \n",
    "                    kernel_size=kernel_size, \n",
    "                    stride=2, \n",
    "                    padding=padding, \n",
    "                    bias=use_bias),\n",
    "                norm_layer(oc),\n",
    "                nn.LeakyReLU(.2, True)\n",
    "            ]\n",
    "            mult_i = mult_o\n",
    "        \n",
    "        mult_o = min(mult_i * 2, 8)\n",
    "        ic, oc = ndf * mult_i, ndf * mult_o\n",
    "        sequence += [\n",
    "            nn.Conv2d(ic, oc, kernel_size=kernel_size, stride=1, padding=padding, bias=use_bias),\n",
    "            norm_layer(oc),\n",
    "            nn.LeakyReLU(.2, True)\n",
    "        ]\n",
    "        \n",
    "        sequence += [nn.Conv2d(oc, 1, kernel_size=kernel_size, stride=1, padding=padding)]\n",
    "        \n",
    "        self.model = nn.Sequential(*sequence)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "     \n",
    "#     def forward(self, x):\n",
    "#         for l, layer in enumerate(self.model):\n",
    "#             y = layer(x)\n",
    "#             print(type(layer))\n",
    "#             print(f'\\n\\tencoder layer = {l}\\n\\tinput shape = {x.shape}\\n\\toutput shape = {y.shape}\\n')\n",
    "#             x = y\n",
    "#         return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [],
   "source": [
    "# class NLayerDiscriminator(nn.Module):\n",
    "#     \"\"\"\n",
    "#     'PatchGAN' classifier described in the original pix2pix paper.\n",
    "#     It can classify whether 70×70 overlapping patches are real or fake.\n",
    "#     Such a patch-level discriminator architecture has fewer parameters\n",
    "#     than a full-image discriminator and can work on arbitrarily-sized images\n",
    "#     in a fully convolutional fashion.\n",
    "#     \"\"\"\n",
    "#     def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.InstanceNorm2d):\n",
    "#         \"\"\"\n",
    "#         Construct a PatchGAN discriminator\n",
    "        \n",
    "#         Parameters:\n",
    "#             1. input_nc (int): the number of channels in input images\n",
    "#             2. ndf (int): the number filters\n",
    "#             3. n_layers (int): the number of conv layers in the discriminator\n",
    "#             4. norm_layer (nn.Module): normalization layer \n",
    "#         \"\"\"\n",
    "#         super(NLayerDiscriminator, self).__init__()\n",
    "#         use_bias = (norm_layer == nn.InstanceNorm2d)\n",
    "    \n",
    "#         kernel_size, padding = 3, 0\n",
    "        \n",
    "#         # Leading Conv\n",
    "#         sequence = [\n",
    "#             nn.Conv2d(input_nc, ndf, kernel_size=kernel_size, stride=2, padding=padding),\n",
    "#             nn.LeakyReLU(.2, True)\n",
    "#         ]\n",
    "        \n",
    "#         # gradually increase the number of filters and decrease the dimension of the image\n",
    "#         mult_i = 1\n",
    "#         for n in range(n_layers):\n",
    "#             mult_o = min(mult_i * 2, 8)\n",
    "#             ic, oc = ndf * mult_i, ndf * mult_o\n",
    "#             sequence += [\n",
    "#                 nn.Conv2d(\n",
    "#                     ic, oc, \n",
    "#                     kernel_size=kernel_size, \n",
    "#                     stride=2, \n",
    "#                     padding=padding, \n",
    "#                     bias=use_bias),\n",
    "#                 norm_layer(oc),\n",
    "#                 nn.LeakyReLU(.2, True)\n",
    "#             ]\n",
    "#             mult_i = mult_o\n",
    "            \n",
    "#         self.fc1 = nn.Linear(32 * 5 * 5, 80)\n",
    "#         self.fc2 = nn.Linear(80, 8)\n",
    "#         self.fc3 = nn.Linear(8, 1)\n",
    "        \n",
    "#         self.model = nn.Sequential(*sequence)\n",
    "        \n",
    "     \n",
    "#     def forward(self, x):\n",
    "#         for l, layer in enumerate(self.model):\n",
    "#             y = layer(x)\n",
    "#             # print(type(layer))\n",
    "#             # print(f'\\n\\tencoder layer = {l}\\n\\tinput shape = {x.shape}\\n\\toutput shape = {y.shape}\\n')\n",
    "#             x = y\n",
    "        \n",
    "#         # print(x.shape)\n",
    "#         x = torch.flatten(x, 1) # flatten all dimensions except batch (start_dim = 1)\n",
    "#         # print(x.shape)\n",
    "#         x = F.relu(self.fc1(x))\n",
    "#         # print(x.shape)\n",
    "#         x = F.relu(self.fc2(x))\n",
    "#         # print(x.shape)\n",
    "#         x = F.relu(self.fc3(x))\n",
    "#         # print(x.shape)\n",
    "#         return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trial(model, input_shape, cuda=True):\n",
    "    x = torch.rand(16, 1, *input_shape, dtype=torch.float32)\n",
    "    if cuda:\n",
    "        model = model.cuda()\n",
    "        x = x.cuda()\n",
    "\n",
    "    y = model.forward(x)\n",
    "    print(f'Output shape:\\n\\t{y.shape}\\n')\n",
    "\n",
    "    num_parameters_trainable = sum([p.numel() for p in model.parameters() if p.requires_grad])\n",
    "    num_parameters = sum([p.numel() for p in model.parameters()])\n",
    "    print(f'Number of parameters:\\n\\tTrainable = {num_parameters_trainable}\\n\\tTotal = {num_parameters}') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output shape:\n",
      "\ttorch.Size([16, 1, 2, 2])\n",
      "\n",
      "Number of parameters:\n",
      "\tTrainable = 44221\n",
      "\tTotal = 44221\n"
     ]
    }
   ],
   "source": [
    "discriminator = NLayerDiscriminator(1, ndf=4, n_layers=4, norm_layer=nn.InstanceNorm2d)\n",
    "input_shape = [128, 128]\n",
    "trial(discriminator, input_shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dataset_ls4gan(Dataset):\n",
    "    \"\"\"\n",
    "    LS4GAN dataset\n",
    "    \"\"\"\n",
    "    def __init__(self, class_paths):\n",
    "        super(Dataset_ls4gan, self).__init__()    \n",
    "        self.image_fnames = []\n",
    "        self.labels = []\n",
    "        for c, class_path in enumerate(class_paths):\n",
    "            fnames = list(Path(class_path).glob('*npz'))\n",
    "            self.image_fnames += fnames\n",
    "            self.labels += [c] * len(fnames)\n",
    "            \n",
    "        # print(len(self.image_fnames) == len(self.labels))    \n",
    "        \n",
    "        shuffle_indices = np.arange(len(self.image_fnames))\n",
    "        np.random.shuffle(shuffle_indices)\n",
    "        self.image_fnames = np.array(self.image_fnames)[shuffle_indices]\n",
    "        self.labels = np.array(self.labels)[shuffle_indices]\n",
    "        self.labels = self.labels.astype(np.float32)\n",
    "        \n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.image_fnames)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        image_fname, label = self.image_fnames[idx], self.labels[idx]\n",
    "        \n",
    "        image = np.load(image_fname)\n",
    "        key = list(image.keys())[0]\n",
    "        image = image[key]\n",
    "        image = np.expand_dims(np.float32(image), 0)\n",
    "        \n",
    "        image_tensor = torch.from_numpy(image)\n",
    "        label_tensor = torch.tensor([label])\n",
    "        return image_tensor, label_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [],
   "source": [
    "class_paths_train = [\n",
    "    '/sdcc/u/yhuang2/PROJs/GAN/datasets/ls4gan/toyzero/U/trainA/',\n",
    "    '/sdcc/u/yhuang2/PROJs/GAN/datasets/ls4gan/toyzero/U/trainB/'\n",
    "]\n",
    "\n",
    "class_paths_test = [\n",
    "    '/sdcc/u/yhuang2/PROJs/GAN/datasets/ls4gan/toyzero/U/testA/',\n",
    "    '/sdcc/u/yhuang2/PROJs/GAN/datasets/ls4gan/toyzero/U/testB/'\n",
    "]\n",
    "\n",
    "bsz = 16\n",
    "dataset_train = Dataset_ls4gan(class_paths_train)\n",
    "dataset_test = Dataset_ls4gan(class_paths_test)\n",
    "train_loader = DataLoader(dataset_train, batch_size=bsz, shuffle=True)\n",
    "test_loader = DataLoader(dataset_test, batch_size=bsz, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1] loss: 0.709\n",
      "[2] loss: 0.698\n",
      "[3] loss: 0.694\n",
      "[4] loss: 0.693\n",
      "[5] loss: 0.694\n",
      "[6] loss: 0.694\n",
      "[7] loss: 0.693\n",
      "[8] loss: 0.693\n",
      "[9] loss: 0.693\n",
      "[10] loss: 0.693\n",
      "[11] loss: 0.693\n",
      "[12] loss: 0.693\n",
      "[13] loss: 0.693\n",
      "[14] loss: 0.693\n",
      "[15] loss: 0.693\n",
      "[16] loss: 0.693\n",
      "[17] loss: 0.693\n",
      "[18] loss: 0.693\n",
      "[19] loss: 0.693\n",
      "[20] loss: 0.693\n",
      "[21] loss: 0.693\n",
      "[22] loss: 0.693\n",
      "[23] loss: 0.693\n",
      "[24] loss: 0.694\n",
      "[25] loss: 0.694\n",
      "[26] loss: 0.693\n",
      "[27] loss: 0.693\n",
      "[28] loss: 0.693\n",
      "[29] loss: 0.693\n",
      "[30] loss: 0.693\n",
      "[31] loss: 0.694\n",
      "[32] loss: 0.694\n",
      "[33] loss: 0.693\n",
      "[34] loss: 0.693\n",
      "[35] loss: 0.693\n",
      "[36] loss: 0.693\n",
      "[37] loss: 0.693\n",
      "[38] loss: 0.693\n",
      "[39] loss: 0.693\n",
      "[40] loss: 0.693\n",
      "[41] loss: 0.693\n",
      "[42] loss: 0.693\n",
      "[43] loss: 0.693\n",
      "[44] loss: 0.693\n",
      "[45] loss: 0.693\n",
      "[46] loss: 0.693\n",
      "[47] loss: 0.693\n",
      "[48] loss: 0.693\n",
      "[49] loss: 0.693\n",
      "[50] loss: 0.693\n",
      "[51] loss: 0.693\n",
      "[52] loss: 0.693\n",
      "[53] loss: 0.693\n",
      "[54] loss: 0.693\n",
      "[55] loss: 0.693\n",
      "[56] loss: 0.693\n",
      "[57] loss: 0.693\n",
      "[58] loss: 0.693\n",
      "[59] loss: 0.693\n",
      "[60] loss: 0.693\n",
      "[61] loss: 0.693\n",
      "[62] loss: 0.693\n",
      "[63] loss: 0.693\n",
      "[64] loss: 0.693\n",
      "[65] loss: 0.693\n",
      "[66] loss: 0.693\n",
      "[67] loss: 0.693\n",
      "[68] loss: 0.693\n",
      "[69] loss: 0.693\n",
      "[70] loss: 0.693\n",
      "[71] loss: 0.693\n",
      "[72] loss: 0.694\n",
      "[73] loss: 0.693\n",
      "[74] loss: 0.693\n",
      "[75] loss: 0.693\n",
      "[76] loss: 0.694\n",
      "[77] loss: 0.693\n",
      "[78] loss: 0.694\n",
      "[79] loss: 0.693\n",
      "[80] loss: 0.693\n",
      "[81] loss: 0.693\n",
      "[82] loss: 0.693\n",
      "[83] loss: 0.693\n",
      "[84] loss: 0.693\n",
      "[85] loss: 0.693\n",
      "[86] loss: 0.694\n",
      "[87] loss: 0.694\n",
      "[88] loss: 0.693\n",
      "[89] loss: 0.693\n",
      "[90] loss: 0.693\n",
      "[91] loss: 0.693\n",
      "[92] loss: 0.693\n",
      "[93] loss: 0.693\n",
      "[94] loss: 0.693\n",
      "[95] loss: 0.693\n",
      "[96] loss: 0.693\n",
      "[97] loss: 0.693\n",
      "[98] loss: 0.693\n",
      "[99] loss: 0.694\n",
      "[100] loss: 0.693\n",
      "[101] loss: 0.693\n",
      "[102] loss: 0.693\n",
      "[103] loss: 0.694\n",
      "[104] loss: 0.693\n",
      "[105] loss: 0.693\n",
      "[106] loss: 0.693\n",
      "[107] loss: 0.694\n",
      "[108] loss: 0.693\n",
      "[109] loss: 0.693\n",
      "[110] loss: 0.693\n",
      "[111] loss: 0.694\n",
      "[112] loss: 0.693\n",
      "[113] loss: 0.693\n",
      "[114] loss: 0.693\n",
      "[115] loss: 0.694\n",
      "[116] loss: 0.693\n",
      "[117] loss: 0.693\n",
      "[118] loss: 0.693\n",
      "[119] loss: 0.693\n",
      "[120] loss: 0.693\n",
      "[121] loss: 0.694\n",
      "[122] loss: 0.693\n",
      "[123] loss: 0.693\n",
      "[124] loss: 0.693\n",
      "[125] loss: 0.693\n",
      "[126] loss: 0.693\n",
      "[127] loss: 0.693\n",
      "[128] loss: 0.693\n",
      "[129] loss: 0.693\n",
      "[130] loss: 0.693\n",
      "[131] loss: 0.693\n",
      "[132] loss: 0.693\n",
      "[133] loss: 0.693\n",
      "[134] loss: 0.693\n",
      "[135] loss: 0.693\n",
      "[136] loss: 0.693\n",
      "[137] loss: 0.693\n",
      "[138] loss: 0.693\n",
      "[139] loss: 0.693\n",
      "[140] loss: 0.693\n",
      "[141] loss: 0.693\n",
      "[142] loss: 0.694\n",
      "[143] loss: 0.694\n",
      "[144] loss: 0.693\n",
      "[145] loss: 0.693\n",
      "[146] loss: 0.693\n",
      "[147] loss: 0.693\n",
      "[148] loss: 0.693\n",
      "[149] loss: 0.693\n",
      "[150] loss: 0.694\n",
      "[151] loss: 0.693\n",
      "[152] loss: 0.693\n",
      "[153] loss: 0.693\n",
      "[154] loss: 0.693\n",
      "[155] loss: 0.694\n",
      "[156] loss: 0.693\n",
      "[157] loss: 0.693\n",
      "[158] loss: 0.694\n",
      "[159] loss: 0.694\n",
      "[160] loss: 0.693\n",
      "[161] loss: 0.693\n",
      "[162] loss: 0.693\n",
      "[163] loss: 0.693\n",
      "[164] loss: 0.693\n",
      "[165] loss: 0.693\n",
      "[166] loss: 0.693\n",
      "[167] loss: 0.693\n",
      "[168] loss: 0.694\n",
      "[169] loss: 0.693\n",
      "[170] loss: 0.693\n",
      "[171] loss: 0.693\n",
      "[172] loss: 0.693\n",
      "[173] loss: 0.693\n",
      "[174] loss: 0.693\n",
      "[175] loss: 0.694\n",
      "[176] loss: 0.694\n",
      "[177] loss: 0.694\n",
      "[178] loss: 0.693\n",
      "[179] loss: 0.693\n",
      "[180] loss: 0.693\n",
      "[181] loss: 0.693\n",
      "[182] loss: 0.693\n",
      "[183] loss: 0.693\n",
      "[184] loss: 0.693\n",
      "[185] loss: 0.693\n",
      "[186] loss: 0.693\n",
      "[187] loss: 0.693\n",
      "[188] loss: 0.693\n",
      "[189] loss: 0.693\n",
      "[190] loss: 0.693\n",
      "[191] loss: 0.693\n",
      "[192] loss: 0.693\n",
      "[193] loss: 0.693\n",
      "[194] loss: 0.693\n",
      "[195] loss: 0.693\n",
      "[196] loss: 0.693\n",
      "[197] loss: 0.694\n",
      "[198] loss: 0.693\n",
      "[199] loss: 0.693\n",
      "[200] loss: 0.693\n",
      "Finished Training\n"
     ]
    }
   ],
   "source": [
    "discriminator = NLayerDiscriminator(1, ndf=4, n_layers=4, norm_layer=nn.InstanceNorm2d)\n",
    "discriminator = discriminator.cuda()\n",
    "\n",
    "criterion = nn.BCEWithLogitsLoss()\n",
    "# criterion = nn.L1Loss()\n",
    "optimizer = torch.optim.AdamW(discriminator.parameters(), lr=0.0001)\n",
    "epochs = 200\n",
    "\n",
    "for epoch in range(epochs):  # loop over the dataset multiple times\n",
    "\n",
    "    epoch_loss = 0\n",
    "    for i, data in enumerate(train_loader):\n",
    "        # get the inputs; data is a list of [inputs, labels]\n",
    "        inputs, labels = data\n",
    "        inputs = inputs.cuda()\n",
    "        labels = labels.cuda()\n",
    "        \n",
    "        # zero the parameter gradients\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # forward + backward + optimize\n",
    "        outputs = discriminator(inputs)\n",
    "        outputs = outputs.reshape_as(labels)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # print statistics\n",
    "        epoch_loss += loss.item()\n",
    "        if i == bsz - 1: \n",
    "            print(f'[{epoch + 1}] loss: {epoch_loss / bsz :.3f}')\n",
    "\n",
    "print('Finished Training')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "yi_test",
   "language": "python",
   "name": "yi_test"
  },
  "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}