CIFAR_10_cSAWGAN2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "CIFAR 10 cSAWGAN2.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "TJAxoFJKhlhW",
        "kmZ9gfjUjs_A",
        "Mm964yyI4AiK",
        "Y8xtngMv2yDK",
        "QfUqmsxmTPsu",
        "ltEqQsIpK1Us",
        "mHL5-Ns5zg1O"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "nINXlbCewTtc",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "tensorboard --logdir='./logs' --port=6006"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dPrgnGl9ZR_E",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive/',force_remount=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N8xfw1bjrLRv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch \n",
        "import torchvision\n",
        "from torch import nn, optim\n",
        "from torch.autograd.variable import Variable\n",
        "from torchvision import transforms, datasets\n",
        "import os\n",
        "from PIL import Image\n",
        "import glob\n",
        "import PIL\n",
        "from PIL import ImageFont\n",
        "from PIL import Image\n",
        "from PIL import ImageDraw\n",
        "import matplotlib.pyplot as plt\n",
        "import random\n",
        "from bokeh.io import curdoc, show, output_notebook\n",
        "from bokeh.layouts import column\n",
        "from bokeh.models import ColumnDataSource\n",
        "from bokeh.plotting import figure\n",
        "from functools import partial\n",
        "from threading import Thread\n",
        "from tornado import gen\n",
        "import time\n",
        "import pickle\n",
        "from tqdm import tqdm\n",
        "import torchvision.utils as vutils"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G81dSJ3Uquzz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "\n",
        "\n",
        "#root folder\n",
        "root_dir=\"/content/drive/My Drive/Fall 2019/Deep Learning CSE 676/Projects/1/DC Gans/\"\n",
        "\n",
        "#data directories\n",
        "#output images\n",
        "output_dir=root_dir+\"output/epoch/\"\n",
        "#input images\n",
        "input_dir=root_dir+\"data/augmented/\"\n",
        "\n",
        "#models\n",
        "model_dir=root_dir+\"model/\"\n",
        "#other resources\n",
        "res_dir=root_dir+\"res/\"\n",
        "#report and logging\n",
        "report_dir=root_dir+\"report/\"\n",
        "\n",
        "\n",
        "#parameters\n",
        "batch_size=100\n",
        "train_split=0.8\n",
        "train_epoch=50\n",
        "\n",
        "#input\n",
        "data_dir=\"data\"\n",
        "inp_width=32\n",
        "inp_height=32\n",
        "inp_channels=3\n",
        "nc=3\n",
        "nz=200\n",
        "\n",
        "#generator\n",
        "ngf=64\n",
        "ndf=64\n",
        "\n",
        "#discriminator\n",
        "ndf=2"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yX9_Iv0XfRKp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model='sagan'\n",
        "adv_loss='wgan-gp'\n",
        "imsize=64\n",
        "g_num=5\n",
        "z_dim=100\n",
        "label_dim=10\n",
        "g_conv_dim=64\n",
        "d_conv_dim=64\n",
        "lambda_gp=10\n",
        "version='sawgan_trail1'\n",
        "\n",
        "#training parameters\n",
        "total_step=100000\n",
        "d_iters=5\n",
        "batch_size=64\n",
        "num_workers=2\n",
        "g_lr=0.0001\n",
        "d_lr=0.0004\n",
        "lr_decay=0.95\n",
        "beta1=0\n",
        "beta2=0.9\n",
        "\n",
        "#pretrained\n",
        "pretrained_model=None\n",
        "\n",
        "#misc\n",
        "train=True\n",
        "parallel=False\n",
        "dataset='cifar'\n",
        "use_tensorboard=True\n",
        "\n",
        "#paths\n",
        "root_dir=\"/content/drive/My Drive/Fall 2019/Deep Learning CSE 676/Projects/1/SA Gans/\"\n",
        "\n",
        "image_path=root_dir+\"data\"\n",
        "log_path=root_dir+\"logs\"\n",
        "model_save_path=root_dir+\"models\"\n",
        "sample_path=root_dir+\"samples\"\n",
        "attn_path=root_dir+\"attn\"\n",
        "log_step=10\n",
        "sample_step=100\n",
        "model_save_step=1\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D_7FLsJSq_NS",
        "colab_type": "text"
      },
      "source": [
        "Importing logger to keep a track of training progress "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TJAxoFJKhlhW",
        "colab_type": "text"
      },
      "source": [
        "##Dataloader"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u7BZTrX5hn2u",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torchvision.datasets as dsets\n",
        "from torchvision import transforms\n",
        "\n",
        "\n",
        "class Data_Loader():\n",
        "    def __init__(self, train, dataset, image_path, image_size, batch_size, shuf=True):\n",
        "        self.dataset = dataset\n",
        "        self.path = image_path\n",
        "        self.imsize = image_size\n",
        "        self.batch = batch_size\n",
        "        self.shuf = shuf\n",
        "        self.train = train\n",
        "\n",
        "    def transform(self, resize, totensor, normalize, centercrop):\n",
        "        options = []\n",
        "        if centercrop:\n",
        "            options.append(transforms.CenterCrop(160))\n",
        "        if resize:\n",
        "            options.append(transforms.Resize((self.imsize,self.imsize)))\n",
        "        if totensor:\n",
        "            options.append(transforms.ToTensor())\n",
        "        if normalize:\n",
        "            options.append(transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)))\n",
        "        transform = transforms.Compose(options)\n",
        "        return transform\n",
        "\n",
        "    def load_lsun(self, classes='church_outdoor_train'):\n",
        "        transforms = self.transform(True, True, True, False)\n",
        "        dataset = dsets.LSUN(self.path, classes=[classes], transform=transforms)\n",
        "        return dataset\n",
        "\n",
        "    def load_celeb(self):\n",
        "        transforms = self.transform(True, True, True, True)\n",
        "        dataset = dsets.ImageFolder(self.path+'/CelebA', transform=transforms)\n",
        "        return dataset\n",
        "\n",
        "\n",
        "    def loader(self):\n",
        "        if self.dataset == 'lsun':\n",
        "            dataset = self.load_lsun()\n",
        "        elif self.dataset == 'celeb':\n",
        "            dataset = self.load_celeb()\n",
        "\n",
        "        loader = torch.utils.data.DataLoader(dataset=self.dataset,\n",
        "                                              batch_size=self.batch,\n",
        "                                              shuffle=self.shuf,\n",
        "                                              num_workers=2,\n",
        "                                              drop_last=True)\n",
        "        return loader\n",
        "      \n",
        "    def loader(self):\n",
        "      transform = transforms.Compose(\n",
        "          [transforms.Resize((self.imsize,self.imsize)),\n",
        "           transforms.ToTensor(),\n",
        "           transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n",
        "\n",
        "      dataset_full = torchvision.datasets.CIFAR10(root='./data', train=True,\n",
        "                                              download=True, transform=transform)\n",
        "      train_split=0.8\n",
        "      train_size=int(train_split*len(dataset_full))\n",
        "      val_size=len(dataset_full)-train_size\n",
        "      trainset, valset=torch.utils.data.random_split(dataset_full,[train_size,val_size])\n",
        "      trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,\n",
        "                                                shuffle=True, num_workers=2)\n",
        "      \n",
        "      return trainloader\n",
        "    def len(self):\n",
        "      return len(self.loader())"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-f1ueGd7hVae",
        "colab_type": "text"
      },
      "source": [
        "##main"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aWP5I_oUhVHo",
        "colab_type": "code",
        "outputId": "d8eed098-61cf-4109-92f8-596204173e57",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "\n",
        "from torch.backends import cudnn\n",
        "\n",
        "cudnn.benchmark=True\n",
        "# Data loader\n",
        "data_loader = Data_Loader(train, dataset, image_path, imsize,\n",
        "                        batch_size, shuf=train)\n",
        "\n",
        "\n",
        "# Create directories if not exist\n",
        "make_folder(model_save_path, version)\n",
        "make_folder(sample_path,version)\n",
        "make_folder(log_path,version)\n",
        "make_folder(attn_path,version)\n",
        "\n",
        "\n",
        "if train:\n",
        "    if model=='sagan':\n",
        "        trainer = Trainer(data_loader.loader())\n",
        "    elif model == 'qgan':\n",
        "        trainer = qgan_trainer(data_loader.loader())\n",
        "    trainer.train()\n",
        "else:\n",
        "    tester = Tester(data_loader.loader())\n",
        "    tester.test()"
      ],
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Files already downloaded and verified\n",
            "Generator(\n",
            "  (l4): Sequential(\n",
            "    (0): SpectralNorm(\n",
            "      (module): ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
            "    )\n",
            "    (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (2): ReLU()\n",
            "  )\n",
            "  (l1): Sequential(\n",
            "    (0): SpectralNorm(\n",
            "      (module): ConvTranspose2d(164, 512, kernel_size=(4, 4), stride=(1, 1))\n",
            "    )\n",
            "    (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (2): ReLU()\n",
            "  )\n",
            "  (l2): Sequential(\n",
            "    (0): SpectralNorm(\n",
            "      (module): ConvTranspose2d(512, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
            "    )\n",
            "    (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (2): ReLU()\n",
            "  )\n",
            "  (l3): Sequential(\n",
            "    (0): SpectralNorm(\n",
            "      (module): ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
            "    )\n",
            "    (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (2): ReLU()\n",
            "  )\n",
            "  (last): Sequential(\n",
            "    (0): ConvTranspose2d(64, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
            "    (1): Tanh()\n",
            "  )\n",
            "  (attn1): Self_Attn(\n",
            "    (query_conv): Conv2d(128, 16, kernel_size=(1, 1), stride=(1, 1))\n",
            "    (key_conv): Conv2d(128, 16, kernel_size=(1, 1), stride=(1, 1))\n",
            "    (value_conv): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1))\n",
            "    (softmax): Softmax(dim=-1)\n",
            "  )\n",
            "  (attn2): Self_Attn(\n",
            "    (query_conv): Conv2d(64, 8, kernel_size=(1, 1), stride=(1, 1))\n",
            "    (key_conv): Conv2d(64, 8, kernel_size=(1, 1), stride=(1, 1))\n",
            "    (value_conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n",
            "    (softmax): Softmax(dim=-1)\n",
            "  )\n",
            ")\n",
            "Discriminator(\n",
            "  (l4): Sequential(\n",
            "    (0): SpectralNorm(\n",
            "      (module): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
            "    )\n",
            "    (1): LeakyReLU(negative_slope=0.1)\n",
            "  )\n",
            "  (l1): Sequential(\n",
            "    (0): SpectralNorm(\n",
            "      (module): Conv2d(3, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
            "    )\n",
            "    (1): LeakyReLU(negative_slope=0.1)\n",
            "  )\n",
            "  (l2): Sequential(\n",
            "    (0): SpectralNorm(\n",
            "      (module): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
            "    )\n",
            "    (1): LeakyReLU(negative_slope=0.1)\n",
            "  )\n",
            "  (l3): Sequential(\n",
            "    (0): SpectralNorm(\n",
            "      (module): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
            "    )\n",
            "    (1): LeakyReLU(negative_slope=0.1)\n",
            "  )\n",
            "  (last): Sequential(\n",
            "    (0): Conv2d(512, 1024, kernel_size=(4, 4), stride=(1, 1))\n",
            "  )\n",
            "  (hidden): Sequential(\n",
            "    (0): Linear(in_features=1034, out_features=1, bias=True)\n",
            "  )\n",
            "  (attn1): Self_Attn(\n",
            "    (query_conv): Conv2d(256, 32, kernel_size=(1, 1), stride=(1, 1))\n",
            "    (key_conv): Conv2d(256, 32, kernel_size=(1, 1), stride=(1, 1))\n",
            "    (value_conv): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))\n",
            "    (softmax): Softmax(dim=-1)\n",
            "  )\n",
            "  (attn2): Self_Attn(\n",
            "    (query_conv): Conv2d(512, 64, kernel_size=(1, 1), stride=(1, 1))\n",
            "    (key_conv): Conv2d(512, 64, kernel_size=(1, 1), stride=(1, 1))\n",
            "    (value_conv): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
            "    (softmax): Softmax(dim=-1)\n",
            "  )\n",
            ")\n",
            "Size of labels while encoding:  torch.Size([64])\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "error",
          "ename": "KeyError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-17-5da1ba338dcc>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     20\u001b[0m     \u001b[0;32melif\u001b[0m \u001b[0mmodel\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'qgan'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m         \u001b[0mtrainer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mqgan_trainer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata_loader\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloader\u001b[0m\u001b[0;34m(\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[0;32m---> 22\u001b[0;31m     \u001b[0mtrainer\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     23\u001b[0m \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     24\u001b[0m     \u001b[0mtester\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mTester\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata_loader\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloader\u001b[0m\u001b[0;34m(\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[0;32m<ipython-input-14-4d2b9a04ac96>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     98\u001b[0m             \u001b[0;31m# dr1, dr2, df1, df2, gf1, gf2 are attention scores\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     99\u001b[0m             \u001b[0mreal_images\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtensor2var\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mreal_images\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 100\u001b[0;31m             \u001b[0mlabels\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtensor2var\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mencode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlabels\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    101\u001b[0m             \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Size of labels: \"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mlabels\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m(\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    102\u001b[0m             \u001b[0md_out_real\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mdr1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mdr2\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mD\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mreal_images\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mlabels\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m#labels not added in generator, generator still not sorted.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-16-6392cc13db70>\u001b[0m in \u001b[0;36mencode\u001b[0;34m(label)\u001b[0m\n\u001b[1;32m     37\u001b[0m           \u001b[0;34m'ship'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;36m8\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     38\u001b[0m           'truck':9}\n\u001b[0;32m---> 39\u001b[0;31m   \u001b[0mres\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mlabels\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mlabel\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     40\u001b[0m   \u001b[0;32mreturn\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mres\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyError\u001b[0m: tensor([4, 6, 4, 8, 5, 4, 6, 3, 8, 0, 1, 4, 1, 9, 0, 0, 9, 7, 8, 8, 8, 1, 6, 4,\n        7, 6, 8, 9, 0, 1, 0, 2, 8, 6, 0, 4, 5, 7, 6, 7, 9, 1, 7, 1, 4, 4, 9, 0,\n        0, 8, 1, 3, 1, 2, 0, 9, 8, 2, 8, 4, 7, 1, 0, 5])"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kmZ9gfjUjs_A",
        "colab_type": "text"
      },
      "source": [
        "##Spectral\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eknZPG2jjv_4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "from torch.optim.optimizer import Optimizer, required\n",
        "\n",
        "from torch.autograd import Variable\n",
        "import torch.nn.functional as F\n",
        "from torch import nn\n",
        "from torch import Tensor\n",
        "from torch.nn import Parameter\n",
        "\n",
        "def l2normalize(v, eps=1e-12):\n",
        "    return v / (v.norm() + eps)\n",
        "\n",
        "\n",
        "class SpectralNorm(nn.Module):\n",
        "    def __init__(self, module, name='weight', power_iterations=1):\n",
        "        super(SpectralNorm, self).__init__()\n",
        "        self.module = module\n",
        "        self.name = name\n",
        "        self.power_iterations = power_iterations\n",
        "        if not self._made_params():\n",
        "            self._make_params()\n",
        "\n",
        "    def _update_u_v(self):\n",
        "        u = getattr(self.module, self.name + \"_u\")\n",
        "        v = getattr(self.module, self.name + \"_v\")\n",
        "        w = getattr(self.module, self.name + \"_bar\")\n",
        "\n",
        "        height = w.data.shape[0]\n",
        "        for _ in range(self.power_iterations):\n",
        "            v.data = l2normalize(torch.mv(torch.t(w.view(height,-1).data), u.data))\n",
        "            u.data = l2normalize(torch.mv(w.view(height,-1).data, v.data))\n",
        "\n",
        "        # sigma = torch.dot(u.data, torch.mv(w.view(height,-1).data, v.data))\n",
        "        sigma = u.dot(w.view(height, -1).mv(v))\n",
        "        setattr(self.module, self.name, w / sigma.expand_as(w))\n",
        "\n",
        "    def _made_params(self):\n",
        "        try:\n",
        "            u = getattr(self.module, self.name + \"_u\")\n",
        "            v = getattr(self.module, self.name + \"_v\")\n",
        "            w = getattr(self.module, self.name + \"_bar\")\n",
        "            return True\n",
        "        except AttributeError:\n",
        "            return False\n",
        "\n",
        "\n",
        "    def _make_params(self):\n",
        "        w = getattr(self.module, self.name)\n",
        "\n",
        "        height = w.data.shape[0]\n",
        "        width = w.view(height, -1).data.shape[1]\n",
        "\n",
        "        u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)\n",
        "        v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False)\n",
        "        u.data = l2normalize(u.data)\n",
        "        v.data = l2normalize(v.data)\n",
        "        w_bar = Parameter(w.data)\n",
        "\n",
        "        del self.module._parameters[self.name]\n",
        "\n",
        "        self.module.register_parameter(self.name + \"_u\", u)\n",
        "        self.module.register_parameter(self.name + \"_v\", v)\n",
        "        self.module.register_parameter(self.name + \"_bar\", w_bar)\n",
        "\n",
        "\n",
        "    def forward(self, *args):\n",
        "        self._update_u_v()\n",
        "        return self.module.forward(*args)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KLzOWR3EjUs3",
        "colab_type": "text"
      },
      "source": [
        "##models"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QiXGEDQ1jWwG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.autograd import Variable\n",
        "import numpy as np\n",
        "\n",
        "class Self_Attn(nn.Module):\n",
        "    \"\"\" Self attention Layer\"\"\"\n",
        "    def __init__(self,in_dim,activation):\n",
        "        super(Self_Attn,self).__init__()\n",
        "        self.chanel_in = in_dim\n",
        "        self.activation = activation\n",
        "        \n",
        "        self.query_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)\n",
        "        self.key_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)\n",
        "        self.value_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim , kernel_size= 1)\n",
        "        self.gamma = nn.Parameter(torch.zeros(1))\n",
        "\n",
        "        self.softmax  = nn.Softmax(dim=-1) #\n",
        "    def forward(self,x):\n",
        "        \"\"\"\n",
        "            inputs :\n",
        "                x : input feature maps( B X C X W X H)\n",
        "            returns :\n",
        "                out : self attention value + input feature \n",
        "                attention: B X N X N (N is Width*Height)\n",
        "        \"\"\"\n",
        "        m_batchsize,C,width ,height = x.size()\n",
        "        proj_query  = self.query_conv(x).view(m_batchsize,-1,width*height).permute(0,2,1) # B X CX(N)\n",
        "        proj_key =  self.key_conv(x).view(m_batchsize,-1,width*height) # B X C x (*W*H)\n",
        "        energy =  torch.bmm(proj_query,proj_key) # transpose check\n",
        "        attention = self.softmax(energy) # BX (N) X (N) \n",
        "        proj_value = self.value_conv(x).view(m_batchsize,-1,width*height) # B X C X N\n",
        "\n",
        "        out = torch.bmm(proj_value,attention.permute(0,2,1) )\n",
        "        out = out.view(m_batchsize,C,width,height)\n",
        "        \n",
        "        out = self.gamma*out + x\n",
        "        return out,attention\n",
        "\n",
        "class Generator(nn.Module):\n",
        "    \"\"\"Generator.\"\"\"\n",
        "\n",
        "    def __init__(self, batch_size, image_size=64, z_dim=100, label_dim=10,conv_dim=64):\n",
        "        super(Generator, self).__init__()\n",
        "        self.imsize = image_size\n",
        "        layer1 = []\n",
        "        layer2 = []\n",
        "        layer3 = []\n",
        "        last = []\n",
        "\n",
        "        repeat_num = int(np.log2(self.imsize)) - 3\n",
        "        mult = 2 ** repeat_num # 8\n",
        "        layer1.append(SpectralNorm(nn.ConvTranspose2d(z_dim+label_dim, (conv_dim * mult), 4)))\n",
        "        layer1.append(nn.BatchNorm2d(conv_dim * mult))\n",
        "        layer1.append(nn.ReLU())\n",
        "        \n",
        "\n",
        "        curr_dim = conv_dim * mult\n",
        "\n",
        "        layer2.append(SpectralNorm(nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 4, 2, 1)))\n",
        "        layer2.append(nn.BatchNorm2d(int(curr_dim / 2)))\n",
        "        layer2.append(nn.ReLU())\n",
        "\n",
        "        curr_dim = int(curr_dim / 2)\n",
        "\n",
        "        layer3.append(SpectralNorm(nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 4, 2, 1)))\n",
        "        layer3.append(nn.BatchNorm2d(int(curr_dim / 2)))\n",
        "        layer3.append(nn.ReLU())\n",
        "\n",
        "        if self.imsize == 64:\n",
        "            layer4 = []\n",
        "            curr_dim = int(curr_dim / 2)\n",
        "            layer4.append(SpectralNorm(nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 4, 2, 1)))\n",
        "            layer4.append(nn.BatchNorm2d(int(curr_dim / 2)))\n",
        "            layer4.append(nn.ReLU())\n",
        "            self.l4 = nn.Sequential(*layer4)\n",
        "            curr_dim = int(curr_dim / 2)\n",
        "\n",
        "        self.l1 = nn.Sequential(*layer1)\n",
        "        self.l2 = nn.Sequential(*layer2)\n",
        "        self.l3 = nn.Sequential(*layer3)\n",
        "\n",
        "        last.append(nn.ConvTranspose2d(curr_dim, 3, 4, 2, 1))\n",
        "        last.append(nn.Tanh())\n",
        "        self.last = nn.Sequential(*last)\n",
        "\n",
        "        self.attn1 = Self_Attn( 128, 'relu')\n",
        "        self.attn2 = Self_Attn( 64,  'relu')\n",
        "\n",
        "    def forward(self, z,label=None):\n",
        "        z=torch.cat([z,label],1)\n",
        "        z = z.view(z.size(0), z.size(1), 1, 1)\n",
        "        out=self.l1(z)\n",
        "        out=self.l2(out)\n",
        "        out=self.l3(out)\n",
        "        out,p1 = self.attn1(out)\n",
        "        out=self.l4(out)\n",
        "        out,p2 = self.attn2(out)\n",
        "        out=self.last(out)\n",
        "\n",
        "        return out, p1, p2\n",
        "\n",
        "\n",
        "class Discriminator(nn.Module):\n",
        "    \"\"\"Discriminator, Auxiliary Classifier.\"\"\"\n",
        "\n",
        "    def __init__(self, batch_size=64, image_size=64,label_dim=10, conv_dim=64):\n",
        "        super(Discriminator, self).__init__()\n",
        "        self.imsize = image_size\n",
        "        layer1 = []\n",
        "        layer2 = []\n",
        "        layer3 = []\n",
        "        last = []\n",
        "        hidden=[]\n",
        "\n",
        "        layer1.append(SpectralNorm(nn.Conv2d(3, (conv_dim), 4, 2, 1)))\n",
        "        layer1.append(nn.LeakyReLU(0.1))\n",
        "\n",
        "\n",
        "        curr_dim = conv_dim\n",
        "\n",
        "        layer2.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))\n",
        "        layer2.append(nn.LeakyReLU(0.1))\n",
        "        curr_dim = curr_dim * 2\n",
        "\n",
        "        layer3.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))\n",
        "        layer3.append(nn.LeakyReLU(0.1))\n",
        "        curr_dim = curr_dim * 2\n",
        "\n",
        "        if self.imsize == 64:\n",
        "            layer4 = []\n",
        "            layer4.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))\n",
        "            layer4.append(nn.LeakyReLU(0.1))\n",
        "            self.l4 = nn.Sequential(*layer4)\n",
        "            curr_dim = curr_dim*2\n",
        "        self.l1 = nn.Sequential(*layer1)\n",
        "        self.l2 = nn.Sequential(*layer2)\n",
        "        self.l3 = nn.Sequential(*layer3)\n",
        "\n",
        "        last.append(nn.Conv2d(curr_dim, curr_dim*2, 4))\n",
        "        curr_dim = curr_dim * 2\n",
        "        hidden.append(nn.Linear(1034,1))\n",
        "        \n",
        "        self.last = nn.Sequential(*last)\n",
        "        self.hidden=nn.Sequential(*hidden)\n",
        "        \n",
        "        self.attn1 = Self_Attn(256, 'relu')\n",
        "        self.attn2 = Self_Attn(512, 'relu')\n",
        "\n",
        "    def forward(self, x,label):\n",
        "        out = self.l1(x)\n",
        "        out = self.l2(out)\n",
        "        out = self.l3(out)\n",
        "        out,p1 = self.attn1(out)\n",
        "        out=self.l4(out)\n",
        "        out,p2 = self.attn2(out)\n",
        "        out=self.last(out)\n",
        "        print(out.size())\n",
        "        #out=out.view(-1)\n",
        "        \n",
        "        out = out.reshape(out.size(0), -1)\n",
        "        out=torch.cat([out,label],1) # Should it be axis 1 or 0\n",
        "        print(\"Flattened shape= \",out.size()) # doubt of batch size x latent space would be executable?\n",
        "        #check if \n",
        "        out=self.hidden(out)\n",
        "\n",
        "        return out.squeeze(), p1, p2"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yvbTgHvxkImH",
        "colab_type": "text"
      },
      "source": [
        "##utils"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WBtgYjx4kJ6o",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import torch\n",
        "from torch.autograd import Variable\n",
        "\n",
        "\n",
        "def make_folder(path, version):\n",
        "        if not os.path.exists(os.path.join(path, version)):\n",
        "            os.makedirs(os.path.join(path, version))\n",
        "\n",
        "\n",
        "def tensor2var(x, grad=False):\n",
        "    if torch.cuda.is_available():\n",
        "        x = x.cuda()\n",
        "    return Variable(x, requires_grad=grad)\n",
        "\n",
        "def var2tensor(x):\n",
        "    return x.data.cpu()\n",
        "\n",
        "def var2numpy(x):\n",
        "    return x.data.cpu().numpy()\n",
        "\n",
        "def denorm(x):\n",
        "    out = (x + 1) / 2\n",
        "    return out.clamp_(0, 1)\n",
        "  \n",
        "def encode(label):\n",
        "  print(\"Size of labels while encoding: \",label.size())\n",
        "  res=[0]*10\n",
        "  labels={'airplane':0,\t\t\t\t\t\t\t\t\t\t\n",
        "          'automobile':1,\t\t\t\t\t\t\t\t\t\t\n",
        "          'bird':2,\t\t\t\t\t\t\t\t\t\t\n",
        "          'cat':3,\t\t\t\t\t\t\t\t\t\t\n",
        "          'deer':4,\t\t\t\t\t\t\t\t\t\t\n",
        "          'dog':5,\t\t\t\t\t\t\t\t\t\t\n",
        "          'frog':6,\t\t\t\t\t\t\t\t\t\t\n",
        "          'horse':7,\t\t\t\t\t\t\t\t\t\t\n",
        "          'ship':8,\t\t\t\t\t\t\t\t\t\t\n",
        "          'truck':9}\n",
        "  res[labels[label]]=1\n",
        "  return torch.Tensor(res)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WAdfdhZTjMRf",
        "colab_type": "text"
      },
      "source": [
        "#train"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O4Vy5roAjK5U",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 258
        },
        "outputId": "70d9e090-2dc6-4430-d479-8776569a6af8"
      },
      "source": [
        "import os\n",
        "import time\n",
        "import torch\n",
        "import datetime\n",
        "\n",
        "import torch.nn as nn\n",
        "from torch.autograd import Variable\n",
        "from torchvision.utils import save_image\n",
        "\n",
        "class Trainer(object):\n",
        "    def __init__(self, data_loader):\n",
        "\n",
        "        # Data loader\n",
        "        self.data_loader = data_loader\n",
        "\n",
        "        # exact model and loss\n",
        "        self.model = model\n",
        "        self.adv_loss = adv_loss\n",
        "\n",
        "        # Model hyper-parameters\n",
        "        self.imsize = imsize\n",
        "        self.g_num = g_num\n",
        "        self.z_dim = z_dim\n",
        "        self.g_conv_dim = g_conv_dim\n",
        "        self.d_conv_dim = d_conv_dim\n",
        "        self.parallel = parallel\n",
        "\n",
        "        self.lambda_gp = lambda_gp\n",
        "        self.total_step = total_step\n",
        "        self.d_iters = d_iters\n",
        "        self.batch_size = batch_size\n",
        "        self.num_workers = num_workers\n",
        "        self.g_lr = g_lr\n",
        "        self.d_lr = d_lr\n",
        "        self.lr_decay = lr_decay\n",
        "        self.beta1 = beta1\n",
        "        self.beta2 = beta2\n",
        "        self.pretrained_model = pretrained_model\n",
        "\n",
        "        self.dataset = dataset\n",
        "        self.use_tensorboard = use_tensorboard\n",
        "        self.image_path = image_path\n",
        "        self.log_path = log_path\n",
        "        self.model_save_path = model_save_path\n",
        "        self.sample_path = sample_path\n",
        "        self.log_step = log_step\n",
        "        self.sample_step = sample_step\n",
        "        self.model_save_step = model_save_step\n",
        "        self.version = version\n",
        "\n",
        "        # Path\n",
        "        self.log_path = os.path.join(log_path, self.version)\n",
        "        self.sample_path = os.path.join(sample_path, self.version)\n",
        "        self.model_save_path = os.path.join(model_save_path, self.version)\n",
        "\n",
        "        self.build_model()\n",
        "\n",
        "        if self.use_tensorboard:\n",
        "            self.build_tensorboard()\n",
        "\n",
        "        # Start with trained model\n",
        "        if self.pretrained_model:\n",
        "            self.load_pretrained_model()\n",
        "\n",
        "\n",
        "\n",
        "    def train(self):\n",
        "\n",
        "        # Data iterator\n",
        "        data_iter = iter(self.data_loader)\n",
        "        step_per_epoch = len(self.data_loader)\n",
        "        model_save_step = int(self.model_save_step * step_per_epoch)\n",
        "\n",
        "        # Fixed input for debugging\n",
        "        fixed_z = tensor2var(torch.randn(self.batch_size, self.z_dim))\n",
        "\n",
        "        # Start with trained model\n",
        "        if self.pretrained_model:\n",
        "            start = self.pretrained_model + 1\n",
        "        else:\n",
        "            start = 0\n",
        "\n",
        "        # Start time\n",
        "        start_time = time.time()\n",
        "        for step in range(start, self.total_step):\n",
        "\n",
        "            # ================== Train D ================== #\n",
        "            self.D.train()\n",
        "            self.G.train()\n",
        "\n",
        "            try:\n",
        "                real_images, labels = next(data_iter)\n",
        "            except:\n",
        "                data_iter = iter(self.data_loader)\n",
        "                real_images, labels = next(data_iter)\n",
        "\n",
        "            # Compute loss with real images\n",
        "            # dr1, dr2, df1, df2, gf1, gf2 are attention scores\n",
        "            real_images = tensor2var(real_images)\n",
        "            labels=tensor2var(encode(labels))\n",
        "            print(\"Size of labels: \",labels.size()) \n",
        "            d_out_real,dr1,dr2 = self.D(real_images,labels) #labels not added in generator, generator still not sorted.\n",
        "            if self.adv_loss == 'wgan-gp':\n",
        "                d_loss_real = - torch.mean(d_out_real)\n",
        "            elif self.adv_loss == 'hinge':\n",
        "                d_loss_real = torch.nn.ReLU()(1.0 - d_out_real).mean()\n",
        "\n",
        "            # apply Gumbel Softmax\n",
        "            z = tensor2var(torch.randn(real_images.size(0), self.z_dim))\n",
        "            fake_images,gf1,gf2 = self.G(z)\n",
        "            d_out_fake,df1,df2 = self.D(fake_images)\n",
        "\n",
        "            if self.adv_loss == 'wgan-gp':\n",
        "                d_loss_fake = d_out_fake.mean()\n",
        "            elif self.adv_loss == 'hinge':\n",
        "                d_loss_fake = torch.nn.ReLU()(1.0 + d_out_fake).mean()\n",
        "\n",
        "\n",
        "            # Backward + Optimize\n",
        "            d_loss = d_loss_real + d_loss_fake\n",
        "            self.reset_grad()\n",
        "            d_loss.backward()\n",
        "            self.d_optimizer.step()\n",
        "\n",
        "\n",
        "            if self.adv_loss == 'wgan-gp':\n",
        "                # Compute gradient penalty\n",
        "                alpha = torch.rand(real_images.size(0), 1, 1, 1).cuda().expand_as(real_images)\n",
        "                interpolated = Variable(alpha * real_images.data + (1 - alpha) * fake_images.data, requires_grad=True)\n",
        "                out,_,_ = self.D(interpolated)\n",
        "\n",
        "                grad = torch.autograd.grad(outputs=out,\n",
        "                                           inputs=interpolated,\n",
        "                                           grad_outputs=torch.ones(out.size()).cuda(),\n",
        "                                           retain_graph=True,\n",
        "                                           create_graph=True,\n",
        "                                           only_inputs=True)[0]\n",
        "\n",
        "                grad = grad.view(grad.size(0), -1)\n",
        "                grad_l2norm = torch.sqrt(torch.sum(grad ** 2, dim=1))\n",
        "                d_loss_gp = torch.mean((grad_l2norm - 1) ** 2)\n",
        "\n",
        "                # Backward + Optimize\n",
        "                d_loss = self.lambda_gp * d_loss_gp\n",
        "\n",
        "                self.reset_grad()\n",
        "                d_loss.backward()\n",
        "                self.d_optimizer.step()\n",
        "\n",
        "            # ================== Train G and gumbel ================== #\n",
        "            # Create random noise\n",
        "            z = tensor2var(torch.randn(real_images.size(0), self.z_dim))\n",
        "            fake_images,_,_ = self.G(z)\n",
        "\n",
        "            # Compute loss with fake images\n",
        "            g_out_fake,_,_ = self.D(fake_images)  # batch x n\n",
        "            if self.adv_loss == 'wgan-gp':\n",
        "                g_loss_fake = - g_out_fake.mean()\n",
        "            elif self.adv_loss == 'hinge':\n",
        "                g_loss_fake = - g_out_fake.mean()\n",
        "\n",
        "            self.reset_grad()\n",
        "            g_loss_fake.backward()\n",
        "            self.g_optimizer.step()\n",
        "\n",
        "\n",
        "            # Print out log info\n",
        "            if (step + 1) % self.log_step == 0:\n",
        "                elapsed = time.time() - start_time\n",
        "                elapsed = str(datetime.timedelta(seconds=elapsed))\n",
        "                print(d_loss_real)\n",
        "                print(\"Elapsed [{}], G_step [{}/{}], D_step[{}/{}], d_out_real: {:.4f}, \"\n",
        "                      \" ave_gamma_l3: {:.4f}, ave_gamma_l4: {:.4f}\".\n",
        "                      format(elapsed, step + 1, self.total_step, (step + 1),\n",
        "                             self.total_step , d_loss_real.item(),\n",
        "                             self.G.attn1.gamma.mean().item(), self.G.attn2.gamma.mean().item()))\n",
        "\n",
        "            # Sample images\n",
        "            if (step + 1) % self.sample_step == 0:\n",
        "                fake_images,_,_= self.G(fixed_z)\n",
        "                save_image(denorm(fake_images.data),\n",
        "                           os.path.join(self.sample_path, '{}_fake.png'.format(step + 1)))\n",
        "\n",
        "            if (step+1) % model_save_step==0:\n",
        "                torch.save(self.G.state_dict(),\n",
        "                           os.path.join(self.model_save_path, '{}_G.pth'.format(step + 1)))\n",
        "                torch.save(self.D.state_dict(),\n",
        "                           os.path.join(self.model_save_path, '{}_D.pth'.format(step + 1)))\n",
        "\n",
        "    def build_model(self):\n",
        "\n",
        "        self.G = Generator(self.batch_size,self.imsize, self.z_dim, self.g_conv_dim).cuda()\n",
        "        self.D = Discriminator(self.batch_size,self.imsize, self.d_conv_dim).cuda()\n",
        "        if self.parallel:\n",
        "            self.G = nn.DataParallel(self.G)\n",
        "            self.D = nn.DataParallel(self.D)\n",
        "\n",
        "        # Loss and optimizer\n",
        "        # self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr, [self.beta1, self.beta2])\n",
        "        self.g_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, self.G.parameters()), self.g_lr, [self.beta1, self.beta2])\n",
        "        self.d_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, self.D.parameters()), self.d_lr, [self.beta1, self.beta2])\n",
        "\n",
        "        self.c_loss = torch.nn.CrossEntropyLoss()\n",
        "\n",
        "\n",
        "    def build_tensorboard(self):\n",
        "        return\n",
        "        #from logger import Logger\n",
        "        #self.logger = Logger(self.log_path)\n",
        "\n",
        "    def load_pretrained_model(self):\n",
        "        self.G.load_state_dict(torch.load(os.path.join(\n",
        "            self.model_save_path, '{}_G.pth'.format(self.pretrained_model))))\n",
        "        self.D.load_state_dict(torch.load(os.path.join(\n",
        "            self.model_save_path, '{}_D.pth'.format(self.pretrained_model))))\n",
        "        print('loaded trained models (step: {})..!'.format(self.pretrained_model))\n",
        "\n",
        "    def reset_grad(self):\n",
        "        self.d_optimizer.zero_grad()\n",
        "        self.g_optimizer.zero_grad()\n",
        "\n",
        "    def save_sample(self, data_iter):\n",
        "        real_images, _ = next(data_iter)\n",
        "        save_image(denorm(real_images), os.path.join(self.sample_path, 'real.png'))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Exception ignored in: <bound method _MultiProcessingDataLoaderIter.__del__ of <torch.utils.data.dataloader._MultiProcessingDataLoaderIter object at 0x7f99809d5898>>\n",
            "Traceback (most recent call last):\n",
            "  File \"/usr/local/lib/python3.6/dist-packages/torch/utils/data/dataloader.py\", line 926, in __del__\n",
            "    self._shutdown_workers()\n",
            "  File \"/usr/local/lib/python3.6/dist-packages/torch/utils/data/dataloader.py\", line 906, in _shutdown_workers\n",
            "    w.join()\n",
            "  File \"/usr/lib/python3.6/multiprocessing/process.py\", line 124, in join\n",
            "    res = self._popen.wait(timeout)\n",
            "  File \"/usr/lib/python3.6/multiprocessing/popen_fork.py\", line 50, in wait\n",
            "    return self.poll(os.WNOHANG if timeout == 0.0 else 0)\n",
            "  File \"/usr/lib/python3.6/multiprocessing/popen_fork.py\", line 28, in poll\n",
            "    pid, sts = os.waitpid(self.pid, flag)\n",
            "KeyboardInterrupt: \n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mm964yyI4AiK",
        "colab_type": "text"
      },
      "source": [
        "##utils"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m6pzQjPG4B2S",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "import itertools, imageio, torch\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "from torchvision import datasets\n",
        "#from scipy.misc import imresize\n",
        "#import PIL\n",
        "#import pilutil\n",
        "\n",
        "\n",
        "def ones_target(size):\n",
        "  #returning tensor variable containing ones, with shape =size\n",
        "  data=Variable(torch.ones(size,1))\n",
        "  return data\n",
        "\n",
        "def zeros_target(size):\n",
        "  data=Variable(torch.zeros(size,1))\n",
        "  return data\n",
        "\n",
        "def images_to_vectors(images):\n",
        "  return images.view(images.size(0),3072) #parameters: (size to flatten(28 here, row size), no. of images to be flattened)\n",
        "\n",
        "def vectors_to_images(vectors):\n",
        "  return vectors.view(vectors.size(0),3,32,32) # to convert each row of vectors to 28x28 pixels images \n",
        "# parameteres: (initial flattened dim, no. of sub vectors to be converted to 2d, 2d dimensions of output )\n",
        "\n",
        "def show_result(G, x_, num_epoch, show = False, save = False, path = 'result.png'):\n",
        "    # G.eval()\n",
        "    y_ = G(x_)\n",
        "    y_=y_[:4]\n",
        "    y_ = y_.cpu()\n",
        "\n",
        "    size_figure_grid = 3\n",
        "    fig, ax = plt.subplots(x_.size()[0], size_figure_grid, figsize=(5, 5))\n",
        "    for i, j in itertools.product(range(x_.size()[0]), range(size_figure_grid)):\n",
        "        ax[i, j].get_xaxis().set_visible(False)\n",
        "        ax[i, j].get_yaxis().set_visible(False)\n",
        "\n",
        "    for i in range(y_.size()[0],2):\n",
        "\n",
        "        ax[i, 0].cla()\n",
        "        ax[i, 0].imshow((y_[i].numpy().transpose(1, 2, 0) + 1) / 2)\n",
        "        ax[i, 1].cla()\n",
        "        ax[i, 1].imshow((y_[i+1].numpy().transpose(1, 2, 0) + 1) / 2)\n",
        "\n",
        "    label = 'Epoch {0}'.format(num_epoch)\n",
        "    fig.text(0.5, 0.04, label, ha='center')\n",
        "\n",
        "    if save:\n",
        "        plt.savefig(path)\n",
        "\n",
        "    if show:\n",
        "        plt.show()\n",
        "    else:\n",
        "        plt.close()\n",
        "\n",
        "def show_train_hist(hist, show = False, save = False, path = 'Train_hist.png'):\n",
        "    x = range(len(hist['D_losses']))\n",
        "\n",
        "    y1 = hist['D_losses']\n",
        "    y2 = hist['G_losses']\n",
        "\n",
        "    plt.plot(x, y1, label='D_loss')\n",
        "    plt.plot(x, y2, label='G_loss')\n",
        "\n",
        "    plt.xlabel('Iter')\n",
        "    plt.ylabel('Loss')\n",
        "\n",
        "    plt.legend(loc=4)\n",
        "    plt.grid(True)\n",
        "    plt.tight_layout()\n",
        "\n",
        "    if save:\n",
        "        plt.savefig(path)\n",
        "\n",
        "    if show:\n",
        "        plt.show()\n",
        "    else:\n",
        "        plt.close()\n",
        "\n",
        "def generate_animation(root, model, opt):\n",
        "    images = []\n",
        "    for e in range(opt.train_epoch):\n",
        "        img_name = root + 'Fixed_results/' + model + str(e + 1) + '.png'\n",
        "        images.append(imageio.imread(img_name))\n",
        "    imageio.mimsave(root + model + 'generate_animation.gif', images, fps=5)\n",
        "\n",
        "def data_load(path, subfolder, transform, batch_size, shuffle=True):\n",
        "    dset = datasets.ImageFolder(path, transform)\n",
        "    ind = dset.class_to_idx[subfolder]\n",
        "\n",
        "    n = 0\n",
        "    for i in range(dset.__len__()):\n",
        "        if ind != dset.imgs[n][1]:\n",
        "            del dset.imgs[n]\n",
        "            n -= 1\n",
        "\n",
        "        n += 1\n",
        "\n",
        "    return torch.utils.data.DataLoader(dset, batch_size=batch_size, shuffle=shuffle)\n",
        "\n",
        "def imgs_resize(imgs, resize_scale = 286):\n",
        "    outputs = torch.FloatTensor(imgs.size()[0], imgs.size()[1], resize_scale, resize_scale)\n",
        "    for i in range(imgs.size()[0]):\n",
        "        img = imresize(imgs[i].numpy(), [resize_scale, resize_scale])\n",
        "        outputs[i] = torch.FloatTensor((img.transpose(2, 0, 1).astype(np.float32).reshape(-1, imgs.size()[1], resize_scale, resize_scale) - 127.5) / 127.5)\n",
        "\n",
        "    return outputs\n",
        "\n",
        "def random_crop(imgs1, imgs2, crop_size = 256):\n",
        "    outputs1 = torch.FloatTensor(imgs1.size()[0], imgs1.size()[1], crop_size, crop_size)\n",
        "    outputs2 = torch.FloatTensor(imgs2.size()[0], imgs2.size()[1], crop_size, crop_size)\n",
        "    for i in range(imgs1.size()[0]):\n",
        "        img1 = imgs1[i]\n",
        "        img2 = imgs2[i]\n",
        "        rand1 = np.random.randint(0, imgs1.size()[2] - crop_size)\n",
        "        rand2 = np.random.randint(0, imgs2.size()[2] - crop_size)\n",
        "        outputs1[i] = img1[:, rand1: crop_size + rand1, rand2: crop_size + rand2]\n",
        "        outputs2[i] = img2[:, rand1: crop_size + rand1, rand2: crop_size + rand2]\n",
        "\n",
        "    return outputs1, outputs2\n",
        "\n",
        "def random_fliplr(imgs1, imgs2):\n",
        "    outputs1 = torch.FloatTensor(imgs1.size())\n",
        "    outputs2 = torch.FloatTensor(imgs2.size())\n",
        "    for i in range(imgs1.size()[0]):\n",
        "        if torch.rand(1)[0] < 0.5:\n",
        "            img1 = torch.FloatTensor(\n",
        "                (np.fliplr(imgs1[i].numpy().transpose(1, 2, 0)).transpose(2, 0, 1).reshape(-1, imgs1.size()[1], imgs1.size()[2], imgs1.size()[3]) + 1) / 2)\n",
        "            outputs1[i] = (img1 - 0.5) / 0.5\n",
        "            img2 = torch.FloatTensor(\n",
        "                (np.fliplr(imgs2[i].numpy().transpose(1, 2, 0)).transpose(2, 0, 1).reshape(-1, imgs2.size()[1], imgs2.size()[2], imgs2.size()[3]) + 1) / 2)\n",
        "            outputs2[i] = (img2 - 0.5) / 0.5\n",
        "        else:\n",
        "            outputs1[i] = imgs1[i]\n",
        "            outputs2[i] = imgs2[i]\n",
        "\n",
        "    return outputs1, outputs2\n",
        "\n",
        "def imshow(img):\n",
        "    img = img / 2 + 0.5     # unnormalize\n",
        "    npimg = img.numpy()\n",
        "    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
        "    plt.show()\n",
        "    \n",
        "def noise(size):\n",
        "  n=Variable(torch.randn(size,200))\n",
        "  return n\n",
        " \n",
        "def weights_init(m):\n",
        "    classname = m.__class__.__name__\n",
        "    if classname.find('Conv') != -1:\n",
        "        m.weight.data.normal_(0.0, 0.02)\n",
        "    elif classname.find('BatchNorm') != -1:\n",
        "        m.weight.data.normal_(1.0, 0.02)\n",
        "        m.bias.data.fill_(0)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}