DCGANv2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "j4KxRoZqp5ir"
   },
   "source": [
    "### Import statements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     },
     "base_uri": "https://localhost:8080/",
     "height": 224
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 3720,
     "status": "ok",
     "timestamp": 1533387411668,
     "user": {
      "displayName": "Kristine Faith Roque",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s128",
      "userId": "101942258231543881539"
     },
     "user_tz": -540
    },
    "id": "dFyc-I899Ylc",
    "outputId": "4030756f-cee0-4b4b-ebfd-d501428e82ad"
   },
   "outputs": [],
   "source": [
    "!pip3 install torch torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "fufN73BNp5is"
   },
   "outputs": [],
   "source": [
    "from __future__ import division\n",
    "from IPython.display import clear_output\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.autograd import Variable\n",
    "from torch.optim import Adam\n",
    "from torchvision import transforms, datasets\n",
    "import torchvision.utils as vutils\n",
    "\n",
    "from google.colab import files\n",
    "import os\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     },
     "base_uri": "https://localhost:8080/",
     "height": 190
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 4229,
     "status": "ok",
     "timestamp": 1533387417597,
     "user": {
      "displayName": "Kristine Faith Roque",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s128",
      "userId": "101942258231543881539"
     },
     "user_tz": -540
    },
    "id": "5t3gblD2_Led",
    "outputId": "10b58c33-956a-48d8-ef8e-0292f20c2233"
   },
   "outputs": [],
   "source": [
    "!pip install Pillow==4.0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "TR_FeRH_Iqq-"
   },
   "outputs": [],
   "source": [
    "!pip install -U -q PyDrive\n",
    "\n",
    "from pydrive.auth import GoogleAuth\n",
    "from pydrive.drive import GoogleDrive\n",
    "from google.colab import auth\n",
    "from oauth2client.client import GoogleCredentials\n",
    "\n",
    "# 1. Authenticate and create the PyDrive client.\n",
    "auth.authenticate_user()\n",
    "gauth = GoogleAuth()\n",
    "gauth.credentials = GoogleCredentials.get_application_default()\n",
    "drive = GoogleDrive(gauth)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "kq61KCJ9IsV-"
   },
   "outputs": [],
   "source": [
    "#Create the folder\n",
    "# If no parent ID is provided this will automatically go to the root or My Drive 'directory'\n",
    "results = drive.CreateFile({'title': 'results',\n",
    "                                    \"mimeType\": \"application/vnd.google-apps.folder\"})\n",
    "# Upload the file to your drive\n",
    "results.Upload()\n",
    "# Grab the ID of the folder we just created\n",
    "result_id = results['id']\n",
    "\n",
    "# Create a sub-directory\n",
    "# Make sure to assign it the proper parent ID\n",
    "real = drive.CreateFile({'title': 'real', \n",
    "                         'parents':[{'id':result_id}],\n",
    "                        \"mimeType\": \"application/vnd.google-apps.folder\"} )\n",
    "real.Upload()\n",
    "real_id = real['id']\n",
    "\n",
    "fake = drive.CreateFile({'title': 'fake', \n",
    "                         'parents':[{'id':result_id}],\n",
    "                        \"mimeType\": \"application/vnd.google-apps.folder\"} )\n",
    "fake.Upload()\n",
    "fake_id = fake['id']\n",
    "\n",
    "loss_logs = drive.CreateFile({'title': 'loss_logs', \n",
    "                         'parents':[{'id':result_id}],\n",
    "                        \"mimeType\": \"application/vnd.google-apps.folder\"} )\n",
    "loss_logs.Upload()\n",
    "loss_id = loss_logs['id']\n",
    "\n",
    "models = drive.CreateFile({'title': 'models', \n",
    "                         'parents':[{'id':result_id}],\n",
    "                        \"mimeType\": \"application/vnd.google-apps.folder\"} )\n",
    "\n",
    "models.Upload()\n",
    "models_id = models['id']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KAJosSOIp5iv"
   },
   "source": [
    "### Define discriminator and generator networks\n",
    "\n",
    "Discriminator --> just a CNN \n",
    "\n",
    "Generator --> basically a reverse CNN\n",
    "\n",
    "The DCGAN network in the paper by Radford et al is implemented. \n",
    "Note:\n",
    "\n",
    "_''The first layer of the GAN, which takes a uniform noise distribution Z as input, could be called fully connected as it is just a matrix multiplication, but the result is reshaped into a 4-dimensional tensor and used as the start of the convolution stack. For the discriminator, the last convolution layer is flattened and then fed into a single sigmoid output. See Fig. 1 for a visualization of an example model architecture_''.\n",
    "\n",
    "<img src=\"figures/Generator_network.png\" width=\"800\">\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "rUhpRT40p5iv"
   },
   "outputs": [],
   "source": [
    "class Generator(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    This generator network is based on the original DCGAN paper by Radford et al.\n",
    "    The generator takes as input a 100-dimensional noise vector (z) and maps it to the data space \n",
    "    (which in this case is the image space) via a series of 'deconvolution'* blocks.\n",
    "    Hence, from the input random noise, the generator outputs an image \n",
    "    \n",
    "    *Note: These 'deconvolution' blocks do not perform the mathematical deconvolution operator, \n",
    "    it is just a reverse operation for the convolution operation used in deep learning models.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self):\n",
    "        super(Generator, self).__init__()\n",
    "        self.nz = 100\n",
    "        self.linear = torch.nn.Linear(self.nz, 1024*4*4)\n",
    "        \n",
    "        #first 'deconvolution' block\n",
    "        self.Conv1 = nn.Sequential(\n",
    "            nn.ConvTranspose2d(in_channels=1024, out_channels=512, kernel_size=4,stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(inplace=True))\n",
    "        \n",
    "        #second 'deconvolution' block\n",
    "        self.Conv2 = nn.Sequential(\n",
    "            nn.ConvTranspose2d(in_channels=512, out_channels=256, kernel_size=4, stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "        \n",
    "        #third 'deconvolution' block\n",
    "        self.Conv3 = nn.Sequential(\n",
    "            nn.ConvTranspose2d(in_channels=256, out_channels=128, kernel_size=4, stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "        \n",
    "        #fourth 'deconvolution' block\n",
    "        self.Conv4 = nn.Sequential(\n",
    "            nn.ConvTranspose2d(in_channels=128, out_channels=3, kernel_size=4, stride=2, padding=1, bias=False)\n",
    "        )\n",
    "        self.out = torch.nn.Tanh()\n",
    "\n",
    "    def forward(self, z):\n",
    "        \"\"\"\n",
    "        Perform forward calculation for generator output, given random noise input z\n",
    "        \"\"\"\n",
    "        # Project and reshape\n",
    "        X = self.linear(z)\n",
    "        X = X.view(X.shape[0], 1024, 4, 4)\n",
    "        # conv blocks\n",
    "        X = self.Conv1(X)\n",
    "        X = self.Conv2(X)\n",
    "        X = self.Conv3(X)\n",
    "        X = self.Conv4(X)\n",
    "        # tanh activation\n",
    "        return self.out(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "RHoFNo_Up5iy"
   },
   "outputs": [],
   "source": [
    "class Discriminator(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    This discriminator network is based on the original DCGAN paper by Radford et al.\n",
    "    The discriminator is a CNN which takes as input a 3-channel image data (i.e. RGB image)\n",
    "    and outputs a probability,p(real), that the image is from the real dataset.\n",
    "    \"\"\"\n",
    "    def __init__(self):\n",
    "        super(Discriminator, self).__init__()\n",
    "        #Conv block 1\n",
    "        self.Conv1 = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=3, out_channels=128, kernel_size=4, stride=2, padding=1, bias=False),\n",
    "            nn.LeakyReLU(0.2, inplace=True)\n",
    "        )\n",
    "        #Conv block 2\n",
    "        self.Conv2 = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.LeakyReLU(0.2, inplace=True)\n",
    "        )\n",
    "        #Conv block 3\n",
    "        self.Conv3 = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=4, stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.LeakyReLU(0.2, inplace=True)\n",
    "        )\n",
    "        #Conv block 4\n",
    "        self.Conv4 = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=512, out_channels=1024, kernel_size=4,stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(1024),\n",
    "            nn.LeakyReLU(0.2, inplace=True)\n",
    "        )\n",
    "        self.Out = nn.Sequential(\n",
    "            nn.Linear(1024*4*4, 1),\n",
    "            nn.Sigmoid(),\n",
    "        )\n",
    "\n",
    "    def forward(self, I):\n",
    "        # Convolutional layers\n",
    "        X = self.Conv1(I)\n",
    "        X = self.Conv2(X)\n",
    "        X = self.Conv3(X)\n",
    "        X = self.Conv4(X)\n",
    "        # reshape and apply sigmoid activation\n",
    "        X = X.view(-1, 1024*4*4)\n",
    "        X = self.Out(X)\n",
    "        return X"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MJBETlrDp5i0"
   },
   "source": [
    "### Training of DCGAN \n",
    "The discriminator and generator are like two agents playing a minimax game with value function V(G, D):\n",
    "\n",
    "$$\\min_{G}\\max_{D}V(G, D)  = \\textbf{E}_{x\\sim p_{data}(x)}[\\log{D(x)}] + \\textbf{E}_{z\\sim p_{z}(z)}[\\log{(1-D(G(z)))}]$$\n",
    "\n",
    "That is, the generator is trying to fool the discriminator by maximizing the probability that its output is recognized as real, which is mathematically equivalent to minimizing $log(1-D(G(z))$. Meanwhile the discriminator tries to increase its accuracy in distinguishing between real and fake data by minimizing this same quantity, which is equivalent to maximizing the above value function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "Wb1DZWeqp5i1"
   },
   "outputs": [],
   "source": [
    "def train(discriminator, generator, d_optimizer, g_optimizer, real_data, fake_data, real_target, fake_target, loss):\n",
    "    #===============================================================\n",
    "    #1. Train Discriminator by maximizing log(D(x)) + log(1 - D(G(z)))\n",
    "    #===============================================================\n",
    "    # Reset gradients\n",
    "    d_optimizer.zero_grad()\n",
    "    \n",
    "    # 1.1 Train on Real Data\n",
    "    d_prediction_real = discriminator(real_data)\n",
    "    # Calculate error and backpropagate\n",
    "    d_error_real = loss(d_prediction_real, real_target)\n",
    "    d_error_real.backward()\n",
    "\n",
    "    # 1.2 Train on Fake Data\n",
    "    d_prediction_fake = discriminator(fake_data.detach())\n",
    "    # Calculate error and backpropagate\n",
    "    d_error_fake = loss(d_prediction_fake, fake_target)\n",
    "    d_error_fake.backward()\n",
    "    \n",
    "    # 1.3 Update weights with gradients\n",
    "    d_optimizer.step()\n",
    "    \n",
    "    #===============================================================\n",
    "    #2. Train Generator by maximizing log(D(G(z)))\n",
    "    #===============================================================\n",
    "    # Reset gradients\n",
    "    g_optimizer.zero_grad()\n",
    "    # Sample noise and generate fake data\n",
    "    g_prediction = discriminator(fake_data)\n",
    "    # Calculate error and backpropagate\n",
    "    g_error = loss(g_prediction, real_target)\n",
    "    g_error.backward()\n",
    "    # Update weights with gradients\n",
    "    g_optimizer.step()\n",
    "\n",
    "    # Return error\n",
    "    return d_error_real + d_error_fake, d_prediction_real, d_prediction_fake, g_error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "Z-LGdBC_p5i3"
   },
   "outputs": [],
   "source": [
    "def train_discriminator(optimizer, real_data, fake_data):\n",
    "    # Reset gradients\n",
    "    optimizer.zero_grad()\n",
    "    \n",
    "    # 1.1 Train on Real Data\n",
    "    prediction_real = netD(real_data)\n",
    "    # Calculate error and backpropagate\n",
    "    error_real = loss(prediction_real, real_data_target(real_data.size(0)))\n",
    "    error_real.backward()\n",
    "\n",
    "    # 1.2 Train on Fake Data\n",
    "    prediction_fake = netD(fake_data)\n",
    "    # Calculate error and backpropagate\n",
    "    error_fake = loss(prediction_fake, fake_data_target(real_data.size(0)))\n",
    "    error_fake.backward()\n",
    "    \n",
    "    # 1.3 Update weights with gradients\n",
    "    optimizer.step()\n",
    "    \n",
    "    # Return error\n",
    "    return error_real + error_fake, prediction_real, prediction_fake\n",
    "\n",
    "def train_generator(optimizer, fake_data):\n",
    "    # 2. Train Generator\n",
    "    # Reset gradients\n",
    "    optimizer.zero_grad()\n",
    "    # Sample noise and generate fake data\n",
    "    prediction = netD(fake_data)\n",
    "    # Calculate error and backpropagate\n",
    "    error = loss(prediction, real_data_target(prediction.size(0)))\n",
    "    error.backward()\n",
    "    # Update weights with gradients\n",
    "    optimizer.step()\n",
    "    # Return error\n",
    "    return error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "gKItb5Rgp5i4"
   },
   "outputs": [],
   "source": [
    "def noise(s):\n",
    "    \"\"\"\n",
    "    Generate s-dimensional noise vector from random normal distribution with mean zero and std one\n",
    "    \"\"\"\n",
    "    z = Variable(torch.randn(s, 100))\n",
    "    if torch.cuda.is_available(): \n",
    "        return z.cuda()\n",
    "    return z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "Dv6SEEb_p5i7"
   },
   "outputs": [],
   "source": [
    "# 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)\n",
    "        \n",
    "def weights_init(m):\n",
    "    classname = m.__class__.__name__\n",
    "    if classname.find('Conv') != -1 or classname.find('BatchNorm') != -1:\n",
    "        m.weight.data.normal_(0.00, 0.02)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "4nVol2c2p5i9"
   },
   "outputs": [],
   "source": [
    "#load cifar 10 data\n",
    "def cifar_data():\n",
    "    compose = transforms.Compose(\n",
    "        [\n",
    "            transforms.Resize(64),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((.5, .5, .5), (.5, .5, .5))\n",
    "        ])\n",
    "    out_dir = './data/'\n",
    "    return datasets.CIFAR10(root=out_dir, train=True, transform=compose, download=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     },
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 2440,
     "status": "ok",
     "timestamp": 1533387429934,
     "user": {
      "displayName": "Kristine Faith Roque",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s128",
      "userId": "101942258231543881539"
     },
     "user_tz": -540
    },
    "id": "s9xQxOSFp5i_",
    "outputId": "6c566d9c-540b-48b5-8235-7ee93fa0e5f1"
   },
   "outputs": [],
   "source": [
    "data = cifar_data()\n",
    "batch_size = 100\n",
    "data_loader = torch.utils.data.DataLoader(data, batch_size=batch_size, shuffle=True)\n",
    "num_batches = len(data_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "nCNTs9tip5jD"
   },
   "outputs": [],
   "source": [
    "# create network instances\n",
    "netG = Generator()\n",
    "netD = Discriminator()\n",
    "\n",
    "#initialize weights\n",
    "netD.apply(weights_init)\n",
    "netG.apply(weights_init)\n",
    "\n",
    "#use cuda if available\n",
    "if torch.cuda.is_available():\n",
    "    netG.cuda()\n",
    "    netD.cuda()\n",
    "\n",
    "#generate samples for testing\n",
    "num_test_samples = 16\n",
    "test_noise = noise(num_test_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "1fvOkebHp5jF"
   },
   "outputs": [],
   "source": [
    "# learning rate\n",
    "lr = 0.0002\n",
    "\n",
    "# Number of training epochs\n",
    "num_epochs = 200\n",
    "\n",
    "#setup optimizers\n",
    "optD = Adam(netD.parameters(), lr=lr, betas=(0.5, 0.999))\n",
    "optG = Adam(netG.parameters(), lr=lr, betas=(0.5, 0.999))\n",
    "\n",
    "#test noise\n",
    "test_noise = noise(20)\n",
    "\n",
    "# loss function\n",
    "loss = nn.BCELoss()\n",
    "\n",
    "real_label = 1\n",
    "fake_label = 0\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "Mh4mR3l0p5jH"
   },
   "outputs": [],
   "source": [
    "def real_data_target(size):\n",
    "    '''\n",
    "    Tensor containing ones, with shape = size\n",
    "    '''\n",
    "    data = Variable(torch.ones(size, 1))\n",
    "    if torch.cuda.is_available(): return data.cuda()\n",
    "    return data\n",
    "\n",
    "def fake_data_target(size):\n",
    "    '''\n",
    "    Tensor containing zeros, with shape = size\n",
    "    '''\n",
    "    data = Variable(torch.zeros(size, 1))\n",
    "    if torch.cuda.is_available(): return data.cuda()\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     }
    },
    "colab_type": "code",
    "id": "Nc4niU0kDUJ0"
   },
   "outputs": [],
   "source": [
    "def save_file(fname, file, id_):\n",
    "    # 2. Create & upload a file text file\n",
    "    uploaded = drive.CreateFile({'title': fname, 'parents':[{'id': id_}]})\n",
    "    uploaded.SetContentFile(file)\n",
    "    uploaded.Upload()\n",
    "    print('File %s uploaded'%fname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     },
     "base_uri": "https://localhost:8080/",
     "height": 119
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 8871,
     "status": "ok",
     "timestamp": 1533387446168,
     "user": {
      "displayName": "Kristine Faith Roque",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s128",
      "userId": "101942258231543881539"
     },
     "user_tz": -540
    },
    "id": "yPByaZJ9qchA",
    "outputId": "d7be3c6e-1ea2-41e2-90b3-e27acc2aeed8"
   },
   "outputs": [],
   "source": [
    "# memory footprint support libraries/code\n",
    "!ln -sf /opt/bin/nvidia-smi /usr/bin/nvidia-smi\n",
    "!pip install gputil\n",
    "!pip install psutil\n",
    "!pip install humanize\n",
    "import psutil\n",
    "import humanize\n",
    "import os\n",
    "import GPUtil as GPU\n",
    "GPUs = GPU.getGPUs()\n",
    "# XXX: only one GPU on Colab and isn’t guaranteed\n",
    "gpu = GPUs[0]\n",
    "def printm():\n",
    "    process = psutil.Process(os.getpid())\n",
    "    print(\"Gen RAM Free: \" + humanize.naturalsize( psutil.virtual_memory().available ), \" | Proc size: \" + humanize.naturalsize( process.memory_info().rss))\n",
    "    print(\"GPU RAM Free: {0:.0f}MB | Used: {1:.0f}MB | Util {2:3.0f}% | Total {3:.0f}MB\".format(gpu.memoryFree, gpu.memoryUsed, gpu.memoryUtil*100, gpu.memoryTotal))\n",
    "    printm()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     },
     "base_uri": "https://localhost:8080/",
     "height": 85
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 5501,
     "status": "ok",
     "timestamp": 1533387451699,
     "user": {
      "displayName": "Kristine Faith Roque",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s128",
      "userId": "101942258231543881539"
     },
     "user_tz": -540
    },
    "id": "rr2baaybqiz6",
    "outputId": "1c4c4be4-fc7f-4dbf-bf1d-b00d722bbbd1"
   },
   "outputs": [],
   "source": [
    "!mkdir results\n",
    "!mkdir results/models\n",
    "!mkdir results/loss_logs\n",
    "!mkdir results/fake"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "autoexec": {
      "startup": false,
      "wait_interval": 0
     },
     "base_uri": "https://localhost:8080/",
     "height": 810
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 43499,
     "status": "error",
     "timestamp": 1533390525097,
     "user": {
      "displayName": "Kristine Faith Roque",
      "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s128",
      "userId": "101942258231543881539"
     },
     "user_tz": -540
    },
    "id": "RQDp_qHop5jK",
    "outputId": "fb41a628-eb2e-45b3-ff6e-ce541d871d08"
   },
   "outputs": [],
   "source": [
    "for epoch in range(num_epochs):\n",
    "    dloss_log = []\n",
    "    gloss_log = []\n",
    "    for i, data in enumerate(data_loader, 0):\n",
    "        real_data = Variable(data[0])        \n",
    "        # 1. Train Discriminator\n",
    "        if torch.cuda.is_available(): real_data = real_data.cuda()\n",
    "        # Generate fake data\n",
    "        fake_data = netG(noise(real_data.size(0))).detach()\n",
    "        # Train D\n",
    "        d_error, d_pred_real, d_pred_fake = train_discriminator(optD, \n",
    "                                                                real_data, fake_data)\n",
    "\n",
    "        # 2. Train Generator\n",
    "        # Generate fake data\n",
    "        fake_data = netG(noise(real_data.size(0)))\n",
    "        # Train G\n",
    "        g_error = train_generator(optG, fake_data)\n",
    "\n",
    "\n",
    "        print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G:%.4f'%\n",
    "              (epoch, num_epochs, i, len(data_loader), d_error.item(), g_error.item()))\n",
    "\n",
    "        dloss_log.append(d_error.item())\n",
    "        gloss_log.append(g_error.item())\n",
    "        \n",
    "        \n",
    "        if i%int(batch_size) == 0: \n",
    "            fname_real = 'real_samples_epoch_%03d_%03d.png'% (epoch, i)\n",
    "            real_path = 'results/' + fname_real\n",
    "            \n",
    "            vutils.save_image(real_data,real_path, normalize = True)\n",
    "            fake = netG(test_noise)\n",
    "            fname_fake = 'fake_samples2_epoch_%03d_%03d.png'% (epoch, i)\n",
    "            fake_path = 'results/fake/' + fname_fake\n",
    "            \n",
    "            vutils.save_image(fake.detach(), fake_path, normalize=True)\n",
    "            save_file(fname_fake, fake_path, fake_id)\n",
    "            save_file(fname_real, real_path, real_id)\n",
    "\n",
    "            if i != 0:\n",
    "                clear_output()\n",
    "\n",
    "            \n",
    "            #log loss per epoch of training\n",
    "            with open(\"results/loss_logs/dloss2_%03d.pickle\"%(epoch), \"wb\") as output_file:\n",
    "                pickle.dump(dloss_log, output_file)\n",
    "    \n",
    "            with open(\"results/loss_logs/gloss2_%03d.pickle\"%(epoch), \"wb\") as output_file:\n",
    "                pickle.dump(gloss_log, output_file)\n",
    " \n",
    "            torch.save(netG.state_dict(), 'results/models/netG2_epoch_%d.pth' % (epoch))\n",
    "            torch.save(netD.state_dict(), 'results/models/netD2_epoch_%d.pth' % (epoch))\n",
    "    \n",
    "    #log loss per epoch of training\n",
    "    dlogs_fname = 'dloss2_%03d.pickle'%epoch\n",
    "    glogs_fname = 'gloss2_%03d.pickle'%epoch\n",
    "    \n",
    "    with open(\"results/loss_logs/\"+dlogs_fname, \"wb\") as output_file:\n",
    "        pickle.dump(dloss_log, output_file)\n",
    "    \n",
    "    with open('results/loss_logs/'+glogs_fname, \"wb\") as output_file:\n",
    "        pickle.dump(gloss_log, output_file)\n",
    "        \n",
    "    save_file(dlogs_fname, \"results/loss_logs/\"+dlogs_fname, loss_id)\n",
    "    save_file(glogs_fname, \"results/loss_logs/\"+glogs_fname, loss_id)\n",
    "    \n",
    "    #save model state per epoch\n",
    "    Dstate = 'netD2_epoch_%d.pth' % (epoch)\n",
    "    Gstate = 'netG2_epoch_%d.pth' % (epoch)\n",
    "    torch.save(netG.state_dict(), 'results/models/'+Gstate)\n",
    "    torch.save(netD.state_dict(), 'results/models/'+Dstate)\n",
    "\n",
    "    save_file(Dstate, \"results/models/\"+Dstate, models_id)\n",
    "    save_file(Gstate, \"results/models/\"+Gstate, models_id)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "default_view": {},
   "name": "DCGANv2.ipynb",
   "provenance": [],
   "toc_visible": true,
   "version": "0.3.2",
   "views": {}
  },
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}