ProGans.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "90BXN6vYSDQJ"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bf8GKvVySJkt"
      },
      "source": [
        "## Model Components"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zXa5c80UeYEe"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from math import log2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5deyHxALLnzd"
      },
      "source": [
        ""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H8aPG_6WnfGl"
      },
      "outputs": [],
      "source": [
        "class WSConv2d(nn.Module):\n",
        "  def __init__(self, in_channel, out_channel, kernel_size=3, stride=1, padding=1, gain=2):\n",
        "    super(WSConv2d, self).__init__()\n",
        "    self.conv = nn.Conv2d(in_channel, out_channel, kernel_size, stride, padding)\n",
        "    self.scale = (gain / (in_channel * (kernel_size ** 2))) ** 0.5\n",
        "    self.bias = self.conv.bias\n",
        "    self.conv.bias = None\n",
        "\n",
        "    #initializae conv layer\n",
        "    nn.init.normal_(self.conv.weight)\n",
        "    nn.init.zeros_(self.bias)\n",
        "\n",
        "  def forward(self,x):\n",
        "    return self.conv(x * self.scale) + self.bias.view(1, self.bias.shape[0], 1, 1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VaNjkpannn2E"
      },
      "outputs": [],
      "source": [
        "class PixelNorm(nn.Module):\n",
        "  def __init__(self):\n",
        "    super(PixelNorm, self).__init__()\n",
        "    self.epsilon = 1e-8\n",
        "\n",
        "  def forward(self, x):\n",
        "    return x / torch.sqrt(torch.mean(x**2, dim=1, keepdim=True) + self.epsilon)\n",
        "  pass"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vu6whD_ynoV0"
      },
      "outputs": [],
      "source": [
        "class ConvBlock(nn.Module):\n",
        "  def __init__(self, in_channel, out_channel, use_pixelnorm=True):\n",
        "    super(ConvBlock, self).__init__()\n",
        "    self.conv1 = WSConv2d(in_channel, out_channel)\n",
        "    self.conv2 = WSConv2d(out_channel, out_channel)\n",
        "    self.leaky = nn.LeakyReLU(0.2)\n",
        "    self.pn = PixelNorm()\n",
        "    self.use_pn = use_pixelnorm\n",
        "\n",
        "  def forward(self, x, use_pixelnorm=True):\n",
        "    x = self.leaky(self.conv1(x))\n",
        "    x = self.pn(x) if self.use_pn else x\n",
        "    x = self.leaky(self.conv2(x))\n",
        "    x = self.pn(x) if self.use_pn else x\n",
        "    return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x6Mav3Rhnt2g"
      },
      "source": [
        "## Model Classes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1K-89Fu3nonD"
      },
      "outputs": [],
      "source": [
        "class Generator(nn.Module):\n",
        "  def __init__(self, z_dim, in_channel, img_channel=3):\n",
        "    super(Generator,self).__init__()\n",
        "    self.initial = nn.Sequential(\n",
        "        PixelNorm(),\n",
        "        nn.ConvTranspose2d(z_dim, in_channel, 4, 1, 0),#1x1 -> 4x4\n",
        "        nn.LeakyReLU(0.2),\n",
        "        WSConv2d(in_channel, in_channel, kernel_size=3, stride=1, padding=1),\n",
        "        nn.LeakyReLU(0.2),\n",
        "        PixelNorm(),\n",
        "    )\n",
        "\n",
        "    self.initial_rgb = WSConv2d(in_channel, img_channel, kernel_size=1, stride=1, padding=0)\n",
        "    self.prog_blocks, self.rgb_layer = nn.ModuleList([]), nn.ModuleList([self.initial_rgb])\n",
        "\n",
        "    for i in range(len(factors) - 1):# was declared out of this classes\n",
        "      conv_in_c = int(in_channel * factors[i])\n",
        "      conv_out_c = int(in_channel * factors[i + 1])\n",
        "\n",
        "      self.prog_blocks.append(ConvBlock(conv_in_c, conv_out_c))\n",
        "      self.rgb_layer.append(WSConv2d(conv_out_c, img_channel, kernel_size=1, stride=1, padding=0))\n",
        "\n",
        "  def fade_in(self, alpha, upscaled, generated):\n",
        "    return torch.tanh(alpha * generated + (1 - alpha) * upscaled)\n",
        "\n",
        "  def forward(self, x, alpha, steps):\n",
        "    out = self.initial(x)\n",
        "\n",
        "    if steps == 0:\n",
        "      return self.initial_rgb(out)\n",
        "    \n",
        "    for step in range(steps):\n",
        "      upscaled = F.interpolate(out, scale_factor=2, mode=\"nearest\")\n",
        "      out = self.prog_blocks[step](upscaled) # use procbloc number STEP as a function for UPSCALED \"OUT\"\n",
        "    \n",
        "    final_upscaled = self.rgb_layer[steps - 1](upscaled)\n",
        "    final_out = self.rgb_layer[steps](out)\n",
        "    return self.fade_in(alpha, final_upscaled, final_out)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U-gSlCp9nxam"
      },
      "outputs": [],
      "source": [
        "from torch.nn.modules.activation import LeakyReLU\n",
        "class Discriminator(nn.Module):\n",
        "  def __init__(self, in_channel, img_channel=3):\n",
        "    super(Discriminator,self).__init__()\n",
        "    \n",
        "    self.prog_blocks, self.rgb_layers = nn.ModuleList([]), nn.ModuleList([])\n",
        "    self.Leaky = nn.LeakyReLU(0.2)\n",
        "\n",
        "    for i in range(len(factors) - 1, 0, -1):\n",
        "      conv_in_c = int(in_channel * factors[i])\n",
        "      conv_out_c = int(in_channel * factors[i - 1])     \n",
        "      self.prog_blocks.append(ConvBlock(conv_in_c, conv_out_c, use_pixelnorm=False))\n",
        "      self.rgb_layers.append(WSConv2d(img_channel, conv_in_c, kernel_size=1, stride=1, padding=0))\n",
        "\n",
        "    self.initial_rgb = WSConv2d(img_channel, in_channel, kernel_size=1, stride=1, padding=0)\n",
        "    self.rgb_layers.append(self.initial_rgb)\n",
        "    self.avg_pool = nn.AvgPool2d(kernel_size=2, stride=2)\n",
        "\n",
        "    #block for 4x4 resolution\n",
        "    self.final_block =  nn.Sequential(\n",
        "        WSConv2d(in_channel + 1, in_channel, kernel_size=3, stride=1, padding=1),\n",
        "        nn.LeakyReLU(0.2),\n",
        "        WSConv2d(in_channel, in_channel, kernel_size=4, stride=1, padding=0),\n",
        "        nn.LeakyReLU(0.2),\n",
        "        WSConv2d(in_channel, 1, kernel_size=1, stride=1, padding=0),\n",
        "    )\n",
        "\n",
        "  def fade_in(self, alpha, downscaled, out):\n",
        "    return alpha * out + (1 - alpha) * downscaled\n",
        "\n",
        "  def minibatch_std(self, x):\n",
        "    batch_statistics = torch.std(x, dim=0).mean().repeat(x.shape[0], 1, x.shape[2], x.shape[3])\n",
        "    return torch.cat([x, batch_statistics], dim=1)\n",
        "\n",
        "  def forward(self, x, alpha, step):\n",
        "    cur_step = len(self.prog_blocks) - step\n",
        "    out = self.Leaky(self.rgb_layers[cur_step](x))\n",
        "\n",
        "    if step == 0:\n",
        "      \n",
        "      out = self.minibatch_std(out)\n",
        "      return self.final_block(out).view(out.shape[0], -1)\n",
        "\n",
        "    downscaled = self.Leaky(self.rgb_layers[cur_step + 1](self.avg_pool(x)))\n",
        "    out = self.avg_pool(self.prog_blocks[cur_step](out))\n",
        "    out = self.fade_in(alpha, downscaled, out)\n",
        "\n",
        "    for step in range(cur_step + 1, len(self.prog_blocks)):\n",
        "      out = self.prog_blocks[step](out)\n",
        "      out = self.avg_pool(out)\n",
        "\n",
        "    out = self.minibatch_std(out)\n",
        "    return self.final_block(out).view(out.shape[0], -1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RoGf570jQCk_"
      },
      "source": [
        "# Train"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ua1674O2Zgl_"
      },
      "source": [
        "## Library"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gOG57qZyoazv"
      },
      "outputs": [],
      "source": [
        "from torch.utils.tensorboard import SummaryWriter\n",
        "from torch.utils import tensorboard\n",
        "from math import log2\n",
        "from tqdm import tqdm\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import torchvision\n",
        "import numpy as np\n",
        "import random\n",
        "import torch\n",
        "import cv2\n",
        "import os\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ob2HSmfdQc41"
      },
      "source": [
        "## Config & Utils"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gMlldguZb6yv"
      },
      "source": [
        "### Config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eHixgjxrQex5"
      },
      "outputs": [],
      "source": [
        "\n",
        "class configure():\n",
        "  def __init__(self):\n",
        "    self.START_TRAIN_AT_IMG_SIZE = 4\n",
        "    self.DATASET = '/content/drive/MyDrive/Datasets/resized_ganime_fullbody_cleaned_ds/'\n",
        "    self.DATASET_type = 'img' #numpy, img\n",
        "    self.CHECKPOINT_GEN = \"/content/drive/MyDrive/Datasets/save/128x128_progans_fullbody/working/gen.pth\"\n",
        "    self.CHECKPOINT_DIS = \"/content/drive/MyDrive/Datasets/save/128x128_progans_fullbody/working/dis.pth\"\n",
        "    self.DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "    self.SAVE_MODEL = True\n",
        "    self.LOAD_MODEL = False\n",
        "    self.LEARNING_RATE = 1e-8\n",
        "    self.ALPHA = 1e-5\n",
        "    self.BETAS = (0.0, 0.99)\n",
        "    self.BATCH_SIZES = [16, 16, 16, 8, 8, 8, 8, 4, 2]\n",
        "    self.CHANNELS_IMG = 3\n",
        "    self.Z_DIM = 1024  # should be 512 in original paper\n",
        "    self.IN_CHANNELS = 1028  # should be 512 in original paper\n",
        "    self.CRITIC_ITERATIONS = 1\n",
        "    self.LAMBDA_GP = 10\n",
        "    self.PROGRESSIVE_EPOCHS = [30] * len(self.BATCH_SIZES)\n",
        "    self.FIXED_NOISE = torch.randn(8, self.Z_DIM, 1, 1).to(self.DEVICE)\n",
        "    self.NUM_WORKERS = 4\n",
        "    self.SEED = 0\n",
        "\n",
        "config = configure()\n",
        "visualize_noise = torch.randn(16, config.IN_CHANNELS, device=config.DEVICE, dtype=torch.float) #don't touch this noise since we're using it for visualizations"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y4EcycwqedmY"
      },
      "outputs": [],
      "source": [
        "# factors = [1, 1, 1, 1, 1 / 2, 1 / 4, 1 / 8, 1 / 16, 1 / 32] \n",
        "factors = [1, 1, 1, 1, 1 / 2, 1 / 4, 1 / 8] # for 256x256"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tivz1LTdRibc"
      },
      "source": [
        "### Utils"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "368R0_pJmtjN"
      },
      "outputs": [],
      "source": [
        "def plotter(data):\n",
        "  data = torch.moveaxis(data*127.5+127.5,1,-1).numpy().astype(int)\n",
        "  fig,ax = plt.subplots(2,2,figsize=(5,5))\n",
        "  for i in range(2):\n",
        "    for j in range(2):\n",
        "      ax[i,j].imshow(data[i+2*j])\n",
        "      ax[i,j].axis('off')\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dSwwukUFu6hv"
      },
      "outputs": [],
      "source": [
        "def random_sample(data):\n",
        "  for images in data:\n",
        "    for img in images:\n",
        "      print(images.shape)\n",
        "      img = np.transpose(img.numpy(), (1, 2, 0))\n",
        "      print(img.shape)\n",
        "      plt.imshow(img)\n",
        "      \n",
        "      break\n",
        "    break"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ozdu7ghWVoph"
      },
      "outputs": [],
      "source": [
        "def set_seed(seed):\n",
        "    seed = config.SEED\n",
        "    random.seed(seed)\n",
        "    np.random.seed(seed)\n",
        "    torch.manual_seed(seed)\n",
        "    torch.cuda.manual_seed(seed)\n",
        "    torch.cuda.manual_seed_all(seed)\n",
        "    torch.backends.cudnn.deterministic = True\n",
        "    torch.backends.cudnn.benchmark = False"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u58b0XUIVJNe"
      },
      "outputs": [],
      "source": [
        "def load_checkpoint(checkpoint_file, model, optimizer, lr):\n",
        "    print(\"=> Loading checkpoint\")\n",
        "    checkpoint = torch.load(checkpoint_file, map_location=\"cuda\")\n",
        "    model.load_state_dict(checkpoint[\"state_dict\"])\n",
        "    optimizer.load_state_dict(checkpoint[\"optimizer\"])\n",
        "\n",
        "    # If we don't do this then it will just have learning rate of old checkpoint\n",
        "    # and it will lead to many hours of debugging \\:\n",
        "    for param_group in optimizer.param_groups:\n",
        "        param_group[\"lr\"] = lr"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uOxC9RRqVRbF"
      },
      "outputs": [],
      "source": [
        "def save_checkpoint(model, optimizer, filename=\"my_checkpoint.pth.tar\"):\n",
        "    print(\"=> Saving checkpoint\")\n",
        "    checkpoint = {\n",
        "        \"state_dict\": model.state_dict(),\n",
        "        \"optimizer\": optimizer.state_dict(),\n",
        "    }\n",
        "    torch.save(checkpoint, filename)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y6qDHhuxVhhV"
      },
      "outputs": [],
      "source": [
        "# Print losses occasionally and print to tensorboard\n",
        "def plot_to_tensorboard(\n",
        "    writer, loss_critic, loss_gen, real, fake, tensorboard_step\n",
        "):\n",
        "    writer.add_scalar(\"Loss Critic\", loss_critic, global_step=tensorboard_step)\n",
        "\n",
        "    with torch.no_grad():\n",
        "        # take out (up to) 8 examples to plot\n",
        "        img_grid_real = torchvision.utils.make_grid(real[:8], normalize=True)\n",
        "        img_grid_fake = torchvision.utils.make_grid(fake[:8], normalize=True)\n",
        "        writer.add_image(\"Real\", img_grid_real, global_step=tensorboard_step)\n",
        "        writer.add_image(\"Fake\", img_grid_fake, global_step=tensorboard_step)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_15CorM2QEE0"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eG3NYtgDSXGB"
      },
      "source": [
        "### Data Loader"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BrAfvyCimeAs"
      },
      "outputs": [],
      "source": [
        "def dataloader(size,batch_size):\n",
        "  np_array = []\n",
        "  loader_size = None\n",
        "\n",
        "  if config.DATASET_type == 'img':\n",
        "    print(f'Loading Data at size {size}...')\n",
        "    for image in tqdm(os.listdir(config.DATASET)):\n",
        "      img_path = config.DATASET + image\n",
        "      img = cv2.imread(img_path)\n",
        "      img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
        "\n",
        "      to_size = (size,size)\n",
        "      img = cv2.resize(img,to_size)\n",
        "      np_array.append(img)\n",
        "\n",
        "  np_array = np.array(np_array)/127.5-1\n",
        "  loader_size = np_array.shape[0]#number of images\n",
        "\n",
        "  temp0 = np.moveaxis(np_array,-1,1)[0:]#[0:x] -> x = limit of datasets\n",
        "  del np_array\n",
        "  temp = torch.tensor(temp0).float()\n",
        "  del temp0\n",
        "  train_loader = torch.utils.data.DataLoader(temp, \n",
        "                                           batch_size=batch_size, \n",
        "                                           num_workers=config.NUM_WORKERS, \n",
        "                                           shuffle=True,\n",
        "                                           )\n",
        "  del temp\n",
        "\n",
        "  return train_loader, loader_size"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IbxAojXcVW8c"
      },
      "source": [
        "### loss & all"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yCSewAYXVc3u"
      },
      "outputs": [],
      "source": [
        "def gradient_penalty(critic, real, fake, alpha, train_step, device=\"cpu\"):\n",
        "    BATCH_SIZE, C, H, W = real.shape\n",
        "    beta = torch.rand((BATCH_SIZE, 1, 1, 1)).repeat(1, C, H, W).to(device)\n",
        "    interpolated_images = real * beta + fake.detach() * (1 - beta)\n",
        "    interpolated_images.requires_grad_(True)\n",
        "\n",
        "    # Calculate critic scores\n",
        "    mixed_scores = critic(interpolated_images, alpha, train_step)\n",
        "\n",
        "    # Take the gradient of the scores with respect to the images\n",
        "    gradient = torch.autograd.grad(\n",
        "        inputs=interpolated_images,\n",
        "        outputs=mixed_scores,\n",
        "        grad_outputs=torch.ones_like(mixed_scores),\n",
        "        create_graph=True,\n",
        "        retain_graph=True,\n",
        "    )[0]\n",
        "    gradient = gradient.view(gradient.shape[0], -1)\n",
        "    gradient_norm = gradient.norm(2, dim=1)\n",
        "    gradient_penalty = torch.mean((gradient_norm - 1) ** 2)\n",
        "    return gradient_penalty"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sJq7Dd5sbRIX"
      },
      "source": [
        "### Trainer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aup-GwsKbUyI"
      },
      "outputs": [],
      "source": [
        "def trainer(\n",
        "    critic,\n",
        "    gen,\n",
        "    loader,\n",
        "    data_len,\n",
        "    step,\n",
        "    alpha,\n",
        "    opt_critic,\n",
        "    opt_gen,\n",
        "    tensorboard_step,\n",
        "    writer,\n",
        "    scaler_gen,\n",
        "    scaler_critic,\n",
        "):\n",
        "    loop = tqdm(loader, leave=True)\n",
        "    for batch_idx, real in enumerate(loop):\n",
        "        real = real.to(config.DEVICE)\n",
        "        cur_batch_size = real.shape[0]\n",
        "\n",
        "        \n",
        "        \n",
        "\n",
        "        # Train Critic: max E[critic(real)] - E[critic(fake)] <-> min -E[critic(real)] + E[critic(fake)]\n",
        "        # which is equivalent to minimizing the negative of the expression\n",
        "        noise = torch.randn(cur_batch_size, config.Z_DIM, 1, 1).to(config.DEVICE)\n",
        "\n",
        "        with torch.cuda.amp.autocast():\n",
        "            fake = gen(noise, alpha, step)\n",
        "            critic_real = critic(real, alpha, step)\n",
        "            critic_fake = critic(fake.detach(), alpha, step)\n",
        "            gp = gradient_penalty(critic, real, fake, alpha, step, device=config.DEVICE)\n",
        "            loss_critic = (\n",
        "                -(torch.mean(critic_real) - torch.mean(critic_fake))\n",
        "                + config.LAMBDA_GP * gp\n",
        "                + (0.001 * torch.mean(critic_real ** 2))\n",
        "            )\n",
        "\n",
        "        opt_critic.zero_grad()\n",
        "        scaler_critic.scale(loss_critic).backward()\n",
        "        scaler_critic.step(opt_critic)\n",
        "        scaler_critic.update()\n",
        "\n",
        "        # Train Generator: max E[critic(gen_fake)] <-> min -E[critic(gen_fake)]\n",
        "        with torch.cuda.amp.autocast():\n",
        "            gen_fake = critic(fake, alpha, step)\n",
        "            loss_gen = -torch.mean(gen_fake)\n",
        "\n",
        "        opt_gen.zero_grad()\n",
        "        scaler_gen.scale(loss_gen).backward()\n",
        "        scaler_gen.step(opt_gen)\n",
        "        scaler_gen.update()\n",
        "\n",
        "        # Update alpha and ensure less than 1\n",
        "        alpha += cur_batch_size / (\n",
        "            (config.PROGRESSIVE_EPOCHS[step] * 0.5) * data_len\n",
        "        )\n",
        "        alpha = min(alpha, 1)\n",
        "\n",
        "        if batch_idx % 100 == 0:\n",
        "            with torch.no_grad():\n",
        "                fixed_fakes = gen(config.FIXED_NOISE, alpha, step) * 0.5 + 0.5\n",
        "            plot_to_tensorboard(\n",
        "                writer,\n",
        "                loss_critic.item(),\n",
        "                loss_gen.item(),\n",
        "                real.detach(),\n",
        "                fixed_fakes.detach(),\n",
        "                tensorboard_step,\n",
        "            )\n",
        "            tensorboard_step += 1\n",
        "\n",
        "            plotter(fixed_fakes.detach().cpu())\n",
        "\n",
        "        loop.set_postfix(\n",
        "            gp=gp.item(),\n",
        "            loss_gen=loss_gen.item(),\n",
        "            loss_dis=loss_critic.item(),\n",
        "        )\n",
        "\n",
        "    return tensorboard_step, alpha"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6LJVA0FMXZg2"
      },
      "source": [
        "## Train"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C5-Xv7UJS8o6"
      },
      "source": [
        "### Model Init"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ta6l1nrtZKCZ"
      },
      "outputs": [],
      "source": [
        "#optimizers\n",
        "seed_size = config.Z_DIM\n",
        "in_channels = config.IN_CHANNELS\n",
        "device = torch.device(config.DEVICE)\n",
        "\n",
        "gen = Generator(seed_size, in_channels)\n",
        "dis = Discriminator(in_channels)\n",
        "\n",
        "gen.to(device)\n",
        "dis.to(device)\n",
        "\n",
        "gen_optimizer = torch.optim.Adam(\n",
        "    dis.parameters(), lr=config.LEARNING_RATE, betas=config.BETAS\n",
        ")\n",
        "dis_optimizer = torch.optim.Adam(\n",
        "    dis.parameters(), lr=config.LEARNING_RATE, betas=config.BETAS\n",
        ")\n",
        "\n",
        "dis_scaler = torch.cuda.amp.GradScaler()\n",
        "gen_scaler = torch.cuda.amp.GradScaler()\n",
        "\n",
        "#tensorboard plotting\n",
        "writer = SummaryWriter(f\"logs/gan1\")\n",
        "\n",
        "if config.LOAD_MODEL:\n",
        "  load_checkpoint(\n",
        "      config.CHECKPOINT_GEN, gen, gen_optimizer, config.LEARNING_RATE\n",
        "  )\n",
        "  load_checkpoint(\n",
        "      config.CHECKPOINT_DIS, dis, dis_optimizer, config.LEARNING_RATE\n",
        "  )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TXLLaY7dUJ2E"
      },
      "source": [
        "### Train"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9FQatmpHpclJ"
      },
      "outputs": [],
      "source": [
        "num_epochs = config.PROGRESSIVE_EPOCHS\n",
        "batch_size = config.BATCH_SIZES\n",
        "reg_lambda = config.LAMBDA_GP\n",
        "\n",
        "torch.backends.cudnn.benchmarks = True #idk why they said it improve performance"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 235
        },
        "id": "OZhN4tDFp1WL",
        "outputId": "f33da795-6b96-41c1-e78c-2e7030f71c78"
      },
      "outputs": [
        {
          "output_type": "error",
          "ename": "NameError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-2-ffeab97f918d>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m#enable traning\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mgen\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m \u001b[0mdis\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0mtensorboard_step\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mNameError\u001b[0m: name 'gen' is not defined"
          ]
        }
      ],
      "source": [
        "#enable traning\n",
        "gen.train()\n",
        "dis.train()\n",
        "\n",
        "tensorboard_step = 0\n",
        "\n",
        "step = int(log2(config.START_TRAIN_AT_IMG_SIZE / 4))\n",
        "\n",
        "for num_epochs in config.PROGRESSIVE_EPOCHS[step:]:\n",
        "  alpha = config.ALPHA\n",
        "  cur_batch_size = config.BATCH_SIZES[step]\n",
        "  loader, data_len = dataloader(4*2 ** step, cur_batch_size)\n",
        "\n",
        "  print(f'Images Size : {4*2 ** step}')\n",
        "\n",
        "  for epoch in range(num_epochs):\n",
        "            print(f\"Epoch [{epoch+1}/{num_epochs}]\")\n",
        "\n",
        "            tensorboard_step, alpha = trainer(\n",
        "                dis,\n",
        "                gen,\n",
        "                loader,\n",
        "                data_len,\n",
        "                step,\n",
        "                alpha,\n",
        "                dis_optimizer,\n",
        "                gen_optimizer,\n",
        "                tensorboard_step,\n",
        "                writer,\n",
        "                gen_scaler,\n",
        "                dis_scaler,\n",
        "            )\n",
        "\n",
        "            if config.SAVE_MODEL:\n",
        "                # if epoch % 2:#save every 2 epochs\n",
        "                save_checkpoint(gen, gen_optimizer, filename=config.CHECKPOINT_GEN)\n",
        "                save_checkpoint(dis, dis_optimizer, filename=config.CHECKPOINT_DIS)\n",
        "\n",
        "  step += 1  # progress to the next img size\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "14RgnVqrIFA8"
      },
      "source": [
        "## Show Tensorboard"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vOSNSoF6IXEz"
      },
      "outputs": [],
      "source": [
        "%load_ext tensorboard"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ILlDl2PXK29h"
      },
      "outputs": [],
      "source": [
        "%tensorboard --logdir 'logs/gan1'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7hrtvJMR53Na"
      },
      "source": [
        "# Test"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## test model code"
      ],
      "metadata": {
        "id": "RggXh41XhN3y"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j466ksDJn1xP"
      },
      "outputs": [],
      "source": [
        "def test_code():\n",
        "  z_dim = 100\n",
        "  in_channel = 256\n",
        "  imgchannel = 3\n",
        "  gen = Generator(z_dim, in_channel, imgchannel)\n",
        "  dis = Discriminator(in_channel, imgchannel)\n",
        "\n",
        "  # for image_size in [4, 8, 16, 32, 64, 128, 256, 512, 1024]:\n",
        "  for image_size in [4, 8, 16, 32, 64, 128, 256]:\n",
        "    num_steps = int(log2(image_size / 4))\n",
        "    x = torch.randn((1, z_dim, 1, 1))\n",
        "    z = gen(x, 0.5, steps=num_steps)\n",
        "    assert z.shape == (1, 3, image_size, image_size)\n",
        "    out = dis(z, alpha=0.5, step=num_steps)\n",
        "    assert out.shape == (1, 1)\n",
        "    print(f'Success At Images Size of {image_size}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3IRVZ3ta7LE_"
      },
      "outputs": [],
      "source": [
        "#test_code()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "40r49-TFSxLu"
      },
      "outputs": [],
      "source": [
        ""
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Test Model Result"
      ],
      "metadata": {
        "id": "5YL_5zcjhP6r"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "load_checkpoint(\n",
        "    config.CHECKPOINT_GEN, gen, gen_optimizer, config.LEARNING_RATE\n",
        ")\n",
        "load_checkpoint(\n",
        "    config.CHECKPOINT_DIS, dis, dis_optimizer, config.LEARNING_RATE\n",
        ")\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "B2ZldEybhSM-",
        "outputId": "4dbd6fa8-954b-4cb4-bfa7-b6c293c68148"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "=> Loading checkpoint\n",
            "=> Loading checkpoint\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "noise = torch.randn(4, config.Z_DIM, 1, 1).to(config.DEVICE)\n",
        "img_path = config.DATASET\n",
        "img_path = os.listdir(img_path)"
      ],
      "metadata": {
        "id": "PRXayPORiMiJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "img_path[0]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "ggZ3_aLmv5Wk",
        "outputId": "805fc3e0-0738-4d7a-c522-f19582bb5fc1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'9039.png'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 23
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def test_smallbatch_img(img_path, to_size, num):\n",
        "  np_array = []\n",
        "  \n",
        "  for i in range(num):\n",
        "    img = cv2.imread(config.DATASET + img_path[i])\n",
        "    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
        "    img = cv2.resize(img,(to_size,to_size))\n",
        "    np_array.append(img)\n",
        "\n",
        "  np_array = np.array(np_array)/127.5-1\n",
        "  loader_size = np_array.shape[0]#number of images\n",
        "\n",
        "  temp0 = np.moveaxis(np_array,-1,1)[0:]#[0:x] -> x = limit of datasets\n",
        "  temp = torch.tensor(temp0).float()\n",
        "  train_loader = torch.utils.data.DataLoader(temp, \n",
        "                                      batch_size=4, \n",
        "                                      num_workers=config.NUM_WORKERS,)\n",
        "  return train_loader"
      ],
      "metadata": {
        "id": "LYdTmR7fuwCI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "fake = gen(config.FIXED_NOISE, 0.5, 2) * 0.5 + 0.5\n",
        "plotter(fake.detach().cpu())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 303
        },
        "id": "51Y7FaWOibci",
        "outputId": "cd8c5e28-c4eb-4a20-deb0-eed622d7dd91"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 360x360 with 4 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAASIAAAEeCAYAAAAuBhhgAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAXgElEQVR4nO3dSY9c53XG8XNvzdVVXT2PFEmJpElKlGRZjmQBCYIsbCiIN/kC+S7OLh8mQZJFEiCbOECMxJEHQbIscZBIUexms+eh5rr3ZpH1eS5gtHgE+P/bHr41vn1YwPvgvElRFAYAkdLoFwAANCIA4WhEAMLRiACEoxEBCEcjAhCuqoo/+8d/kGf7yfmeW2tduSGfuH34lqwf1n8h67ad+6/r9W29tujK8o3muqwPc79/H0yfyrXpeEPWC8tkvTpXd2u1wZJc+37F/8zMzN7r1BL5Dy7Zzx78Tu6v9JWxW2tX78vHbiR6f701uCnrxdT/rNqjx3Ltg+OZrD85/Y2s21V/fW/hTC6dZg1ZXyz57XH9kf+ZP2jo/bWx9n1Z/6vVN9z9xS8iAOFoRADC0YgAhKMRAQhHIwIQjkYEIByNCEA4mSNKN9+Ti2vHv3drsy/bcu109ETW7Z2Brg8X3VIxOJJLi6aOyzw511mffOfCrbXG8iO1zvxQ1lfqfo7DzOx49sKt1YY/kmvrzZGsW6em65csWe3Iet7/xq0N6025dtTS3/GzQue1Vsz/nM/zvn7uwt+bZmbF4FVZr5z4WbGkot93IvJPZmbT6gNZP+zsuLX+gd6bRcv/vszMbPUNt8QvIgDhaEQAwtGIAISjEQEIRyMCEI5GBCCcPGsuHutjyLxyz60lLX3UN0v8I0ozs9fPP5L1x+v+8f7kdE2uzYYTWa8f6aPf+qdfu7Uk08ejb1zXx8aNLf2Zdwp/RMSTdX2s/ORsTtb1EIdvwUwfRacN/3ssTvUolyTX0YB2Wx9zH2f+93BR09GUSVfv7WSxJetZzV8/eab3br46lfW9C72/Tur+aJ/Rhh6t8vnZNVn/qajxiwhAOBoRgHA0IgDhaEQAwtGIAISjEQEIRyMCEE7niM6ey8VZa8WtLayVjCu4X5LlmVuW9c2Fc7fWb2/KtReVLVnvrugrWb6/8Ypbu392INfWK/IGHZs99x/bzKy27H9ljYnOzhyNv2P/76Q6C1Tkfl6nWdfv1S70e73Rrsh6L/VzRk9SPeqlEFkvM7NfDXSeayoeP32uM261oX7u/ESPuFnq+nv/sKPzea2K/lyU79jOBPDHiEYEIByNCEA4GhGAcDQiAOFoRADC0YgAhJMH/7XTY7m4OfEzLytdndOY6EiCFTc+lPXugp/H6eyeybWjZT0TaHum5/pMxK082Z6ec/PNCz2Lxv7spiy/WPJzJFlLXxfUa+p8lJn+zi5bzfTcnrziX7vTK/kvtF7V3/Gjsd7b10d+zu2kvyvXHua3ZL35X/p9N/x4nt2+qfN5rUy/7/27OsPUzfzPfONU7+3OhX/NVhl+EQEIRyMCEI5GBCAcjQhAOBoRgHA0IgDhaEQAwskcUbNak4uLmj/7pJro2STVkhxR2tY9Mh/5WYz9/aFc21/Uj32Q6Vk39U0/LzEb6gxS/YXOgYxyPa9oXPg5kX5WlsuR5Zeuqkfr2Ey8nXyq5+40Sh77sND37p03/HlY43Rerh3X9Hfcf1N/EXNT/ztubevXXTF951qlrvf+5xN//1UP9ete7Jdk5AR+EQEIRyMCEI5GBCAcjQhAOBoRgHA0IgDh5Bn7+fv+SAAzs+Sxf1Sdz+kxC3Mjfbzfmejz1+bAX7/5UB+/P9zVj71YLRnTMPGvI8rv/KVcW7+lj36Tlv6/4eqJf7z6cF6vHVa+W//vLD/U38Pshn9cfFFym1BZDGLU35P16ol/hF55MpVrG039Hc891uNY7iz54zb2vtDH99mx/rt6fVVHco4f+ONPanf0GJDtUz1+x2zBrXy3diaAP0o0IgDhaEQAwtGIAISjEQEIRyMCEI5GBCCcDh30dF7CrvtZjcrcQC5dGumRAvOnun4h8jTjc50DWj3V76uW6DEOae6PoKh3dX6lvayzM7WWro8G/v8daUXnPGam6y/7/6WpvpXHlkTcZrShR71cbev99yB/LOvTHbG3L/T+yXb0/ts811me+cTPMM36em021NcJnfb0+JTzRX+PzObLMkx6/3wgavwiAhCORgQgHI0IQDgaEYBwNCIA4WhEAMLRiACEkzmidFPnTqqbfqah2tQ9rjOdk/WniX7us0U/y9Hv68c+fqGzFHVxnYuZ2fKhX7u225Vr24s6J7TZ1fV6za8fPRzJtV/paI3ZD/Vrv2ynDX310lHhb8/kXO+v3Vxnedb6etbWaWfVrdWXn8i1zRc6bzP7RJZtdOFvsJ2Bf82RmdlwpnNGp2v6cxs99euDN0/k2vxdHUtU+EUEIByNCEA4GhGAcDQiAOFoRADC0YgAhKMRAQgnD/6rJfNiauv+TJjtVN9rtlzTWYuvip6sF03/8cclWZ36C10vZjoPURMZqUTHOOxirJ97ONb5qWTsz8lZ7JfMOsq/W//vTI50nsu6bb/W0Z9TZVfPs2rvrcl6s+rv7a3FI7nWvlmR5cPZU1lPzc9XjU3PyhrXSuYVNXW9K+6DmzZ1jqhd1X/zZn5267u1MwH8UaIRAQhHIwIQjkYEIByNCEA4GhGAcPKcevNjfTXOldfESIpXluTaoj4v65/V9RgHm/j17JXf6bXXrspy79+/kPVk2HJrtcqWXDu/p4+sD8d6lMfhxP+/Y2dTH9/nnz6XdXv/pq5fsvr8nqxna/72XDjWx/Nv7OnP8fH8C1lvNCdu7bo+nbds8bGsD5b0+z5o+Ufs54v676bI9PH8s4/1i69O/Ped51fk2vpM/12ZuB6KX0QAwtGIAISjEQEIRyMCEI5GBCAcjQhAOBoRgHAyR1RPfy0Xj47fdWtPlv2sjZnZ7ZJxBfWqzomk5j/+tKFHPGSp7r9L7/pZCjOz3q7/sW0/02s7pnNEn7zQVxlNzv3Prfe2/sxObuhc2MtWWzyV9Vm26dbSls6ZHXf0/qvP6ZxbWvhjQPaO9ZiY2VBnmPbW9R45F/uzONTfYbag/66SM733i5E/XqWod+TaNLuQdbn2D14JAJeERgQgHI0IQDgaEYBwNCIA4WhEAMLRiACEk4GIfO+aXPxs37/aZPqxzsOcNnXmoPOhzpgMRdbi5spjufYofU3Wx7t+fsXMbL/wcyCVse7tG0M/n2JmNj/UGZX6wL/SZa+YyrWv3j2UdbNbJfXLNf1Mz27KW/6VU+MVPXvp+VRfu5Otfi3rlZm/fmWxK9c+7eoc0VnrHVm3ws+DDacl+bpUXMFkZgsl+b1Zxc8Krd7XGabr2YGs2w/9Er+IAISjEQEIRyMCEI5GBCAcjQhAOBoRgHA0IgDhZGilvbksFw/b/uyS/AudIzo0Pw9jZnZR6KyGTfy8xGy6L5e2qzpjkuT6tc3Ojt3adL0n19pQ12df6te22PA/84uS/Mr+sZ6F9EHJS79sxZUjXV/zZw4lxYZcO9BxLKuu6EzVYn/BrS31dJ7meDon6yeF/qCbNTFLaex//2Zmud4+lnT1a6+KLdSe93NdZmaDwUA/ucAvIgDhaEQAwtGIAISjEQEIRyMCEI5GBCCcPOR8u69HVvym5h8lVjf00f/NkiPM/u91j5w2/TEQ7w3uyrXnizpaUG3q4//Kl/4YkN6qPpKuvqZjCbUDPcbBxBSIBXHFkplZ2r+uH/sl+9MH57L+5MI/ot96W4+z6OhTbrv58LasN0/E3t/W3+F6U2cHPhvPy/pBw9/7B8d6/EmyIss2f0/Xr537j9+p6PjH9pm+xkvhFxGAcDQiAOFoRADC0YgAhKMRAQhHIwIQjkYEIJwMPHwz0qGE2u/8MQ3vJzpLsdeQZTtu6XEF++bXn3V0PqV45F+ZYmY2muprlLa3/ZDKvdeW5Nqluh4/sbKgR3lkF35+ZvyFDs989YmffzIz+9Fd//v8Nmws/EDW5479/ye3PtfX6pyOdP2rV1+V9a3Ev5qpPtaPXZz1ZT2v6izQydTf29MNvdb0V2wn3+h6z59wY1c+0n+0+U19RZjCLyIA4WhEAMLRiACEoxEBCEcjAhCORgQgHI0IQLikKHReBwC+bfwiAhCORgQgHI0IQDgaEYBwNCIA4WhEAMLRiACEoxEBCEcjAhCORgQgHI0IQDgaEYBwNCIA4WhEAMLRiACEoxEBCEcjAhCORgQgHI0IQDgaEYBwNCIA4WhEAMLRiACEq6ri3w6eykvP6om/PK2cyCdOq2uyPs77sm6DVbdUO5zKpcnTsaznvX+V9Vbhv7cfbN+Ra/cbG7I+LTqy/r206dbO5vyamdnhSH8uf91eT+Q/uGR/9/Njub+q94ZubZKsy8ee1fX/sY1PJ7I+P83d2srBsVzbyfRj/7qpv+O8W3Fr2du7cq2dnMvy0pz/2GZmc/Pzbu0vJodybSfT38nqwmvu/uIXEYBwNCIA4WhEAMLRiACEoxEBCEcjAhCORgQgnMwR2Z6f1TEzy+Yzt1aczem1izVZzxslPXJy4Zfq/usyM0sX9EMnS/pjGWTbbq1d6Pe9fPaprD8c35L1F3X/tR2MdQzoWWtT1l+2acXP6piZtTI/0zKpyAiSVVP92N2OzvrMEj9vM67pvNZ40NKPnfv5KDOzYui/9mLna7nW9vy/CzOzk3uLsj593nZrv7gYyLW9Db33fypq/CICEI5GBCAcjQhAOBoRgHA0IgDhaEQAwulz6lz3qUrdP+LMh/oYcfZCHzUn7T1ZLw789dWvZ3Lt8p/oESPFxpuyfrXwR3n8b6HHMAwH+iO/NtLjUYqef6SdpSO51kqOtF+29s26rNd7/tiSi0SvreX6vZ4vHch6MvIfv5o9lGuzvn5t7Z7eI4n5Y2Yuavr43V69Lct5T9cvFvy/+dGe/kznT/Txvi37JX4RAQhHIwIQjkYEIByNCEA4GhGAcDQiAOFoRADCyVBL0StZPRDjCvRNMWazkhEQ5/rqkuHAn+VR6+vxJeNnuj55tSHrz80fITHN9XiTdLYk6/XUzwmZmd0e+19Za04/98VYj68wPb3i0mWZvtqmP/HHZeTNkhxayZiQ+kSvb+/7668+0XmtRw39HVcP9RiQ4W0xxqb7XK61o5IboYqrupz6+ytZ3i957pL9JfCLCEA4GhGAcDQiAOFoRADC0YgAhKMRAQhHIwIQTg/H+bmeKZRX/dxK7anOtDRbuge2rh7Jem3Zz/qs3Xsg116s7Mj6oP6BrPfzM7+4K2pmVnusszOdjZIrnIb+rKWlQudTfjIqyZgsdHT9ktV39PyapUf+3J7Rm3rrvtLUmZb7L8ayPtnxc2yfdd+Qa4dzJXOW7vjzrMzM8u66W2tU7su1SVd/LtOK3vvVxA8PLjb0dUFza11ZV/hFBCAcjQhAOBoRgHA0IgDhaEQAwtGIAISjEQEIJ0MH6WbJTKFdf3l75t9J9f8PLmaumFnlqZ6b0twUc3ve0HNTrlZ1/z0T84bMzEYVP0+RLZTMI+rpjEna1LOQipr/2gaFnoPTKZnh87K1a3q2zvDwa7eWVN6Va79O9Lyr0y39WdTa/mdZ6K/YKg393EmzJEtW9fdXPb8n12YVPURsOtV3/qVVPwdXSfXfRWP6jaxbzc9H8YsIQDgaEYBwNCIA4WhEAMLRiACEoxEBCCeP7+tjfVyXDPw+Nr5ZMjKgLctWG96W9eFN//h/WOgRIgcz/yoiMzM70/15Onvh1vIv9fFo8Wlf1n9585Gsz6X+sXLn+uty7bRkDMj1NVm+dINr/ynr3Zb/WV4/fiLXfjXnHxWbmS219Zu97k8gsS9qb8u1taaOvSRf6QhHY8uPcLz+W32V0cK+fuz7V3R0IL3qj/75+qreP6t1HVuQz/sHrwSAS0IjAhCORgQgHI0IQDgaEYBwNCIA4WhEAMLJHFG+qDMv07p/HUxlXgeF0q7OJBQzncVQoxZOa7q/JknJc3d01sJO/KxG5bxkRkRbj2lIDvQohUl1162dde7ItbvLJe/rJZtN9CiYZsO/WmlpXV91tTvUV/YMavrqm53eiVvLdvR31Ozr97VkV2S9fuLv305Vj4kZ3NbXKE07OuvT7rziP7b5Y1nMzLJEX6Wl8IsIQDgaEYBwNCIA4WhEAMLRiACEoxEBCEcjAhBO5oi6+3q2zvnIzyxU93UOqN7UWYujYUnmZeK/9PzHLbk0beiMieV6DlNRE69tS2ev8iWdM5r8z31Zt6n/vmuLQ7m0biWf6a2SDNQl6490lufxkX+lVLai511d9Jdlvd3Ve3uu4efgjjp6f53XDmV9vKOfe2vqf4+LR+IaLTM7eCbL9vq2ziFNK35262FPr+0WO/rJxfbjFxGAcDQiAOFoRADC0YgAhKMRAQhHIwIQjkYEIJzMEZ0/1Fmf2difbdJakg9tjWOdM2oN9NyUdMWfKdT3rx0zM7OkouempEslc1dEBiq74WdfzMyKic6B2EhnNZKnfgYqXdAzerKKzjiZlby2S1aMt2Q9y/2ZQJNTfX/Xwq5+r+u7JRfrXfif5W5D55/GVT+LY2Y2PdE5opOav7cnL/Tacabrw1d1Vqwprk27SPVndn+q9/6PRY1fRADC0YgAhKMRAQhHIwIQjkYEIByNCEA4ecY+fWtRLi6e+MfYyS3/2hszs6o+vbfVeXGOaGbpsn+836zoI8pW5VzWK6OOrA/P/DEQByXjKWyij+cblXdkPfneglurtf3rnczMOvZrWTfT1xFdtvTf3pP1SubHLM7ufCzXLi3pvftFT49EyVJ/xM3KJyX5kIq+0md3zb+yx8zMTvy9/au7OrZwVnJNV1ayPUdikszWmf7dcudz/b7tA7/ELyIA4WhEAMLRiACEoxEBCEcjAhCORgQgHI0IQDg9q+OoZPVgzS1NamO59GSisz7TOZ2HSBv+lS2d2Z5cm5geZ9AZ6REkS0d+vuXYNuXa2bgk53Hak/Xkjp9h6m3pL2yuZPzEy1Yd6EyMpf61Tp3c33tmZnmm91e1pbd+VvEDN5OeDuMMr+rnzvp6D0zFbUXDtZJruib6sWdJyXpR3sz1Y9eX9N+Nwi8iAOFoRADC0YgAhKMRAQhHIwIQjkYEIByNCEA4GaZIbj0qWf2RW5qdX5FLpzURljAz++dtWW5v+zODsuLv5dpZsa6fu6dnBg2f77i17++uyLX9RAx8MbNvFq7J+lV/HJEtrOuMyHJJbOxlKwp9/VE658+kOr22L9eu1vR3eHdBX5VVVPzP8pPp9+TacaY/52Km8zazhl8vVo/l2tFMv6+x6b+7xrL/vvW0K7NnO3oG2Q9v33Vr/CICEI5GBCAcjQhAOBoRgHA0IgDhaEQAwtGIAISTgYe0ou/3yid+bqBX6JlA/Yme6TLJS7IWW/5MocpkWa4dZnoe0UIyL+vdsZ9vycc6v1JJdM5o7rnOgTSO/DTH8xV9X9vi4XVZNx3dunRZqy/redX/HruzX8q106GeCfTlut4jlvp7v18yy6hIdVYsnerXNqr5+2ta9Wc0mZmJ+JOZmU0m+r7AvOHf99Yr9N9Nuqj3vlz7B68EgEtCIwIQjkYEIByNCEA4GhGAcDQiAOHkOeTd3+urcSYTv/7jrj5H7Jdcm/NP/ZLrYP7Dv9qk85M/l2sng6asv3Wqn7tz9x23NlrXR5z912TZXt/XR7927o+B+JfDe3LpzkiPgHhLP/Ole60kZnE89b/j9/77b+Tajumrb7481PGQ4RV//URPL7GkrY/YbxZfyfps7P/trBY69tKo6r1dz8UcGTNbSf0rp35T0+9r8RVdV/hFBCAcjQhAOBoRgHA0IgDhaEQAwtGIAISjEQEIJ3NE5xd6DMjKkZ9L+SjROaJu7ucVzMwGOgZis8L/B2eZzstMG3pcwUeHY1lv9Py8RH+xLtfunejP5caczmJs7vvXNLUe6fzT502dnflQVi9fp+OPnDAzK079z6qS6v9DT6f6ueuHev2zY39/ZQf6c0y2dZasuKKvs+pe+Hug9aX+jtORHiNz8lS/71vr/t/GzQ39R/nbtGQGiYiN8YsIQDgaEYBwNCIA4WhEAMLRiACEoxEBCEcjAhAuKYqSs38A+JbxiwhAOBoRgHA0IgDhaEQAwtGIAISjEQEI93/vG6WqLwR19gAAAABJRU5ErkJggg==\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "data = test_smallbatch_img(img_path, 16, 4)\n",
        "real = None\n",
        "for image in data:\n",
        "  real = image\n",
        "  real = real.cuda()\n",
        "  plotter(image)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 303
        },
        "id": "pPM19eVxwatf",
        "outputId": "0b8f88b2-0a8e-4eb5-d45d-c2b83f39c238"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 360x360 with 4 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAASIAAAEeCAYAAAAuBhhgAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAL/klEQVR4nO3dfZBWVR0H8Ocuu7yYAoumGIaapLwIBWQBQy+m5GTNmMmYk5ZmvjRkM9ZMOZZjWFZjpTYFhZKko5OFVljKlGj2oqAllSKiBWVkimuoKSMiuk9/NU0z3t9d9tlnf88un8+/X865Z5fLlztzz5xb1Ov1GkCmtuwFACgiIJ0iAtIpIiCdIgLSKSIgXXsUFkXR63f7udsCqq5d9Msq+ltRNPZz1ev1fv3FDNz7a/fUzPvLExGQThEB6RQRkE4RAekUEZBOEQHpFBGQLtxHNFDt/NumMP/6D5eH+YRDXh/mxx8xszRrP+h14Vhaw4aH/hHmkya+ttdz37yhK8z3Gj0szN++/6heX/u623aE+SlHx9fO4okISKeIgHSKCEiniIB0ighIp4iAdEV0nMJAPaZh9d13xn9gR/yK87G/bQzzorOzNDvhuBPjazeRY0D6zi/++UJpdsy44eHYy1f8KcyHtce/5gXvfUNpdssjL4Zj33PQ0DBvhGNAgEFNEQHpFBGQThEB6RQRkE4RAekUEZBuUB4DMmfW3DD//W9+Hubvn398mC+/7Y5dXhOt5c7VD4b5MXMm93rujXevC/PLP/m2Xs9dtU/ozjUbwnzu7Em9vnYzeSIC0ikiIJ0iAtIpIiCdIgLSKSIgnSIC0g3K84jOu/r7Yd629emKGeK1d+89pjS75LQPVszdPM4j6jtr/1S+z+jg8WMbmruodVf8gfKfrbPz1Q1duxHOIwIGNUUEpFNEQDpFBKRTREA6RQSkU0RAukF5HlGV7r3Lv0tWq9Vq//7XE2Heue25vlzO/1l27RVhfvqHzm7atfmfzvGHlmbHnbckHHvKrAPD/Oa/xt8mGzai/Ltp758R7yOaduC+YX74lIPCPIsnIiCdIgLSKSIgnSIC0ikiIJ0iAtKlvb6//NtfDfNPLvhMxQzlRyV0V/TrtBfvDPM/bHo5zMccOTXMG+H1fP/47eo/hvnkCfuUZuPGHxyOPX3+jDD/wOZHw3z+d8s/CfTA4+Wv9mu1Wm3EjmfC3Ot7gBKKCEiniIB0ighIp4iAdIoISKeIgHRp+4iq9wlVKf+0ychHN4YjV63fGs88tCPMu266qzSrn3BmPPeQeG76x1vnTA/zDQ+V30NdT/4rHPvcy8PCfMlN8R6mV40sP+rj3ZPKP2VVq+V+xqsRnoiAdIoISKeIgHSKCEiniIB0ighIp4iAdIPyc0IP3n9/mHcMj890qdqKMW7cuNLMPqHBYedL5WdSXTB/cji2qHeH+TuPmR3mJ4x6VWl22KGvC8cOVJ6IgHSKCEiniIB0ighIp4iAdIoISKeIgHQtu4/oqKOOCvPbb7+9NLt2yaXh2Ne8eVKY3/KDW8P8N79cGeYMfNMOP6zXYzf8bl2Yj463GQ3avUIRT0RAOkUEpFNEQDpFBKRTREA6RQSka9nX99Hr+SpD2oeG+SFj49f3R7xpTphPH19+DAi0d8RHwYzuHNVPKxk4PBEB6RQRkE4RAekUEZBOEQHpFBGQThEB6Vp2H9HixYvDPNpntPr++BiPLRu39WpN/1Wvl39qBto74n9WxZAh/bSSgcMTEZBOEQHpFBGQThEB6RQRkE4RAekUEZCuqNfr5WFRlIcVonn7wluOnF6anfHhs8OxZ37kY329nB6bf8r7wvzG61b0eu6iKHo9tlar1er1emMT7KJWvr8asenhjWG+9+gxYT56vzjP0sz7yxMRkE4RAekUEZBOEQHpFBGQThEB6RQRkK5lzyNaevWSMJ84YWpptu2xv/T1cnrsmm9fHOZf+NyX+mklZDnksAlhvunBeJ9Rq+4jaiZPREA6RQSkU0RAOkUEpFNEQDpFBKRr2df3Z54WH9UR5SedenxfL6fHunbEx1NMnjSln1ZClkWLFoX5vWvvDfNvLl1WmnXteCkcO6I7Pqpj3F4dYZ7FExGQThEB6RQRkE4RAekUEZBOEQHpFBGQbsB+Tijywo7tYT582IiG5r/4yvIjSi44K+9TRT4nNDCcf96nwvwrl1zWTyvZNT4nBAxqighIp4iAdIoISKeIgHSKCEiniIB04T4igP7giQhIp4iAdIoISKeIgHSKCEiniIB0ighIp4iAdIoISKeIgHSKCEiniIB0ighIp4iAdIoISKeIgHSKCEiniIB0ighIp4iAdIoISKeIgHSKCEjXHoVFUTTto2et/D21oijCPHPtVWtrRL1eb97kr6CZ91eVVr7/mqmZ90+V6P7yRASkU0RAOkUEpFNEQDpFBKRTREA6RQSkC/cR7a6+872VYT5z5szSbO3atX29HFrMsquuCvMTZ44K8z3fOL8vlzMoeCIC0ikiIJ0iAtIpIiCdIgLSKSIg3W75+v7Dnzg3zI/ueCnMP/DRBaVZKx8hQt8YuqMrzM+98PowP/mcjjA/8l3H7fKaBjpPREA6RQSkU0RAOkUEpFNEQDpFBKRTREC6QbmPqGiL9/J0P7kozLd1vTnM99y3/KiP/fa/LhzLwHfKgvPDvHPrY2G+O+4TquKJCEiniIB0ighIp4iAdIoISKeIgHSKCEhXROfjFEXRtMNzXu7uDvO2inN9GrHs1Hgfx5j3nR7m79x2b2k28kNf7NWaeqrqvKNG1Ov15k3+Cpp5f1Vp5FyoT51zVpifNntKmK9Y//cwv/DLl+3ymnqqmfdPlej+8kQEpFNEQDpFBKRTREA6RQSkU0RAOkUEpEs7j6hqn9CP1jwQ5q9++YnSbMOWp8Oxsy74RphP278zzB9f/KvSbGQ4ksHg/IWXhHlbW/z/+/a1n+3L5QwKnoiAdIoISKeIgHSKCEiniIB0ighIl3YMSNUxDN0VeSPHhDz4jy1hPmX8/vG19ynf9bD98efDsUPbO8K8imNA+kYjx4CsWPjxMB8+fESYL7vn4TBf/pOf7fKaesoxIAAlFBGQThEB6RQRkE4RAekUEZBOEQHp0o4BqVK1T+jFZx4tzYaOPiAc+9Pbfh3mb5s4PswXLF1Rfu0G9wnRGnbu3FmaPfPA8HDsSRftF+Y3/HFdr9Y0mHkiAtIpIiCdIgLSKSIgnSIC0ikiIJ0iAtK17D6iKtFeoS1P/Tsce9Chk8J8n7fMC/PNV15Umt3Rdm449sg57whzWkNHR/l+sKFz4382zxZHhPnmJ28O8y1bys/LGjt2bDh2oPJEBKRTREA6RQSkU0RAOkUEpFNEQDpFBKQbsPuIFl13Q2m24OT54dh/PrIpzH98zVVhfkN3d5hHfrBmfZifOGtyr+emf+zxzLYwX728/N6s1Wq14+dOD/POHY+Uh1u7wrHnf/7SMG9VnoiAdIoISKeIgHSKCEiniIB0ighIV9Tr9fKwKMrDBkXX7Yki+NzQ5iefDse2B0c81Gq12n4j9+jVmmq1Wm1IW9zt16/ZEOYnzZoY5tHP3ah6vd68yV9BM++vKo3cfy+s/1qYD5/y6fjaf45fsa+49anSrH1IfH89vG2vMP/0Z84L82aK7i9PREA6RQSkU0RAOkUEpFNEQDpFBKRTREC6AbuP6KXgKI62ir0231x+S5jfternYf7Dpd8qzdZveigcO3VC/CmjKvYR9Y1G7r9169aF+dSpU+MJNn0jzg+JP0nViGbeP1XsIwJamiIC0ikiIJ0iAtIpIiCdIgLSKSIg3YDdR9TQtbc+EebPd90U5t37zivN9hw2Jhxb7DkqzKvYR9Q3Mu+/++9bG+aHTy3/3FBbxXlXVewjAiihiIB0ighIp4iAdIoISKeIgHSKCEi3W+4j+voV14T527ffE+YHHHhgaXbFfdvDsQsXLgzzKvYR9Y3M+69K/ak1pVkxZnZDc9tHBFBCEQHpFBGQThEB6RQRkE4RAel2y9f382dMDvMTZkwM89kzp5VmBy+4KBzb6M/t9X3faOTv4dhjjw3zlStX9nruWq1WW7Xq1tJs3rx3NTS31/cAJRQRkE4RAekUEZBOEQHpFBGQThEB6XbLfURVeykaWVsz5+7J/I2wj6hnzjnjjDB/7tlnw7y7uzvMr73xxl1eU0/ZRwRQQhEB6RQRkE4RAekUEZBOEQHpFBGQLtxHBNAfPBEB6RQRkE4RAekUEZBOEQHpFBGQ7j9Iyph+IDLkYAAAAABJRU5ErkJggg==\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Fake Pred"
      ],
      "metadata": {
        "id": "j8nKDTCYzuyl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pred = dis(fake[0:4].detach(), 0.5, 2)"
      ],
      "metadata": {
        "id": "Gn69ssn2he34"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(pred)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OUcV7wiejBS2",
        "outputId": "f5c364a4-4703-43e5-e84c-c6989ff5e2ba"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([[1.0674],\n",
            "        [1.0293],\n",
            "        [1.0155],\n",
            "        [1.0903]], device='cuda:0', grad_fn=<ViewBackward0>)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "real.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NoAXbEhS0TAo",
        "outputId": "0f5398ca-e4ae-4a70-c335-016744068f16"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([4, 3, 32, 32])"
            ]
          },
          "metadata": {},
          "execution_count": 97
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "predict = dis(real, 0.5, 2)"
      ],
      "metadata": {
        "id": "iou4z6dw0PAa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(predict)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-P9CSTyylNzh",
        "outputId": "0e9c61b0-0c46-4aa2-be69-5e424848b1bd"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([[0.6053],\n",
            "        [0.4921],\n",
            "        [0.8547],\n",
            "        [0.3865]], device='cuda:0', grad_fn=<ViewBackward0>)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "pred.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DWsz8XCTjCh7",
        "outputId": "44454567-91ae-4b75-c823-0a9ea20a973f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([10, 1])"
            ]
          },
          "metadata": {},
          "execution_count": 50
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "Z8FskUHBjlUf"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "90BXN6vYSDQJ",
        "Bf8GKvVySJkt",
        "x6Mav3Rhnt2g",
        "Ua1674O2Zgl_",
        "ob2HSmfdQc41",
        "eG3NYtgDSXGB",
        "IbxAojXcVW8c",
        "C5-Xv7UJS8o6",
        "TXLLaY7dUJ2E",
        "14RgnVqrIFA8",
        "7hrtvJMR53Na"
      ],
      "machine_shape": "hm",
      "name": "PgGans.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}