StarGAN2_multiGPU_MAI.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "import time\n",
    "import csv\n",
    "\n",
    "from stargan2.solver_multi_GPU import Solver\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torchvision.utils as vutils\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "\n",
    "import torchvision.models.resnet as resnet\n",
    "\n",
    "from torch.backends import cudnn\n",
    "cudnn.benchmark = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision.datasets as datasets\n",
    "from itertools import groupby\n",
    "import random\n",
    "from torch.utils.data.sampler import WeightedRandomSampler\n",
    "\n",
    "class ReferenceDataset(datasets.DatasetFolder):\n",
    "    def __init__(\n",
    "            self,\n",
    "            root,\n",
    "            transform = None,\n",
    "            target_transform = None,\n",
    "            loader = datasets.folder.default_loader,\n",
    "            is_valid_file = None,\n",
    "    ):\n",
    "        super(ReferenceDataset, self).__init__(root, loader, \n",
    "                                          datasets.folder.IMG_EXTENSIONS if is_valid_file is None else None,\n",
    "                                          transform=transform,\n",
    "                                          target_transform=target_transform,\n",
    "                                          is_valid_file=is_valid_file)\n",
    "        \n",
    "        # group samples by label\n",
    "        groupped = {k: list(v) for k, v in  groupby(self.samples, lambda x: x[1])}\n",
    "        \n",
    "        # crate reference images\n",
    "        references = []\n",
    "        targets = []\n",
    "        for domain, values in groupped.items():\n",
    "            # unzip samples and targets\n",
    "            samples, _ = zip(*values)\n",
    "            # shuffle second reference images\n",
    "            samples2 = random.sample(samples, len(samples))\n",
    "            # repeat labels\n",
    "            labels = [domain]*len(samples)\n",
    "            targets+= labels\n",
    "            references+= list(zip(samples, samples2, labels))\n",
    "        # override samples\n",
    "        self.samples = references\n",
    "        # override targets to make sure that the samples have corresponding labels\n",
    "        self.targets = targets\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            index (int): Index\n",
    "        Returns:\n",
    "            tuple: (ref sample 1, ref sample 2, ref target)\n",
    "        \"\"\"\n",
    "        ref_path1, ref_path2, ref_target = self.samples[index]\n",
    "\n",
    "        ref1 = self.loader(ref_path1)\n",
    "        ref2 = self.loader(ref_path2)\n",
    "        \n",
    "        if self.transform is not None:\n",
    "            ref1 = self.transform(ref1)\n",
    "            ref2 = self.transform(ref2)\n",
    "            \n",
    "        if self.target_transform is not None:\n",
    "            ref_target = self.target_transform(ref_target)\n",
    "\n",
    "        return ref1, ref2, ref_target\n",
    "\n",
    "def make_balanced_sampler(labels):\n",
    "    class_counts = np.bincount(labels)\n",
    "    class_weights = 1. / class_counts\n",
    "    weights = class_weights[labels]\n",
    "    WeightedRandomSampler(weights, len(weights))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision import transforms\n",
    "from torchvision.datasets import ImageFolder\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "img_size = 128\n",
    "batch_size = 8\n",
    "num_workers = 4\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize([img_size, img_size]),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.5, 0.5, 0.5],\n",
    "                         std=[0.5, 0.5, 0.5]),\n",
    "])\n",
    "\n",
    "ds = ImageFolder('data/mai', transform)\n",
    "ds_ref = ReferenceDataset('data/mai', transform)\n",
    "\n",
    "sampler = make_balanced_sampler(ds.targets)\n",
    "loader = DataLoader(dataset=ds,\n",
    "                    batch_size=batch_size,\n",
    "                    sampler=sampler,\n",
    "                    shuffle=True,\n",
    "                    num_workers=num_workers,\n",
    "                    pin_memory=True,\n",
    "                    drop_last=True)\n",
    "\n",
    "sampler_ref = make_balanced_sampler(ds_ref.targets)\n",
    "loader_ref = DataLoader(dataset=ds_ref,\n",
    "                        batch_size=batch_size,\n",
    "                        sampler=sampler_ref,\n",
    "                        shuffle=True,\n",
    "                        num_workers=num_workers,\n",
    "                        pin_memory=True,\n",
    "                        drop_last=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n_domains: 2\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'wsi1_tiles': 0, 'wsi2_tiles': 1}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "domains, domains_mapping = ds._find_classes(ds.root)\n",
    "n_domains = len(domains)\n",
    "print('n_domains:', n_domains)\n",
    "domains_mapping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "solver = Solver('StarGAN2_MAI_128', 'r1_ds1_1', img_size, n_domains=2, lambda_ds=.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "solver.load_models('results/StarGAN2_MAI_128_ds1/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33marray\u001b[0m (use `wandb login --relogin` to force relogin)\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: wandb version 0.10.29 is available!  To upgrade, please run:\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m:  $ pip install wandb --upgrade\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                Tracking run with wandb version 0.10.13<br/>\n",
       "                Syncing run <strong style=\"color:#cdcd00\">helpful-snowflake-9</strong> to <a href=\"https://wandb.ai\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
       "                Project page: <a href=\"https://wandb.ai/array/StarGAN2_MAI_128\" target=\"_blank\">https://wandb.ai/array/StarGAN2_MAI_128</a><br/>\n",
       "                Run page: <a href=\"https://wandb.ai/array/StarGAN2_MAI_128/runs/3hcqbqlf\" target=\"_blank\">https://wandb.ai/array/StarGAN2_MAI_128/runs/3hcqbqlf</a><br/>\n",
       "                Run data is saved locally in <code>/data/aray/pathology_gan/wandb/run-20210505_170822-3hcqbqlf</code><br/><br/>\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Calling run.save without any arguments is deprecated.Changes to attributes are automatically persisted.\n"
     ]
    }
   ],
   "source": [
    "solver.train(300000, loader, loader_ref, val=ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchsummary import summary\n",
    "from stargan2.model import Discriminator, Generator, MappingNetwork, StyleEncoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [8, 128, 128, 128]           3,584\n",
      "            Conv2d-2         [8, 256, 128, 128]          32,768\n",
      "    InstanceNorm2d-3         [8, 128, 128, 128]             256\n",
      "         LeakyReLU-4         [8, 128, 128, 128]               0\n",
      "         LeakyReLU-5         [8, 128, 128, 128]               0\n",
      "         LeakyReLU-6         [8, 128, 128, 128]               0\n",
      "         LeakyReLU-7         [8, 128, 128, 128]               0\n",
      "         LeakyReLU-8         [8, 128, 128, 128]               0\n",
      "            Conv2d-9         [8, 128, 128, 128]         147,584\n",
      "   InstanceNorm2d-10           [8, 128, 64, 64]             256\n",
      "        LeakyReLU-11           [8, 128, 64, 64]               0\n",
      "        LeakyReLU-12           [8, 128, 64, 64]               0\n",
      "        LeakyReLU-13           [8, 128, 64, 64]               0\n",
      "        LeakyReLU-14           [8, 128, 64, 64]               0\n",
      "        LeakyReLU-15           [8, 128, 64, 64]               0\n",
      "           Conv2d-16           [8, 256, 64, 64]         295,168\n",
      "           ResBlk-17           [8, 256, 64, 64]               0\n",
      "           Conv2d-18           [8, 512, 64, 64]         131,072\n",
      "   InstanceNorm2d-19           [8, 256, 64, 64]             512\n",
      "        LeakyReLU-20           [8, 256, 64, 64]               0\n",
      "        LeakyReLU-21           [8, 256, 64, 64]               0\n",
      "        LeakyReLU-22           [8, 256, 64, 64]               0\n",
      "        LeakyReLU-23           [8, 256, 64, 64]               0\n",
      "        LeakyReLU-24           [8, 256, 64, 64]               0\n",
      "           Conv2d-25           [8, 256, 64, 64]         590,080\n",
      "   InstanceNorm2d-26           [8, 256, 32, 32]             512\n",
      "        LeakyReLU-27           [8, 256, 32, 32]               0\n",
      "        LeakyReLU-28           [8, 256, 32, 32]               0\n",
      "        LeakyReLU-29           [8, 256, 32, 32]               0\n",
      "        LeakyReLU-30           [8, 256, 32, 32]               0\n",
      "        LeakyReLU-31           [8, 256, 32, 32]               0\n",
      "           Conv2d-32           [8, 512, 32, 32]       1,180,160\n",
      "           ResBlk-33           [8, 512, 32, 32]               0\n",
      "   InstanceNorm2d-34           [8, 512, 32, 32]           1,024\n",
      "        LeakyReLU-35           [8, 512, 32, 32]               0\n",
      "        LeakyReLU-36           [8, 512, 32, 32]               0\n",
      "        LeakyReLU-37           [8, 512, 32, 32]               0\n",
      "        LeakyReLU-38           [8, 512, 32, 32]               0\n",
      "        LeakyReLU-39           [8, 512, 32, 32]               0\n",
      "           Conv2d-40           [8, 512, 32, 32]       2,359,808\n",
      "   InstanceNorm2d-41           [8, 512, 16, 16]           1,024\n",
      "        LeakyReLU-42           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-43           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-44           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-45           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-46           [8, 512, 16, 16]               0\n",
      "           Conv2d-47           [8, 512, 16, 16]       2,359,808\n",
      "           ResBlk-48           [8, 512, 16, 16]               0\n",
      "   InstanceNorm2d-49           [8, 512, 16, 16]           1,024\n",
      "        LeakyReLU-50           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-51           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-52           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-53           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-54           [8, 512, 16, 16]               0\n",
      "           Conv2d-55           [8, 512, 16, 16]       2,359,808\n",
      "   InstanceNorm2d-56           [8, 512, 16, 16]           1,024\n",
      "        LeakyReLU-57           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-58           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-59           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-60           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-61           [8, 512, 16, 16]               0\n",
      "           Conv2d-62           [8, 512, 16, 16]       2,359,808\n",
      "           ResBlk-63           [8, 512, 16, 16]               0\n",
      "   InstanceNorm2d-64           [8, 512, 16, 16]           1,024\n",
      "        LeakyReLU-65           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-66           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-67           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-68           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-69           [8, 512, 16, 16]               0\n",
      "           Conv2d-70           [8, 512, 16, 16]       2,359,808\n",
      "   InstanceNorm2d-71           [8, 512, 16, 16]           1,024\n",
      "        LeakyReLU-72           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-73           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-74           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-75           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-76           [8, 512, 16, 16]               0\n",
      "           Conv2d-77           [8, 512, 16, 16]       2,359,808\n",
      "           ResBlk-78           [8, 512, 16, 16]               0\n",
      "           Linear-79                  [8, 1024]          66,560\n",
      "   InstanceNorm2d-80           [8, 512, 16, 16]               0\n",
      "            AdaIN-81           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-82           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-83           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-84           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-85           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-86           [8, 512, 16, 16]               0\n",
      "           Conv2d-87           [8, 512, 16, 16]       2,359,808\n",
      "           Linear-88                  [8, 1024]          66,560\n",
      "   InstanceNorm2d-89           [8, 512, 16, 16]               0\n",
      "            AdaIN-90           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-91           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-92           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-93           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-94           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-95           [8, 512, 16, 16]               0\n",
      "           Conv2d-96           [8, 512, 16, 16]       2,359,808\n",
      "      AdainResBlk-97           [8, 512, 16, 16]               0\n",
      "           Linear-98                  [8, 1024]          66,560\n",
      "   InstanceNorm2d-99           [8, 512, 16, 16]               0\n",
      "           AdaIN-100           [8, 512, 16, 16]               0\n",
      "       LeakyReLU-101           [8, 512, 16, 16]               0\n",
      "       LeakyReLU-102           [8, 512, 16, 16]               0\n",
      "       LeakyReLU-103           [8, 512, 16, 16]               0\n",
      "       LeakyReLU-104           [8, 512, 16, 16]               0\n",
      "       LeakyReLU-105           [8, 512, 16, 16]               0\n",
      "          Conv2d-106           [8, 512, 16, 16]       2,359,808\n",
      "          Linear-107                  [8, 1024]          66,560\n",
      "  InstanceNorm2d-108           [8, 512, 16, 16]               0\n",
      "           AdaIN-109           [8, 512, 16, 16]               0\n",
      "       LeakyReLU-110           [8, 512, 16, 16]               0\n",
      "       LeakyReLU-111           [8, 512, 16, 16]               0\n",
      "       LeakyReLU-112           [8, 512, 16, 16]               0\n",
      "       LeakyReLU-113           [8, 512, 16, 16]               0\n",
      "       LeakyReLU-114           [8, 512, 16, 16]               0\n",
      "          Conv2d-115           [8, 512, 16, 16]       2,359,808\n",
      "     AdainResBlk-116           [8, 512, 16, 16]               0\n",
      "          Linear-117                  [8, 1024]          66,560\n",
      "  InstanceNorm2d-118           [8, 512, 16, 16]               0\n",
      "           AdaIN-119           [8, 512, 16, 16]               0\n",
      "       LeakyReLU-120           [8, 512, 16, 16]               0\n",
      "       LeakyReLU-121           [8, 512, 16, 16]               0\n",
      "       LeakyReLU-122           [8, 512, 16, 16]               0\n",
      "       LeakyReLU-123           [8, 512, 16, 16]               0\n",
      "       LeakyReLU-124           [8, 512, 16, 16]               0\n",
      "          Conv2d-125           [8, 512, 32, 32]       2,359,808\n",
      "          Linear-126                  [8, 1024]          66,560\n",
      "  InstanceNorm2d-127           [8, 512, 32, 32]               0\n",
      "           AdaIN-128           [8, 512, 32, 32]               0\n",
      "       LeakyReLU-129           [8, 512, 32, 32]               0\n",
      "       LeakyReLU-130           [8, 512, 32, 32]               0\n",
      "       LeakyReLU-131           [8, 512, 32, 32]               0\n",
      "       LeakyReLU-132           [8, 512, 32, 32]               0\n",
      "       LeakyReLU-133           [8, 512, 32, 32]               0\n",
      "          Conv2d-134           [8, 512, 32, 32]       2,359,808\n",
      "     AdainResBlk-135           [8, 512, 32, 32]               0\n",
      "          Linear-136                  [8, 1024]          66,560\n",
      "  InstanceNorm2d-137           [8, 512, 32, 32]               0\n",
      "           AdaIN-138           [8, 512, 32, 32]               0\n",
      "       LeakyReLU-139           [8, 512, 32, 32]               0\n",
      "       LeakyReLU-140           [8, 512, 32, 32]               0\n",
      "       LeakyReLU-141           [8, 512, 32, 32]               0\n",
      "       LeakyReLU-142           [8, 512, 32, 32]               0\n",
      "       LeakyReLU-143           [8, 512, 32, 32]               0\n",
      "          Conv2d-144           [8, 256, 64, 64]       1,179,904\n",
      "          Linear-145                   [8, 512]          33,280\n",
      "  InstanceNorm2d-146           [8, 256, 64, 64]               0\n",
      "           AdaIN-147           [8, 256, 64, 64]               0\n",
      "       LeakyReLU-148           [8, 256, 64, 64]               0\n",
      "       LeakyReLU-149           [8, 256, 64, 64]               0\n",
      "       LeakyReLU-150           [8, 256, 64, 64]               0\n",
      "       LeakyReLU-151           [8, 256, 64, 64]               0\n",
      "       LeakyReLU-152           [8, 256, 64, 64]               0\n",
      "          Conv2d-153           [8, 256, 64, 64]         590,080\n",
      "          Conv2d-154           [8, 256, 64, 64]         131,072\n",
      "     AdainResBlk-155           [8, 256, 64, 64]               0\n",
      "          Linear-156                   [8, 512]          33,280\n",
      "  InstanceNorm2d-157           [8, 256, 64, 64]               0\n",
      "           AdaIN-158           [8, 256, 64, 64]               0\n",
      "       LeakyReLU-159           [8, 256, 64, 64]               0\n",
      "       LeakyReLU-160           [8, 256, 64, 64]               0\n",
      "       LeakyReLU-161           [8, 256, 64, 64]               0\n",
      "       LeakyReLU-162           [8, 256, 64, 64]               0\n",
      "       LeakyReLU-163           [8, 256, 64, 64]               0\n",
      "          Conv2d-164         [8, 128, 128, 128]         295,040\n",
      "          Linear-165                   [8, 256]          16,640\n",
      "  InstanceNorm2d-166         [8, 128, 128, 128]               0\n",
      "           AdaIN-167         [8, 128, 128, 128]               0\n",
      "       LeakyReLU-168         [8, 128, 128, 128]               0\n",
      "       LeakyReLU-169         [8, 128, 128, 128]               0\n",
      "       LeakyReLU-170         [8, 128, 128, 128]               0\n",
      "       LeakyReLU-171         [8, 128, 128, 128]               0\n",
      "       LeakyReLU-172         [8, 128, 128, 128]               0\n",
      "          Conv2d-173         [8, 128, 128, 128]         147,584\n",
      "          Conv2d-174         [8, 128, 128, 128]          32,768\n",
      "     AdainResBlk-175         [8, 128, 128, 128]               0\n",
      "  InstanceNorm2d-176         [8, 128, 128, 128]             256\n",
      "       LeakyReLU-177         [8, 128, 128, 128]               0\n",
      "          Conv2d-178           [8, 3, 128, 128]             387\n",
      "(3, 128, 128, 64)\n",
      "================================================================\n",
      "Total params: 33,632,003\n",
      "Trainable params: 33,632,003\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 96.00\n",
      "Forward/backward pass size (MB): 6555.52\n",
      "Params size (MB): 128.30\n",
      "Estimated Total Size (MB): 6779.81\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "summary(Generator(128, 64).cuda(), [(3, 128, 128), ((64,))], batch_size=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [8, 128, 128, 128]           3,584\n",
      "            Conv2d-2         [8, 256, 128, 128]          32,768\n",
      "         LeakyReLU-3         [8, 128, 128, 128]               0\n",
      "         LeakyReLU-4         [8, 128, 128, 128]               0\n",
      "         LeakyReLU-5         [8, 128, 128, 128]               0\n",
      "         LeakyReLU-6         [8, 128, 128, 128]               0\n",
      "         LeakyReLU-7         [8, 128, 128, 128]               0\n",
      "            Conv2d-8         [8, 128, 128, 128]         147,584\n",
      "         LeakyReLU-9           [8, 128, 64, 64]               0\n",
      "        LeakyReLU-10           [8, 128, 64, 64]               0\n",
      "        LeakyReLU-11           [8, 128, 64, 64]               0\n",
      "        LeakyReLU-12           [8, 128, 64, 64]               0\n",
      "        LeakyReLU-13           [8, 128, 64, 64]               0\n",
      "           Conv2d-14           [8, 256, 64, 64]         295,168\n",
      "           ResBlk-15           [8, 256, 64, 64]               0\n",
      "           Conv2d-16           [8, 512, 64, 64]         131,072\n",
      "        LeakyReLU-17           [8, 256, 64, 64]               0\n",
      "        LeakyReLU-18           [8, 256, 64, 64]               0\n",
      "        LeakyReLU-19           [8, 256, 64, 64]               0\n",
      "        LeakyReLU-20           [8, 256, 64, 64]               0\n",
      "        LeakyReLU-21           [8, 256, 64, 64]               0\n",
      "           Conv2d-22           [8, 256, 64, 64]         590,080\n",
      "        LeakyReLU-23           [8, 256, 32, 32]               0\n",
      "        LeakyReLU-24           [8, 256, 32, 32]               0\n",
      "        LeakyReLU-25           [8, 256, 32, 32]               0\n",
      "        LeakyReLU-26           [8, 256, 32, 32]               0\n",
      "        LeakyReLU-27           [8, 256, 32, 32]               0\n",
      "           Conv2d-28           [8, 512, 32, 32]       1,180,160\n",
      "           ResBlk-29           [8, 512, 32, 32]               0\n",
      "        LeakyReLU-30           [8, 512, 32, 32]               0\n",
      "        LeakyReLU-31           [8, 512, 32, 32]               0\n",
      "        LeakyReLU-32           [8, 512, 32, 32]               0\n",
      "        LeakyReLU-33           [8, 512, 32, 32]               0\n",
      "        LeakyReLU-34           [8, 512, 32, 32]               0\n",
      "           Conv2d-35           [8, 512, 32, 32]       2,359,808\n",
      "        LeakyReLU-36           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-37           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-38           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-39           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-40           [8, 512, 16, 16]               0\n",
      "           Conv2d-41           [8, 512, 16, 16]       2,359,808\n",
      "           ResBlk-42           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-43           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-44           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-45           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-46           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-47           [8, 512, 16, 16]               0\n",
      "           Conv2d-48           [8, 512, 16, 16]       2,359,808\n",
      "        LeakyReLU-49             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-50             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-51             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-52             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-53             [8, 512, 8, 8]               0\n",
      "           Conv2d-54             [8, 512, 8, 8]       2,359,808\n",
      "           ResBlk-55             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-56             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-57             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-58             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-59             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-60             [8, 512, 8, 8]               0\n",
      "           Conv2d-61             [8, 512, 8, 8]       2,359,808\n",
      "        LeakyReLU-62             [8, 512, 4, 4]               0\n",
      "        LeakyReLU-63             [8, 512, 4, 4]               0\n",
      "        LeakyReLU-64             [8, 512, 4, 4]               0\n",
      "        LeakyReLU-65             [8, 512, 4, 4]               0\n",
      "        LeakyReLU-66             [8, 512, 4, 4]               0\n",
      "           Conv2d-67             [8, 512, 4, 4]       2,359,808\n",
      "           ResBlk-68             [8, 512, 4, 4]               0\n",
      "        LeakyReLU-69             [8, 512, 4, 4]               0\n",
      "           Conv2d-70             [8, 512, 1, 1]       4,194,816\n",
      "        LeakyReLU-71             [8, 512, 1, 1]               0\n",
      "           Conv2d-72               [8, 2, 1, 1]           1,026\n",
      "================================================================\n",
      "Total params: 20,735,106\n",
      "Trainable params: 20,735,106\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 1.50\n",
      "Forward/backward pass size (MB): 2422.06\n",
      "Params size (MB): 79.10\n",
      "Estimated Total Size (MB): 2502.66\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "summary(Discriminator(128, 2).cuda(), [(3, 128, 128), ((1,))], batch_size=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0, 1], device='cuda:0') tensor([0, 0], device='cuda:0')\n",
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Linear-1           [8, 3, 128, 512]          66,048\n",
      "              ReLU-2           [8, 3, 128, 512]               0\n",
      "            Linear-3           [8, 3, 128, 512]         262,656\n",
      "              ReLU-4           [8, 3, 128, 512]               0\n",
      "            Linear-5           [8, 3, 128, 512]         262,656\n",
      "              ReLU-6           [8, 3, 128, 512]               0\n",
      "            Linear-7           [8, 3, 128, 512]         262,656\n",
      "              ReLU-8           [8, 3, 128, 512]               0\n",
      "            Linear-9           [8, 3, 128, 512]         262,656\n",
      "             ReLU-10           [8, 3, 128, 512]               0\n",
      "           Linear-11           [8, 3, 128, 512]         262,656\n",
      "             ReLU-12           [8, 3, 128, 512]               0\n",
      "           Linear-13           [8, 3, 128, 512]         262,656\n",
      "             ReLU-14           [8, 3, 128, 512]               0\n",
      "           Linear-15            [8, 3, 128, 64]          32,832\n",
      "           Linear-16           [8, 3, 128, 512]         262,656\n",
      "             ReLU-17           [8, 3, 128, 512]               0\n",
      "           Linear-18           [8, 3, 128, 512]         262,656\n",
      "             ReLU-19           [8, 3, 128, 512]               0\n",
      "           Linear-20           [8, 3, 128, 512]         262,656\n",
      "             ReLU-21           [8, 3, 128, 512]               0\n",
      "           Linear-22            [8, 3, 128, 64]          32,832\n",
      "================================================================\n",
      "Total params: 2,495,616\n",
      "Trainable params: 2,495,616\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 1.50\n",
      "Forward/backward pass size (MB): 243.00\n",
      "Params size (MB): 9.52\n",
      "Estimated Total Size (MB): 254.02\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "summary(MappingNetwork(128).cuda(), [(3, 128, 128), ((1,))], batch_size=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [8, 128, 128, 128]           3,584\n",
      "            Conv2d-2         [8, 256, 128, 128]          32,768\n",
      "         LeakyReLU-3         [8, 128, 128, 128]               0\n",
      "         LeakyReLU-4         [8, 128, 128, 128]               0\n",
      "         LeakyReLU-5         [8, 128, 128, 128]               0\n",
      "         LeakyReLU-6         [8, 128, 128, 128]               0\n",
      "         LeakyReLU-7         [8, 128, 128, 128]               0\n",
      "            Conv2d-8         [8, 128, 128, 128]         147,584\n",
      "         LeakyReLU-9           [8, 128, 64, 64]               0\n",
      "        LeakyReLU-10           [8, 128, 64, 64]               0\n",
      "        LeakyReLU-11           [8, 128, 64, 64]               0\n",
      "        LeakyReLU-12           [8, 128, 64, 64]               0\n",
      "        LeakyReLU-13           [8, 128, 64, 64]               0\n",
      "           Conv2d-14           [8, 256, 64, 64]         295,168\n",
      "           ResBlk-15           [8, 256, 64, 64]               0\n",
      "           Conv2d-16           [8, 512, 64, 64]         131,072\n",
      "        LeakyReLU-17           [8, 256, 64, 64]               0\n",
      "        LeakyReLU-18           [8, 256, 64, 64]               0\n",
      "        LeakyReLU-19           [8, 256, 64, 64]               0\n",
      "        LeakyReLU-20           [8, 256, 64, 64]               0\n",
      "        LeakyReLU-21           [8, 256, 64, 64]               0\n",
      "           Conv2d-22           [8, 256, 64, 64]         590,080\n",
      "        LeakyReLU-23           [8, 256, 32, 32]               0\n",
      "        LeakyReLU-24           [8, 256, 32, 32]               0\n",
      "        LeakyReLU-25           [8, 256, 32, 32]               0\n",
      "        LeakyReLU-26           [8, 256, 32, 32]               0\n",
      "        LeakyReLU-27           [8, 256, 32, 32]               0\n",
      "           Conv2d-28           [8, 512, 32, 32]       1,180,160\n",
      "           ResBlk-29           [8, 512, 32, 32]               0\n",
      "        LeakyReLU-30           [8, 512, 32, 32]               0\n",
      "        LeakyReLU-31           [8, 512, 32, 32]               0\n",
      "        LeakyReLU-32           [8, 512, 32, 32]               0\n",
      "        LeakyReLU-33           [8, 512, 32, 32]               0\n",
      "        LeakyReLU-34           [8, 512, 32, 32]               0\n",
      "           Conv2d-35           [8, 512, 32, 32]       2,359,808\n",
      "        LeakyReLU-36           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-37           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-38           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-39           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-40           [8, 512, 16, 16]               0\n",
      "           Conv2d-41           [8, 512, 16, 16]       2,359,808\n",
      "           ResBlk-42           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-43           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-44           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-45           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-46           [8, 512, 16, 16]               0\n",
      "        LeakyReLU-47           [8, 512, 16, 16]               0\n",
      "           Conv2d-48           [8, 512, 16, 16]       2,359,808\n",
      "        LeakyReLU-49             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-50             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-51             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-52             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-53             [8, 512, 8, 8]               0\n",
      "           Conv2d-54             [8, 512, 8, 8]       2,359,808\n",
      "           ResBlk-55             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-56             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-57             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-58             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-59             [8, 512, 8, 8]               0\n",
      "        LeakyReLU-60             [8, 512, 8, 8]               0\n",
      "           Conv2d-61             [8, 512, 8, 8]       2,359,808\n",
      "        LeakyReLU-62             [8, 512, 4, 4]               0\n",
      "        LeakyReLU-63             [8, 512, 4, 4]               0\n",
      "        LeakyReLU-64             [8, 512, 4, 4]               0\n",
      "        LeakyReLU-65             [8, 512, 4, 4]               0\n",
      "        LeakyReLU-66             [8, 512, 4, 4]               0\n",
      "           Conv2d-67             [8, 512, 4, 4]       2,359,808\n",
      "           ResBlk-68             [8, 512, 4, 4]               0\n",
      "        LeakyReLU-69             [8, 512, 4, 4]               0\n",
      "           Conv2d-70             [8, 512, 1, 1]       4,194,816\n",
      "        LeakyReLU-71             [8, 512, 1, 1]               0\n",
      "           Linear-72                    [8, 64]          32,832\n",
      "           Linear-73                    [8, 64]          32,832\n",
      "================================================================\n",
      "Total params: 20,799,744\n",
      "Trainable params: 20,799,744\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 1.50\n",
      "Forward/backward pass size (MB): 2422.07\n",
      "Params size (MB): 79.34\n",
      "Estimated Total Size (MB): 2502.92\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "summary(StyleEncoder(128).cuda(), [(3, 128, 128), ((1,))], batch_size=8)"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}