cDCGAN_Pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "import imageio\n",
    "import numpy as np\n",
    "import itertools\n",
    "import torch.optim as optim\n",
    "import matplotlib.pyplot as plt\n",
    "import os,time\n",
    "\n",
    "from torchvision import transforms, datasets\n",
    "from torch.autograd import Variable\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 0.0002\n",
    "img_size = 32\n",
    "batch_size = 512\n",
    "num_epochs = 40\n",
    "save_dir = '/home/abdullah/Documents/ai/models/Own implementation of Gan/Conditional Gan/Model results'\n",
    "G_input_dim = 100\n",
    "G_output_dim = 1\n",
    "D_input_dim = 1\n",
    "D_output_dim = 1\n",
    "label_dim = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_var(x):\n",
    "    x = x.cuda()\n",
    "    return Variable(x)\n",
    "\n",
    "def denorm(x):\n",
    "    out = (x+1)/2\n",
    "    return out.clamp(0,1)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_loss(d_losses, g_losses, num_epoch, save=True ,show=False):\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.set_xlim(0, num_epochs)\n",
    "    ax.set_ylim(0, max(np.max(g_losses), np.max(d_losses))*1.1)\n",
    "    plt.xlabel('Epoch {0}'.format(num_epoch + 1))\n",
    "    plt.ylabel('Loss values')\n",
    "    plt.plot(d_losses, label='Discriminator')\n",
    "    plt.plot(g_losses, label='Generator')\n",
    "    plt.legend()\n",
    "\n",
    "    # save figure\n",
    "    if save:\n",
    "        save_fn = save_dir + 'MNIST_cDCGAN_losses_epoch_{:d}'.format(num_epoch + 1) + '.png'\n",
    "        plt.savefig(save_fn)\n",
    "\n",
    "    if show:\n",
    "        plt.show()\n",
    "    else:\n",
    "        plt.close()\n",
    "\n",
    "def plot_result(generator, noise, label, num_epoch, save=True, show=False, fig_size=(5, 5)):\n",
    "    generator.eval()\n",
    "\n",
    "    noise = Variable(noise.cuda())\n",
    "    label = Variable(label.cuda())\n",
    "    gen_image = generator(noise, label)\n",
    "    gen_image = denorm(gen_image)\n",
    "\n",
    "    generator.train()\n",
    "\n",
    "    n_rows = np.sqrt(noise.size()[0]).astype(np.int32)\n",
    "    n_cols = np.sqrt(noise.size()[0]).astype(np.int32)\n",
    "    fig, axes = plt.subplots(n_rows, n_cols, figsize=fig_size)\n",
    "    for ax, img in zip(axes.flatten(), gen_image):\n",
    "        ax.axis('off')\n",
    "        ax.set_adjustable('box-forced')\n",
    "        ax.imshow(img.cpu().data.view(img_size, img_size).numpy(), cmap='gray', aspect='equal')\n",
    "    plt.subplots_adjust(wspace=0, hspace=0)\n",
    "    title = 'Epoch {0}'.format(num_epoch+1)\n",
    "    fig.text(0.5, 0.04, title, ha='center')\n",
    "\n",
    "    # save figure\n",
    "    if save:\n",
    "        save_fn = save_dir + 'MNIST_cDCGAN_epoch_{:d}'.format(num_epoch+1) + '.png'\n",
    "        plt.savefig(save_fn)\n",
    "\n",
    "    if show:\n",
    "        plt.show()\n",
    "    else:\n",
    "        plt.close()\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/abdullah/miniconda3/lib/python3.6/site-packages/torchvision/transforms/transforms.py:188: UserWarning: The use of the transforms.Scale transform is deprecated, please use transforms.Resize instead.\n",
      "  \"please use transforms.Resize instead.\")\n"
     ]
    }
   ],
   "source": [
    "transform = transforms.Compose([\n",
    "    transforms.Scale(img_size),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean = (0.5,0.5,.5), std = (.5,.5,.5))\n",
    "])\n",
    "\n",
    "\n",
    "mnist_data = datasets.MNIST(root=\"/home/abdullah/Documents/ai/models/Own implementation of Gan/Untitled Folder/DCGAN/mnist\",\n",
    "                         train=True,\n",
    "                         transform=transform,\n",
    "                         download=False)\n",
    "\n",
    "data_loader = torch.utils.data.DataLoader(dataset=mnist_data,\n",
    "                                          batch_size=batch_size,\n",
    "                                          shuffle=True)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hello\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "discriminator(\n",
       "  (conv1_x): Conv2d(1, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv1_y): Conv2d(10, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv2): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (batch_norm_2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv3): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (batch_norm_3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv4): Conv2d(512, 1, kernel_size=(4, 4), stride=(1, 1))\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class generator(nn.Module):\n",
    "    def __init__(self,d=128):\n",
    "        \n",
    "        super(generator,self).__init__()\n",
    "        \n",
    "        self.deconv1_z = nn.ConvTranspose2d( 100 , d*2 , 4 , 1 , 0)\n",
    "        self.deconv1_y = nn.ConvTranspose2d( 10, d*2 , 4 , 1, 0)\n",
    "        \n",
    "        self.batch_norm_1 = nn.BatchNorm2d(d*2)\n",
    "        \n",
    "        self.deconv2 = nn.ConvTranspose2d(d*4 , d*2, 4, 2, 1)\n",
    "        self.batch_norm_2 = nn.BatchNorm2d(d*2)\n",
    "        \n",
    "        self.deconv3 = nn.ConvTranspose2d(d*2 , d , 4 , 2 , 1)\n",
    "        self.batch_norm_3 = nn.BatchNorm2d(d)\n",
    "        \n",
    "        self.deconv4 = nn.ConvTranspose2d(d , 1 , 4, 2 , 1)\n",
    "        \n",
    "    def weight_init(self, mean, std):\n",
    "        for m in self._modules:\n",
    "            normal_init(self._modules[m], mean , std)\n",
    "            \n",
    "    def forward(self , input , label):\n",
    "        x = F.relu(self.batch_norm_1(self.deconv1_z(input)))\n",
    "        y = F.relu(self.batch_norm_1(self.deconv1_y(label)))\n",
    "        \n",
    "        x = torch.cat([x,y],1)\n",
    "        \n",
    "        x = F.relu(self.batch_norm_2(self.deconv2(x)))\n",
    "        \n",
    "        x = F.relu(self.batch_norm_3(self.deconv3(x)))\n",
    "        \n",
    "        x = F.tanh(self.deconv4(x))\n",
    "        \n",
    "        return x\n",
    "\n",
    "class discriminator(nn.Module):\n",
    "    def __init__(self, d = 128):\n",
    "        \n",
    "        super( discriminator, self).__init__()\n",
    "        print(\"hello\")\n",
    "        self.conv1_x = nn.Conv2d(1 , d//2 , 4, 2, 1)\n",
    "        self.conv1_y = nn.Conv2d(10 , d//2 , 4, 2, 1)\n",
    "        \n",
    "        self.conv2 = nn.Conv2d( d, d*2, 4, 2, 1)\n",
    "        self.batch_norm_2 = nn.BatchNorm2d(d*2)\n",
    "        \n",
    "        self.conv3 = nn.Conv2d( d*2, d*4, 4 , 2, 1)\n",
    "        self.batch_norm_3 = nn.BatchNorm2d(d*4)\n",
    "        \n",
    "        self.conv4 = nn.Conv2d( d*4, 1, 4, 1, 0)\n",
    "        \n",
    "    def weight_init(self, mean , std):\n",
    "        for m in self._modules:\n",
    "            normal_init(self._modules[m], mean , std)\n",
    "            \n",
    "    def forward(self, input  , label):\n",
    "        x = F.leaky_relu(self.conv1_x(input),0.2)\n",
    "        y = F.leaky_relu(self.conv1_y(label),0.2)\n",
    "        \n",
    "        x = torch.cat([x,y],1)\n",
    "        \n",
    "        x = F.leaky_relu(self.batch_norm_2(self.conv2(x)),0.2)\n",
    "        \n",
    "        x = F.leaky_relu(self.batch_norm_3(self.conv3(x)),0.2)\n",
    "        \n",
    "        x = F.sigmoid(self.conv4(x))\n",
    "        \n",
    "        return x\n",
    "    \n",
    "    \n",
    "def normal_init(m, mean, std):\n",
    "    if isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Conv2d):\n",
    "        m.weight.data.normal_(mean, std)\n",
    "        m.bias.data.zero_()       \n",
    "        \n",
    "        \n",
    "generator = generator()\n",
    "discriminator = discriminator()\n",
    "\n",
    "generator.weight_init(mean = 0 , std = 0.02)\n",
    "discriminator.weight_init(mean = 0 , std = 0.02)\n",
    "\n",
    "generator.cuda()\n",
    "discriminator.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimizations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "opt_Gen = optim.Adam(generator.parameters(),lr = lr , betas = (.5, .999))\n",
    "opt_Disc = optim.Adam(discriminator.parameters(),lr = lr , betas = (.5, .999))\n",
    "\n",
    "criterion = torch.nn.BCELoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/abdullah/miniconda3/lib/python3.6/site-packages/ipykernel_launcher.py:83: UserWarning: invalid index of a 0-dim tensor. This will be an error in PyTorch 0.5. Use tensor.item() to convert a 0-dim tensor to a Python number\n",
      "/home/abdullah/miniconda3/lib/python3.6/site-packages/ipykernel_launcher.py:84: UserWarning: invalid index of a 0-dim tensor. This will be an error in PyTorch 0.5. Use tensor.item() to convert a 0-dim tensor to a Python number\n",
      "/home/abdullah/miniconda3/lib/python3.6/site-packages/ipykernel_launcher.py:87: UserWarning: invalid index of a 0-dim tensor. This will be an error in PyTorch 0.5. Use tensor.item() to convert a 0-dim tensor to a Python number\n"
     ]
    },
    
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-11-4fc5c063c2a4>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     74\u001b[0m         \u001b[0mc_fill\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mVariable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfill\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mc_\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcuda\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     75\u001b[0m         \u001b[0mD_fake_decision\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdiscriminator\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgenerator_image\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mc_fill\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 76\u001b[0;31m         \u001b[0mG_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mD_fake_decision\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_real\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     77\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     78\u001b[0m         \u001b[0;31m# Optimization\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.6/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    489\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    490\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 491\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    492\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\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    493\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.6/site-packages/torch/nn/modules/loss.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input, target)\u001b[0m\n\u001b[1;32m    431\u001b[0m         return F.binary_cross_entropy(input, target, weight=self.weight,\n\u001b[1;32m    432\u001b[0m                                       \u001b[0msize_average\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize_average\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 433\u001b[0;31m                                       reduce=self.reduce)\n\u001b[0m\u001b[1;32m    434\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    435\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.6/site-packages/torch/nn/functional.py\u001b[0m in \u001b[0;36mbinary_cross_entropy\u001b[0;34m(input, target, weight, size_average, reduce)\u001b[0m\n\u001b[1;32m   1481\u001b[0m         \u001b[0mweight\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mweight\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexpand\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnew_size\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1482\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1483\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_C\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_nn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbinary_cross_entropy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msize_average\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreduce\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1484\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1485\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# num_test_samples = 10*10\n",
    "temp_noise = torch.randn(label_dim, G_input_dim)\n",
    "fixed_noise = temp_noise\n",
    "fixed_c = torch.zeros(label_dim, 1)\n",
    "for i in range(9):\n",
    "    fixed_noise = torch.cat([fixed_noise, temp_noise], 0)\n",
    "    temp = torch.ones(label_dim, 1) + i\n",
    "    fixed_c = torch.cat([fixed_c, temp], 0)\n",
    "\n",
    "fixed_noise = fixed_noise.view(-1, G_input_dim, 1, 1)\n",
    "fixed_label = torch.zeros(G_input_dim, label_dim)\n",
    "fixed_label.scatter_(1, fixed_c.type(torch.LongTensor), 1)\n",
    "fixed_label = fixed_label.view(-1, label_dim, 1, 1)\n",
    "\n",
    "# label preprocess\n",
    "onehot = torch.zeros(label_dim, label_dim)\n",
    "onehot = onehot.scatter_(1, torch.LongTensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]).view(label_dim, 1), 1).view(label_dim, label_dim, 1, 1)\n",
    "fill = torch.zeros([label_dim, label_dim, img_size, img_size])\n",
    "for i in range(label_dim):\n",
    "    fill[i, i, :, :] = 1\n",
    "\n",
    "D_avg_losses = []\n",
    "G_avg_losses = []\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "for epoch in range(1, num_epochs +1):\n",
    "    D_losses = []\n",
    "    G_losses = []\n",
    "    epoch_start_time = time.time()\n",
    "    generator.train()\n",
    "    if epoch == 5 or epoch == 10:\n",
    "        opt_Gen.param_groups[0]['lr'] /=2\n",
    "        opt_Disc.param_groups[0]['lr'] /=2\n",
    "        \n",
    "    \n",
    "    for i, (real_images,real_labels) in enumerate(data_loader):\n",
    "        minibatch = real_images.size()[0]\n",
    "        real_images = Variable(real_images.cuda())\n",
    "        \n",
    "        #labels\n",
    "        y_real = Variable(torch.ones(minibatch).cuda())\n",
    "        y_fake = Variable(torch.zeros(minibatch).cuda())\n",
    "        c_fill = Variable(fill[real_labels].cuda())\n",
    "        z_ = torch.randn(minibatch, G_input_dim).view(-1, G_input_dim, 1, 1)\n",
    "        z_ = Variable(z_.cuda())\n",
    "        ## Train Discriminator             \n",
    "        # first with real data\n",
    "        \n",
    "        D_real_decision = discriminator(real_images, c_fill).squeeze()\n",
    "        D_real_loss = criterion(D_real_decision, y_real)\n",
    "        \n",
    "        # Then with fake data\n",
    "        \n",
    "        c_ = (torch.rand(minibatch, 1) * label_dim).type(torch.LongTensor).squeeze()\n",
    "        c_onehot_ = Variable(onehot[c_].cuda())\n",
    "        c_fill_ = Variable(fill[c_].cuda())\n",
    "        \n",
    "        generator_image = generator(z_, c_onehot_) \n",
    "        D_fake_decision = discriminator(generator_image,c_fill_).squeeze()\n",
    "        D_fake_loss = criterion(D_fake_decision, y_fake)\n",
    "        \n",
    "        # Optimization\n",
    "        discriminator.zero_grad()\n",
    "        D_loss = D_fake_loss + D_real_loss\n",
    "        D_loss.backward()\n",
    "        opt_Disc.step()\n",
    "        \n",
    "        # Train Generator\n",
    "        \n",
    "        generator_image = generator(z_, c_onehot_)\n",
    "        c_fill = Variable(fill[c_].cuda())\n",
    "        D_fake_decision = discriminator(generator_image,c_fill).squeeze()\n",
    "        G_loss = criterion(D_fake_decision, y_real)\n",
    "        \n",
    "        # Optimization\n",
    "        generator.zero_grad()\n",
    "        G_loss.backward()\n",
    "        opt_Gen.step()\n",
    "        \n",
    "        D_losses.append(D_loss.data[0])\n",
    "        G_losses.append(G_loss.data[0])\n",
    "        \n",
    "        print('Epoch [%d/%d], Step [%d/%d], D_loss: %.4f, G_loss: %.4f'\n",
    "            % (epoch+1, num_epochs, i+1, len(data_loader), D_loss.data[0], G_loss.data[0]))\n",
    "        \n",
    "    torch.save(generator.state_dict(), '/home/abdullah/Documents/ai/models/Own implementation of Gan/Conditional Gan/model_weights/generator_param.pkl')\n",
    "    torch.save(discriminator.state_dict(), '/home/abdullah/Documents/ai/models/Own implementation of Gan/Conditional Gan/model_weights/discriminator_param.pkl')\n",
    "    \n",
    "    D_avg_loss = torch.mean(torch.FloatTensor(D_losses))\n",
    "    G_avg_loss = torch.mean(torch.FloatTensor(G_losses))\n",
    "\n",
    "    # avg loss values for plot\n",
    "    D_avg_losses.append(D_avg_loss)\n",
    "    G_avg_losses.append(G_avg_loss)\n",
    "\n",
    "    plot_loss(D_avg_losses, G_avg_losses, epoch, save=True)\n",
    "\n",
    "    # Show result for fixed noise\n",
    "    plot_result(generator, fixed_noise, fixed_label, epoch, save=True)\n",
    "    \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make gif\n",
    "loss_plots = []\n",
    "gen_image_plots = []\n",
    "for epoch in range(2,24):\n",
    "    # plot for generating gif\n",
    "    save_fn1 = 'Model resultsMNIST_cDCGAN_losses_epoch_{:d}'.format(epoch + 1) + '.png'\n",
    "    loss_plots.append(imageio.imread(save_fn1))\n",
    "\n",
    "    save_fn2 = 'Model resultsMNIST_cDCGAN_epoch_{:d}'.format(epoch + 1) + '.png'\n",
    "    gen_image_plots.append(imageio.imread(save_fn2))\n",
    "\n",
    "imageio.mimsave('MNIST_cDCGAN_losses_epochs_{:d}'.format(num_epochs) + '.gif', loss_plots, fps=5)\n",
    "imageio.mimsave('MNIST_cDCGAN_epochs_{:d}'.format(num_epochs) + '.gif', gen_image_plots, fps=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}