10-celeba-gan-experiment.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xJ9h-C3iE90L"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "if not os.path.isdir(\"apex\"):\n",
    "    !git clone https://github.com/NVIDIA/apex.git\n",
    "    !pip install -v --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" ./apex\n",
    "\n",
    "!pip install matplotlib scipy\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KRKNdbcmE90R"
   },
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "#%matplotlib inline\n",
    "import argparse\n",
    "import random\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.parallel\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.optim as optim\n",
    "import torch.utils.data\n",
    "import torchvision.datasets as dset\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.utils as vutils\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "from IPython.display import HTML\n",
    "from torch.nn.utils import spectral_norm\n",
    "from apex import amp\n",
    "\n",
    "# Set random seem for reproducibility\n",
    "manualSeed = 999\n",
    "#manualSeed = random.randint(1, 10000) # use if you want new results\n",
    "print(\"Random Seed: \", manualSeed)\n",
    "random.seed(manualSeed)\n",
    "torch.manual_seed(manualSeed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "63NvtMygE90V"
   },
   "outputs": [],
   "source": [
    "# Root directory for dataset\n",
    "dataroot = \"/workspace/imgs\"\n",
    "checkpoint_file_path = './checkpoint.tar'\n",
    "\n",
    "# Number of workers for dataloader\n",
    "workers = 4\n",
    "\n",
    "# Batch size during training\n",
    "batch_size = 128\n",
    "\n",
    "# Spatial size of training images. All images will be resized to this\n",
    "#   size using a transformer.\n",
    "image_size = 64\n",
    "\n",
    "# Number of channels in the training images. For color images this is 3\n",
    "nc = 3\n",
    "\n",
    "# Size of z latent vector (i.e. size of generator input)\n",
    "nz = 100\n",
    "\n",
    "# Size of feature maps in generator\n",
    "ngf = 64\n",
    "\n",
    "# Size of feature maps in discriminator\n",
    "ndf = 64\n",
    "\n",
    "# Number of training epochs\n",
    "num_epochs = 50\n",
    "\n",
    "# Learning rate for optimizers\n",
    "lr_d = 2e-4\n",
    "lr_g = 5e-5\n",
    "betas = (0.0, 0.999)\n",
    "d_steps_per_g_step = 2\n",
    "\n",
    "# decay for exponential moving average\n",
    "ema_decay = 0.9999\n",
    "\n",
    "# Number of GPUs available. Use 0 for CPU mode.\n",
    "ngpu = torch.cuda.device_count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tICYdg68F-mi"
   },
   "outputs": [],
   "source": [
    "# get the data\n",
    "!apt update\n",
    "!apt install wget unzip\n",
    "if not os.path.isdir(dataroot):\n",
    "  !wget --load-cookies /tmp/cookies.txt \"https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1g9tAydWEN0PW-BYxdhNuGG9vXrdNzffA' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\\1\\n/p')&id=1g9tAydWEN0PW-BYxdhNuGG9vXrdNzffA\" -O img_align_celeba.zip && rm -rf /tmp/cookies.txt\n",
    "  !unzip img_align_celeba.zip -d /workspace/imgs\n",
    "\n",
    "# load checkpoint data if it exists\n",
    "have_checkpoint = False\n",
    "checkpoint = None\n",
    "completed_epochs = 0\n",
    "if os.path.exists(checkpoint_file_path):\n",
    "  checkpoint = torch.load(checkpoint_file_path)\n",
    "  have_checkpoint = True\n",
    "  completed_epochs = checkpoint['completed_epochs']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate frechet inception distance\n",
    "from torchvision.utils import save_image\n",
    "from shutil import rmtree\n",
    "num_batches = 200\n",
    "real_images_for_fid_dir = \"./fid_images/real/\"\n",
    "fake_images_for_fid_dir = \"./fid_images/fake/\"\n",
    "\n",
    "if not os.path.isdir(\"pytorch-fid\"):\n",
    "    ! git clone https://github.com/mseitzer/pytorch-fid.git\n",
    "\n",
    "def print_frechet_inception_distance():\n",
    "    torch.cuda.empty_cache()\n",
    "    rmtree(real_images_for_fid_dir, ignore_errors=True)\n",
    "    rmtree(fake_images_for_fid_dir, ignore_errors=True)\n",
    "    os.makedirs(real_images_for_fid_dir)\n",
    "    os.makedirs(fake_images_for_fid_dir)\n",
    "\n",
    "    real_images = torch.cat([x[0] for _, x in zip(range(num_batches), iter(dataloader))])\n",
    "\n",
    "    for i in range(real_images.size(0)):\n",
    "        save_image(real_images[i, :, :, :], real_images_for_fid_dir + '{}.png'.format(i))\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for batch_num in range(num_batches):\n",
    "            generator_inputs = torch.randn(batch_size, nz).to(device)\n",
    "            fake_images = netG(generator_inputs).detach().cpu()\n",
    "            for i in range(fake_images.size(0)):\n",
    "                save_image(fake_images[i, :, :, :], fake_images_for_fid_dir + '{}.png'.format(i + batch_num * batch_size))\n",
    "\n",
    "    ! pytorch-fid/fid_score.py --batch-size 128 --gpu 0 ./fid_images/real/ ./fid_images/fake/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ayFbpLj1E90Z"
   },
   "outputs": [],
   "source": [
    "# We can use an image folder dataset the way we have it setup.\n",
    "# Create the dataset\n",
    "dataset = dset.ImageFolder(root=dataroot,\n",
    "                           transform=transforms.Compose([\n",
    "                               transforms.Resize(image_size),\n",
    "                               transforms.CenterCrop(image_size),\n",
    "                               transforms.ToTensor(),\n",
    "                               transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n",
    "                           ]))\n",
    "# Create the dataloader\n",
    "dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,\n",
    "                                         shuffle=True, num_workers=workers)\n",
    "\n",
    "# Decide which device we want to run on\n",
    "device = torch.device(\"cuda:0\" if (torch.cuda.is_available() and ngpu > 0) else \"cpu\")\n",
    "\n",
    "# Plot some training images\n",
    "real_batch = next(iter(dataloader))\n",
    "plt.figure(figsize=(8,8))\n",
    "plt.axis(\"off\")\n",
    "plt.title(\"Training Images\")\n",
    "plt.imshow(np.transpose(vutils.make_grid(real_batch[0].to(device)[:64], padding=2, normalize=True).cpu(),(1,2,0)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BatxsavEE90d"
   },
   "outputs": [],
   "source": [
    "# custom weights initialization called on netG and netD\n",
    "def weights_init(m):\n",
    "    classname = m.__class__.__name__\n",
    "    if classname.find('Conv') != -1:\n",
    "        nn.init.normal_(m.weight.data, 0.0, 0.02)\n",
    "    elif classname.find('BatchNorm') != -1:\n",
    "        nn.init.normal_(m.weight.data, 1.0, 0.02)\n",
    "        nn.init.constant_(m.bias.data, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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 = spectral_norm(nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1))\n",
    "        self.key_conv = spectral_norm(nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1))\n",
    "        self.value_conv = spectral_norm(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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XpueL3UDE90h"
   },
   "outputs": [],
   "source": [
    "# Generator Code\n",
    "class ResBlockUp(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.in_channels = in_channels\n",
    "        self.out_channels = out_channels\n",
    "        \n",
    "        self.skip_conv = spectral_norm(nn.Conv2d(self.in_channels, self.out_channels, 1, 1, 0, bias=False))\n",
    "        \n",
    "        self.main_conv1 = spectral_norm(nn.Conv2d(self.in_channels, self.out_channels, 3, 1, 1, bias=False))\n",
    "        self.main_bn = nn.BatchNorm2d(self.out_channels)\n",
    "        self.main_relu = nn.ReLU(True)\n",
    "        self.main_conv2 = spectral_norm(nn.Conv2d(self.out_channels, self.out_channels, 3, 1, 1, bias=False))\n",
    "        \n",
    "        self.out_bn = nn.BatchNorm2d(self.out_channels)\n",
    "        self.out_relu = nn.ReLU(True)\n",
    "        \n",
    "        \n",
    "    def forward(self, input):\n",
    "        main = input\n",
    "        skip = input\n",
    "        \n",
    "        skip = nn.functional.interpolate(skip, scale_factor=2)\n",
    "        skip = self.skip_conv(skip)\n",
    "        \n",
    "        main = nn.functional.interpolate(main, scale_factor=2)\n",
    "        main = self.main_conv1(main)\n",
    "        main = self.main_bn(main)\n",
    "        main = self.main_relu(main)\n",
    "        main = self.main_conv2(main)\n",
    "        \n",
    "        out = main + skip\n",
    "        out = self.out_bn(out)\n",
    "        out = self.out_relu(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "class Generator(nn.Module):\n",
    "    def __init__(self, ngpu):\n",
    "        super(Generator, self).__init__()\n",
    "        self.ngpu = ngpu\n",
    "        \n",
    "        self.linear = nn.Sequential(\n",
    "            spectral_norm(nn.Linear(nz, 4 * 4 * ngf * 16))\n",
    "        )\n",
    "        \n",
    "        self.main = nn.Sequential(\n",
    "            nn.BatchNorm2d(ngf * 16),\n",
    "            nn.ReLU(True),\n",
    "            ResBlockUp(ngf * 16, ngf * 8),\n",
    "            ResBlockUp(ngf * 8, ngf * 4),\n",
    "            ResBlockUp(ngf * 4, ngf * 2),\n",
    "            Self_Attn(ngf * 2, 'relu'),\n",
    "            ResBlockUp(ngf * 2, ngf),\n",
    "            spectral_norm(nn.Conv2d( ngf, nc, 3, 1, 1, bias=False)),\n",
    "            nn.Tanh()\n",
    "        )\n",
    "        \n",
    "\n",
    "    def forward(self, input):\n",
    "        out = self.linear(input)\n",
    "        out = out.view(-1, 1024, 4, 4)\n",
    "        out = self.main(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Y7zSNkBkE90m"
   },
   "outputs": [],
   "source": [
    "# Create the generator\n",
    "netG = Generator(ngpu).to(device)\n",
    "\n",
    "optimizerG = optim.Adam(netG.parameters(), lr=lr_g, betas=betas)\n",
    "\n",
    "netG, optimizerG = amp.initialize(\n",
    "    netG, optimizerG, opt_level=\"O2\", \n",
    "    keep_batchnorm_fp32=True, loss_scale=\"dynamic\"\n",
    ")\n",
    "\n",
    "# # Create EMA generator\n",
    "# netG_ema = Generator(ngpu).to(device)\n",
    "\n",
    "# Handle multi-gpu if desired\n",
    "if (device.type == 'cuda') and (ngpu > 1):\n",
    "    netG = nn.DataParallel(netG, list(range(ngpu)))\n",
    "\n",
    "if not have_checkpoint:      \n",
    "  # Apply the weights_init function to randomly initialize all weights\n",
    "  #  to mean=0, stdev=0.2.\n",
    "  netG.apply(weights_init)\n",
    "else:\n",
    "  netG.load_state_dict(checkpoint['netG_state_dict'])\n",
    "#   netG_ema.load_state_dict(checkpoint['netG_ema_state_dict'])\n",
    "\n",
    "# Print the model\n",
    "print(netG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZUlsY7HgE90q"
   },
   "outputs": [],
   "source": [
    "class ResBlockDown(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.in_channels = in_channels\n",
    "        self.out_channels = out_channels\n",
    "        \n",
    "        self.skip_conv = spectral_norm(nn.Conv2d(self.in_channels, self.out_channels, 1, 1, 0))\n",
    "        \n",
    "        self.main_conv1 = spectral_norm(nn.Conv2d(self.in_channels, self.out_channels, 3, 1, 1))\n",
    "        self.main_relu = nn.ReLU(True)\n",
    "        self.main_conv2 = spectral_norm(nn.Conv2d(self.out_channels, self.out_channels, 3, 1, 1))\n",
    "        \n",
    "        self.out_relu = nn.ReLU(True)\n",
    "        \n",
    "        \n",
    "    def forward(self, input):\n",
    "        main = input\n",
    "        skip = input\n",
    "        \n",
    "        skip = self.skip_conv(skip)\n",
    "        skip = nn.functional.avg_pool2d(skip, 2)\n",
    "        \n",
    "        main = self.main_conv1(main)\n",
    "        main = self.main_relu(main)\n",
    "        main = self.main_conv2(main)\n",
    "        main = nn.functional.avg_pool2d(main, 2)\n",
    "        \n",
    "        out = main + skip\n",
    "        out = self.out_relu(out)\n",
    "        return out\n",
    "\n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, ngpu):\n",
    "        super(Discriminator, self).__init__()\n",
    "        self.ngpu = ngpu\n",
    "        self.main = nn.Sequential(\n",
    "            # input is (nc) x 64 x 64\n",
    "            ResBlockDown(nc, ndf),\n",
    "            Self_Attn(ndf, 'relu'),\n",
    "            # state size. (ndf) x 32 x 32\n",
    "            ResBlockDown(ndf, ndf * 2),\n",
    "            # state size. (ndf*2) x 16 x 16\n",
    "            ResBlockDown(ndf * 2, ndf * 4),\n",
    "            # state size. (ndf*4) x 8 x 8\n",
    "            ResBlockDown(ndf * 4, ndf * 8),\n",
    "            # state size. (ndf*8) x 4 x 4\n",
    "        )\n",
    "        \n",
    "        self.linear = nn.Sequential(\n",
    "            spectral_norm(nn.Linear(ndf * 8, 1))\n",
    "        )\n",
    "\n",
    "    def forward(self, input):\n",
    "        out = self.main(input)\n",
    "        out = torch.sum(out.view(out.size(0), out.size(1), -1), dim=2)\n",
    "        return self.linear(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "x9HEb0a4E90v"
   },
   "outputs": [],
   "source": [
    "# Create the Discriminator\n",
    "netD = Discriminator(ngpu).to(device)\n",
    "\n",
    "optimizerD = optim.Adam(netD.parameters(), lr=lr_d, betas=betas)\n",
    "\n",
    "netD, optimizerD = amp.initialize(\n",
    "    netD, optimizerD, opt_level=\"O2\", \n",
    "    keep_batchnorm_fp32=True, loss_scale=\"dynamic\"\n",
    ")\n",
    "\n",
    "# Handle multi-gpu if desired\n",
    "if (device.type == 'cuda') and (ngpu > 1):\n",
    "  netD = nn.DataParallel(netD, list(range(ngpu)))\n",
    "\n",
    "if not have_checkpoint:      \n",
    "  # Apply the weights_init function to randomly initialize all weights\n",
    "  #  to mean=0, stdev=0.2.\n",
    "  netD.apply(weights_init)\n",
    "else:\n",
    "  netD.load_state_dict(checkpoint['netD_state_dict'])\n",
    "\n",
    "# Print the model\n",
    "print(netD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Uk04ng_tE900"
   },
   "outputs": [],
   "source": [
    "# load optimizer and amp checkpoints\n",
    "if have_checkpoint:\n",
    "    optimizerD.load_state_dict(checkpoint['optimizerD_state_dict'])\n",
    "    optimizerG.load_state_dict(checkpoint['optimizerG_state_dict'])\n",
    "    amp.load_state_dict(checkpoint['amp'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uPFCTuKaE903",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Training Loop\n",
    "from datetime import datetime\n",
    "import math\n",
    "\n",
    "# Lists to keep track of progress\n",
    "G_losses = []\n",
    "D_losses = []\n",
    "iters = 0\n",
    "\n",
    "final_epoch = num_epochs + completed_epochs\n",
    "\n",
    "print(\"Starting Training Loop...\")\n",
    "# For each epoch\n",
    "for epoch in range(num_epochs):\n",
    "    start_time = datetime.now()\n",
    "    \n",
    "    # For each batch in the dataloader\n",
    "    for i, data in enumerate(dataloader, 0):\n",
    "        for j in range(d_steps_per_g_step):\n",
    "            netD.zero_grad()\n",
    "      \n",
    "            ############################\n",
    "            # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))\n",
    "            ###########################\n",
    "            ## Train with all-real batch\n",
    "            # Format batch\n",
    "            real_cpu = data[0].to(device)\n",
    "            # Forward pass real batch through D\n",
    "            output = netD(real_cpu).view(-1)\n",
    "            # Calculate loss on all-real batch\n",
    "            errD_real = torch.nn.ReLU()(1.0 - output).mean()\n",
    "            D_x = errD_real.item()\n",
    "\n",
    "            ## Train with all-fake batch\n",
    "            # Generate batch of latent vectors\n",
    "            noise = torch.randn(batch_size, nz).to(device)\n",
    "            # Generate fake image batch with G\n",
    "            fake = netG(noise)\n",
    "            # Classify all fake batch with D\n",
    "            output = netD(fake.detach()).view(-1)\n",
    "            # Calculate D's loss on the all-fake batch\n",
    "            errD_fake = torch.nn.ReLU()(1.0 + output).mean()\n",
    "            D_G_z1 = errD_fake.item()\n",
    "\n",
    "            errD = errD_fake + errD_real\n",
    "\n",
    "            with amp.scale_loss(errG, optimizerG) as scaled_loss:\n",
    "                # backward only called for G to prevent triggering DistributedDataParallel reduction on D,\n",
    "                # which could cause problems for two DDP wrapped networks in the same process\n",
    "                G_out_grad = torch.autograd.grad(scaled_loss, fake)\n",
    "                torch.autograd.backward(fake, G_out_grad)\n",
    "            \n",
    "            # Update D\n",
    "            optimizerD.step()\n",
    "\n",
    "        netG.zero_grad()\n",
    "       \n",
    "        ############################\n",
    "        # (2) Update G network: maximize log(D(G(z)))\n",
    "        ###########################\n",
    "\n",
    "        # Since we just updated D, perform another forward pass of all-fake batch through D\n",
    "        output = netD(fake).view(-1)\n",
    "        # Calculate G's loss based on this output\n",
    "        errG = - output.mean()\n",
    "        # Calculate gradients for G\n",
    "        with amp.scale_loss(errG, optimizerG) as scaled_loss:\n",
    "            scaled_loss.backward()\n",
    "        D_G_z2 = errG.item()\n",
    "     \n",
    "        # Update G\n",
    "        optimizerG.step()\n",
    "\n",
    "        # update netG_ema\n",
    "#         with torch.no_grad():\n",
    "#             for key in netG.state_dict():\n",
    "#                 netG_ema.state_dict()[key].data.copy_(netG_ema.state_dict()[key] * ema_decay\n",
    "#                                                      + netG.state_dict()[key] * (1 - ema_decay))\n",
    "        \n",
    "        # Output training stats\n",
    "        if i % math.ceil(len(dataloader) / 5) == 0:\n",
    "            print('[%d/%d][%d/%d]\\tLoss_D: %.4f\\tLoss_G: %.4f\\tD(x): %.4f\\tD(G(z)): %.4f / %.4f'\n",
    "                  % (1 + completed_epochs, final_epoch, i, len(dataloader),\n",
    "                     errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))\n",
    "        \n",
    "        # Save Losses for plotting later\n",
    "        G_losses.append(errG.item())\n",
    "        D_losses.append(errD.item())\n",
    "            \n",
    "        iters += 1\n",
    "\n",
    "    # save models and optimizers\n",
    "    completed_epochs += 1\n",
    "    torch.save({\n",
    "                'netG_state_dict': netG.state_dict(),\n",
    "                'netD_state_dict': netD.state_dict(),\n",
    "                'optimizerG_state_dict': optimizerG.state_dict(),\n",
    "                'optimizerD_state_dict': optimizerD.state_dict(),\n",
    "#                 'netG_ema_state_dict': netG_ema.state_dict(),\n",
    "                'completed_epochs': completed_epochs,\n",
    "                'amp': amp.state_dict()\n",
    "    }, checkpoint_file_path)\n",
    "\n",
    "    print(\"training time: \" + str(datetime.now() - start_time))\n",
    "    \n",
    "    if completed_epochs % 5 == 0:\n",
    "        print_frechet_inception_distance()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8NODnTqSE906"
   },
   "source": [
    "Results\n",
    "-------\n",
    "\n",
    "Finally, lets check out how we did. Here, we will look at three\n",
    "different results. First, we will see how D and G’s losses changed\n",
    "during training. Second, we will visualize G’s output on the fixed_noise\n",
    "batch for every epoch. And third, we will look at a batch of real data\n",
    "next to a batch of fake data from G.\n",
    "\n",
    "**Loss versus training iteration**\n",
    "\n",
    "Below is a plot of D & G’s losses versus training iterations.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZyrdNTyhE906"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10,5))\n",
    "plt.title(\"Generator and Discriminator Loss During Training\")\n",
    "plt.plot(G_losses,label=\"G\")\n",
    "plt.plot(D_losses,label=\"D\")\n",
    "plt.xlabel(\"iterations\")\n",
    "plt.ylabel(\"Loss\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6Vq1Nwi1E91A"
   },
   "source": [
    "**Real Images vs. Fake Images**\n",
    "\n",
    "Finally, lets take a look at some real images and fake images side by\n",
    "side.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MofRFv_JE91B"
   },
   "outputs": [],
   "source": [
    "# Grab a batch of real images from the dataloader\n",
    "real_batch = next(iter(dataloader))\n",
    "\n",
    "# Plot the real images\n",
    "plt.figure(figsize=(15,15))\n",
    "plt.subplot(1,2,1)\n",
    "plt.axis(\"off\")\n",
    "plt.title(\"Real Images\")\n",
    "plt.imshow(np.transpose(vutils.make_grid(real_batch[0].to(device)[:64], padding=5, normalize=True).cpu(),(1,2,0)))\n",
    "\n",
    "# Plot the fake images from the last epoch\n",
    "with torch.no_grad():\n",
    "    fake = netG(torch.randn(64, nz).to(device)).detach().cpu()\n",
    "\n",
    "fakes = vutils.make_grid(fake, padding=2, normalize=True)\n",
    "\n",
    "plt.subplot(1,2,2)\n",
    "plt.axis(\"off\")\n",
    "plt.title(\"Fake Images\")\n",
    "plt.imshow(np.transpose(fakes,(1,2,0)))\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "dcgan_faces_tutorial.ipynb",
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}