Cat DCGAN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cat DCGAN 🐱‍💻"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Outputs from our DCGAN: "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"assets/output.png\" alt=\"CatDCGAN Output\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p> Cat DCGAN is a Deep Convolutional Generative Adversarial Network (DCGAN) <b>that generates pictures of cats</b> </p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is an open source project"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p> If you have any questions, feel free to ask me: </p>\n",
    "<p> 📧: <a href=\"mailto:hello@simoninithomas.com\">hello@simoninithomas.com</a>  </p>\n",
    "<p> Github: https://github.com/simoninithomas/CatDCGAN </p>\n",
    "<p> 🌐 : https://www.simoninithomas.com </p>\n",
    "<p> Twitter: <a href=\"https://twitter.com/ThomasSimonini\">@ThomasSimonini</a> </p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Important note 🤔\n",
    "<b> You can't run this on your computer </b> (except if you have GPUs or wait 10 years 😅), personally I train this DCGAN for 20 hours with Microsoft Azure and their Deep Learning Virtual Machine (they offer 170$)\n",
    "https://azuremarketplace.microsoft.com/en-us/marketplace/apps/microsoft-ads.dsvm-deep-learning\n",
    "<br>\n",
    "⚠️ I don't have any business relations with them. I just loved their excellent customer service.\n",
    "\n",
    "If you have some troubles to use follow the explainations of this excellent article here (without last the part fast.ai): https://medium.com/@manikantayadunanda/setting-up-deeplearning-machine-and-fast-ai-on-azure-a22eb6bd6429"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist 📝\n",
    "- Download the dataset here: https://www.kaggle.com/crawford/cat-dataset\n",
    "- Type `sh start.sh` it will handle extract, remove outliers, normalization and face centering\n",
    "- Change `do_preprocess = True` ⚠️⚠️⚠️ important!\n",
    "\n",
    "### If you want to train from scratch\n",
    "- Change `from_checkpoint = False`\n",
    "\n",
    "### If you want to train from last model saved (you save 20 hours of training 🎉)\n",
    "- Change `from_checkpoint = True`\n",
    "\n",
    "\n",
    "## Acknowledgement 👏\n",
    "This project was made possible thanks to:\n",
    "- Udacity Face Generator Project \n",
    "- The start.sh and preprocess part (modified) made by Alexia Jolicoeur-Martineau https://ajolicoeur.wordpress.com/cats/\n",
    "- Siraj's Raval PokeGan https://github.com/llSourcell/Pokemon_GAN\n",
    "- The choice of learning rate by Alexia Jolicoeur-Martineau https://ajolicoeur.wordpress.com/cats/\n",
    "\n",
    "\n",
    "<br>\n",
    "<br>\n",
    "<img src=\"assets/training2.gif\" alt=\"Training DCGAN\" />"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "\n",
    "import helper\n",
    "from glob import glob\n",
    "import pickle as pkl\n",
    "import scipy.misc\n",
    "\n",
    "import time\n",
    "\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "do_preprocess = False\n",
    "from_checkpoint = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Data 📚\n",
    "### Resize images to 128x128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = './cats_bigger_than_128x128' # Data\n",
    "data_resized_dir = \"./resized_data\"# Resized data\n",
    "\n",
    "if do_preprocess == True:\n",
    "    os.mkdir(data_resized_dir)\n",
    "\n",
    "    for each in os.listdir(data_dir):\n",
    "        image = cv2.imread(os.path.join(data_dir, each))\n",
    "        image = cv2.resize(image, (128, 128))\n",
    "        cv2.imwrite(os.path.join(data_resized_dir, each), image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This part was taken from Udacity Face generator project\n",
    "def get_image(image_path, width, height, mode):\n",
    "    \"\"\"\n",
    "    Read image from image_path\n",
    "    :param image_path: Path of image\n",
    "    :param width: Width of image\n",
    "    :param height: Height of image\n",
    "    :param mode: Mode of image\n",
    "    :return: Image data\n",
    "    \"\"\"\n",
    "    image = Image.open(image_path)\n",
    "\n",
    "    return np.array(image.convert(mode))\n",
    "\n",
    "def get_batch(image_files, width, height, mode):\n",
    "    data_batch = np.array(\n",
    "        [get_image(sample_file, width, height, mode) for sample_file in image_files]).astype(np.float32)\n",
    "\n",
    "    # Make sure the images are in 4 dimensions\n",
    "    if len(data_batch.shape) < 4:\n",
    "        data_batch = data_batch.reshape(data_batch.shape + (1,))\n",
    "\n",
    "    return data_batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Explore the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_n_images = 25\n",
    "mnist_images = helper.get_batch(glob(os.path.join(data_resized_dir, '*.jpg'))[:show_n_images], 64, 64, 'RGB')\n",
    "plt.imshow(helper.images_square_grid(mnist_images, 'RGB'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The DCGAN 🤖"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this part we're going to implement the DCGAN.\n",
    "Our Architecture:<br><br>\n",
    "<img src=\"assets/GDSchema.png\" alt=\"Cat DCGAN Architecture\"/>\n",
    "\n",
    "Cat Icon made by <a href=\"https://www.flaticon.com/authors/vectors-market\">  Vector Market </a> from www.flaticon.com \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check the version of Tensorflow and access to GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Taken from Udacity face generator project\n",
    "from distutils.version import LooseVersion\n",
    "import warnings\n",
    "import tensorflow as tf\n",
    "\n",
    "# Check TensorFlow Version\n",
    "assert LooseVersion(tf.__version__) >= LooseVersion('1.0'), 'Please use TensorFlow version 1.0 or newer.  You are using {}'.format(tf.__version__)\n",
    "print('TensorFlow Version: {}'.format(tf.__version__))\n",
    "\n",
    "# Check for a GPU\n",
    "if not tf.test.gpu_device_name():\n",
    "    warnings.warn('No GPU found. Please use a GPU to train your neural network.')\n",
    "else:\n",
    "    print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Input\n",
    "Create TF placeholders for the Neural Network:\n",
    "- Real input images placeholder `real_dim`.\n",
    "- Z input placeholder `z_dim`.\n",
    "- Learning rate G placeholder.\n",
    "- Learning rate D placeholder.\n",
    "<br><br>\n",
    "Return the placeholders in a tuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_inputs(real_dim, z_dim):\n",
    "    \"\"\"\n",
    "    Create the model inputs\n",
    "    :param real_dim: tuple containing width, height and channels\n",
    "    :param z_dim: The dimension of Z\n",
    "    :return: Tuple of (tensor of real input images, tensor of z data, learning rate G, learning rate D)\n",
    "    \"\"\"\n",
    "    inputs_real = tf.placeholder(tf.float32, (None, *real_dim), name='inputs_real')\n",
    "    inputs_z = tf.placeholder(tf.float32, (None, z_dim), name=\"input_z\")\n",
    "    learning_rate_G = tf.placeholder(tf.float32, name=\"learning_rate_G\")\n",
    "    learning_rate_D = tf.placeholder(tf.float32, name=\"learning_rate_D\")\n",
    "    \n",
    "    return inputs_real, inputs_z, learning_rate_G, learning_rate_D\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generator Network\n",
    "\n",
    "<img src=\"assets/generator.png\" alt=\"Generator\"/>\n",
    "\n",
    "\n",
    "#### Variable Scope\n",
    "Use tf.variable_scope <b> for 2 reasons </b>:\n",
    "<ul>\n",
    "    <li> Make sure all varaibles names start with generator / discriminator (will help out later when training the separate networks </li>\n",
    "    <li> Also want <b> to reuse these networks with different inputs </b></li>\n",
    "        <ul>\n",
    "            <li> For the generator: we're going to train it but also <b>sample from it as we're training after training </b> </li>\n",
    "            <li> For the discriminator: need to share variables between the fake and real input images </li>\n",
    "    </ul>\n",
    "</ul>\n",
    "<p> So we can use the reuse keyword to <b> tell TensorFlow to reuse the var instead of createing new one if we build the graph again</b></p>\n",
    "\n",
    "#### Leaky ReLU\n",
    "Avoid gradient vanishing\n",
    "\n",
    "#### Tanh Output\n",
    "Generator has been found to perform the best <b> with tanh for the generator output </b>\n",
    "<br>\n",
    "\n",
    "- Leaky ReLU in all layers except for the last tanh layer\n",
    "- Normalization on all the transposed convnets except the last one\n",
    "\n",
    "<br>\n",
    "<b>Transposed convnets --> normalization --> leaky ReLU</b>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generator(z, output_channel_dim, is_train=True):\n",
    "    ''' Build the generator network.\n",
    "    \n",
    "        Arguments\n",
    "        ---------\n",
    "        z : Input tensor for the generator\n",
    "        output_channel_dim : Shape of the generator output\n",
    "        n_units : Number of units in hidden layer\n",
    "        reuse : Reuse the variables with tf.variable_scope\n",
    "        alpha : leak parameter for leaky ReLU\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        out: \n",
    "    '''\n",
    "    with tf.variable_scope(\"generator\", reuse= not is_train):\n",
    "        \n",
    "        # First FC layer --> 8x8x1024\n",
    "        fc1 = tf.layers.dense(z, 8*8*1024)\n",
    "        \n",
    "        # Reshape it\n",
    "        fc1 = tf.reshape(fc1, (-1, 8, 8, 1024))\n",
    "        \n",
    "        # Leaky ReLU\n",
    "        fc1 = tf.nn.leaky_relu(fc1, alpha=alpha)\n",
    "\n",
    "        \n",
    "        # Transposed conv 1 --> BatchNorm --> LeakyReLU\n",
    "        # 8x8x1024 --> 16x16x512\n",
    "        trans_conv1 = tf.layers.conv2d_transpose(inputs = fc1,\n",
    "                                  filters = 512,\n",
    "                                  kernel_size = [5,5],\n",
    "                                  strides = [2,2],\n",
    "                                  padding = \"SAME\",\n",
    "                                kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
    "                                name=\"trans_conv1\")\n",
    "        \n",
    "        batch_trans_conv1 = tf.layers.batch_normalization(inputs = trans_conv1, training=is_train, epsilon=1e-5, name=\"batch_trans_conv1\")\n",
    "       \n",
    "        trans_conv1_out = tf.nn.leaky_relu(batch_trans_conv1, alpha=alpha, name=\"trans_conv1_out\")\n",
    "        \n",
    "        \n",
    "        # Transposed conv 2 --> BatchNorm --> LeakyReLU\n",
    "        # 16x16x512 --> 32x32x256\n",
    "        trans_conv2 = tf.layers.conv2d_transpose(inputs = trans_conv1_out,\n",
    "                                  filters = 256,\n",
    "                                  kernel_size = [5,5],\n",
    "                                  strides = [2,2],\n",
    "                                  padding = \"SAME\",\n",
    "                                kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
    "                                name=\"trans_conv2\")\n",
    "        \n",
    "        batch_trans_conv2 = tf.layers.batch_normalization(inputs = trans_conv2, training=is_train, epsilon=1e-5, name=\"batch_trans_conv2\")\n",
    "       \n",
    "        trans_conv2_out = tf.nn.leaky_relu(batch_trans_conv2, alpha=alpha, name=\"trans_conv2_out\")\n",
    "        \n",
    "        \n",
    "        # Transposed conv 3 --> BatchNorm --> LeakyReLU\n",
    "        # 32x32x256 --> 64x64x128\n",
    "        trans_conv3 = tf.layers.conv2d_transpose(inputs = trans_conv2_out,\n",
    "                                  filters = 128,\n",
    "                                  kernel_size = [5,5],\n",
    "                                  strides = [2,2],\n",
    "                                  padding = \"SAME\",\n",
    "                                kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
    "                                name=\"trans_conv3\")\n",
    "        \n",
    "        batch_trans_conv3 = tf.layers.batch_normalization(inputs = trans_conv3, training=is_train, epsilon=1e-5, name=\"batch_trans_conv3\")\n",
    "       \n",
    "        trans_conv3_out = tf.nn.leaky_relu(batch_trans_conv3, alpha=alpha, name=\"trans_conv3_out\")\n",
    "\n",
    "        \n",
    "        # Transposed conv 4 --> BatchNorm --> LeakyReLU\n",
    "        # 64x64x128 --> 128x128x64\n",
    "        trans_conv4 = tf.layers.conv2d_transpose(inputs = trans_conv3_out,\n",
    "                                  filters = 64,\n",
    "                                  kernel_size = [5,5],\n",
    "                                  strides = [2,2],\n",
    "                                  padding = \"SAME\",\n",
    "                                kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
    "                                name=\"trans_conv4\")\n",
    "        \n",
    "        batch_trans_conv4 = tf.layers.batch_normalization(inputs = trans_conv4, training=is_train, epsilon=1e-5, name=\"batch_trans_conv4\")\n",
    "       \n",
    "        trans_conv4_out = tf.nn.leaky_relu(batch_trans_conv4, alpha=alpha, name=\"trans_conv4_out\")\n",
    "\n",
    "        \n",
    "        # Transposed conv 5 --> tanh\n",
    "        # 128x128x64 --> 128x128x3\n",
    "        logits = tf.layers.conv2d_transpose(inputs = trans_conv4_out,\n",
    "                                  filters = 3,\n",
    "                                  kernel_size = [5,5],\n",
    "                                  strides = [1,1],\n",
    "                                  padding = \"SAME\",\n",
    "                                kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
    "                                name=\"logits\")\n",
    "         \n",
    "        out = tf.tanh(logits, name=\"out\")\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Discriminator\n",
    "<img src=\"assets/discriminator.png\" alt=\"Discriminator\"/>\n",
    "\n",
    "- Input is 128x128x3\n",
    "- Depths starting with 32 and then *2 depth as you add layers\n",
    "- No downsampling using only <b> strided conv layers with no maxpool layers </b>\n",
    "- No batchnorm in input layer\n",
    "\n",
    "<b> convolution > batch norm > leaky ReLU </b>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def discriminator(x, is_reuse=False, alpha = 0.2):\n",
    "    ''' Build the discriminator network.\n",
    "    \n",
    "        Arguments\n",
    "        ---------\n",
    "        x : Input tensor for the discriminator\n",
    "        n_units: Number of units in hidden layer\n",
    "        reuse : Reuse the variables with tf.variable_scope\n",
    "        alpha : leak parameter for leaky ReLU\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        out, logits: \n",
    "    '''\n",
    "    with tf.variable_scope(\"discriminator\", reuse = is_reuse): \n",
    "        \n",
    "        # Input layer 128*128*3 --> 64x64x64\n",
    "        # Conv --> BatchNorm --> LeakyReLU   \n",
    "        conv1 = tf.layers.conv2d(inputs = x,\n",
    "                                filters = 64,\n",
    "                                kernel_size = [5,5],\n",
    "                                strides = [2,2],\n",
    "                                padding = \"SAME\",\n",
    "                                kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
    "                                name='conv1')\n",
    "        \n",
    "        batch_norm1 = tf.layers.batch_normalization(conv1,\n",
    "                                                   training = True,\n",
    "                                                   epsilon = 1e-5,\n",
    "                                                     name = 'batch_norm1')\n",
    "\n",
    "        conv1_out = tf.nn.leaky_relu(batch_norm1, alpha=alpha, name=\"conv1_out\")\n",
    "        \n",
    "        \n",
    "        # 64x64x64--> 32x32x128\n",
    "        # Conv --> BatchNorm --> LeakyReLU   \n",
    "        conv2 = tf.layers.conv2d(inputs = conv1_out,\n",
    "                                filters = 128,\n",
    "                                kernel_size = [5, 5],\n",
    "                                strides = [2, 2],\n",
    "                                padding = \"SAME\",\n",
    "                                kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
    "                                name='conv2')\n",
    "        \n",
    "        batch_norm2 = tf.layers.batch_normalization(conv2,\n",
    "                                                   training = True,\n",
    "                                                   epsilon = 1e-5,\n",
    "                                                     name = 'batch_norm2')\n",
    "        \n",
    "        conv2_out = tf.nn.leaky_relu(batch_norm2, alpha=alpha, name=\"conv2_out\")\n",
    "\n",
    "        \n",
    "        \n",
    "        # 32x32x128 --> 16x16x256\n",
    "        # Conv --> BatchNorm --> LeakyReLU   \n",
    "        conv3 = tf.layers.conv2d(inputs = conv2_out,\n",
    "                                filters = 256,\n",
    "                                kernel_size = [5, 5],\n",
    "                                strides = [2, 2],\n",
    "                                padding = \"SAME\",\n",
    "                                kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
    "                                name='conv3')\n",
    "        \n",
    "        batch_norm3 = tf.layers.batch_normalization(conv3,\n",
    "                                                   training = True,\n",
    "                                                   epsilon = 1e-5,\n",
    "                                                name = 'batch_norm3')\n",
    "        \n",
    "        conv3_out = tf.nn.leaky_relu(batch_norm3, alpha=alpha, name=\"conv3_out\")\n",
    "\n",
    "        \n",
    "        \n",
    "        # 16x16x256 --> 16x16x512\n",
    "        # Conv --> BatchNorm --> LeakyReLU   \n",
    "        conv4 = tf.layers.conv2d(inputs = conv3_out,\n",
    "                                filters = 512,\n",
    "                                kernel_size = [5, 5],\n",
    "                                strides = [1, 1],\n",
    "                                padding = \"SAME\",\n",
    "                                kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
    "                                name='conv4')\n",
    "        \n",
    "        batch_norm4 = tf.layers.batch_normalization(conv4,\n",
    "                                                   training = True,\n",
    "                                                   epsilon = 1e-5,\n",
    "                                                name = 'batch_norm4')\n",
    "        \n",
    "        conv4_out = tf.nn.leaky_relu(batch_norm4, alpha=alpha, name=\"conv4_out\")\n",
    "\n",
    "        \n",
    "        \n",
    "        # 16x16x512 --> 8x8x1024\n",
    "        # Conv --> BatchNorm --> LeakyReLU   \n",
    "        conv5 = tf.layers.conv2d(inputs = conv4_out,\n",
    "                                filters = 1024,\n",
    "                                kernel_size = [5, 5],\n",
    "                                strides = [2, 2],\n",
    "                                padding = \"SAME\",\n",
    "                                kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),\n",
    "                                name='conv5')\n",
    "        \n",
    "        batch_norm5 = tf.layers.batch_normalization(conv5,\n",
    "                                                   training = True,\n",
    "                                                   epsilon = 1e-5,\n",
    "                                                name = 'batch_norm5')\n",
    "        \n",
    "        conv5_out = tf.nn.leaky_relu(batch_norm5, alpha=alpha, name=\"conv5_out\")\n",
    "\n",
    "         \n",
    "        # Flatten it\n",
    "        flatten = tf.reshape(conv5_out, (-1, 8*8*1024))\n",
    "        \n",
    "        # Logits\n",
    "        logits = tf.layers.dense(inputs = flatten,\n",
    "                                units = 1,\n",
    "                                activation = None)\n",
    "        \n",
    "        \n",
    "        out = tf.sigmoid(logits)\n",
    "        \n",
    "        return out, logits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Discriminator and generator losses"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We train the g and d <b> at the same time </b> so we need losses for <b> both networks </b>\n",
    "\n",
    "#### Discriminator Loss\n",
    "Sum of loss for real and fake images\n",
    "<br>\n",
    "`d_loss = d_loss_real + d_loss_fake`\n",
    "<br><br>\n",
    "The losses will by <b> sigmoid cross entropy + wrap with tf.reduce_mean to get the mean for all the images in the batch.\n",
    "</b>\n",
    "\n",
    "##### Real image loss\n",
    "- Use `d_logits_real` and labels <b> are all 1 (since all real data is real) </b>\n",
    "- Label smoothing:  To help the discriminator generalize better, the labels are <b>reduced a bit from 1.0 to 0.9</b>\n",
    "`labels = tf.ones_like(tensor) * (1 - smooth)`\n",
    "For the real image loss, use the real logits and (smoothed) labels of ones. \n",
    "\n",
    "##### Fake image loss\n",
    "- Remember that we want the discriminator to output 1 for real images and 0 for fake images, so we need to set up the losses to reflect that\n",
    "- For the fake image loss, use the fake logits with labels of all zeros\n",
    "\n",
    "#### Generator Loss\n",
    "- The generator loss again uses the fake logits from the discriminator, but this time the labels are all ones because the generator wants to fool the discriminator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_loss(input_real, input_z, output_channel_dim, alpha):\n",
    "    \"\"\"\n",
    "    Get the loss for the discriminator and generator\n",
    "    :param input_real: Images from the real dataset\n",
    "    :param input_z: Z input\n",
    "    :param out_channel_dim: The number of channels in the output image\n",
    "    :return: A tuple of (discriminator loss, generator loss)\n",
    "    \"\"\"\n",
    "    # Generator network here\n",
    "    g_model = generator(input_z, output_channel_dim)   \n",
    "    # g_model is the generator output\n",
    "    \n",
    "    # Discriminator network here\n",
    "    d_model_real, d_logits_real = discriminator(input_real, alpha=alpha)\n",
    "    d_model_fake, d_logits_fake = discriminator(g_model,is_reuse=True, alpha=alpha)\n",
    "    \n",
    "    # Calculate losses\n",
    "    d_loss_real = tf.reduce_mean(\n",
    "                  tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_real, \n",
    "                                                          labels=tf.ones_like(d_model_real)))\n",
    "    d_loss_fake = tf.reduce_mean(\n",
    "                  tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, \n",
    "                                                          labels=tf.zeros_like(d_model_fake)))\n",
    "    d_loss = d_loss_real + d_loss_fake\n",
    "\n",
    "    g_loss = tf.reduce_mean(\n",
    "             tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake,\n",
    "                                                     labels=tf.ones_like(d_model_fake)))\n",
    "    \n",
    "    return d_loss, g_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimizers\n",
    "<ul>\n",
    "    <li>Update the generator and discriminator <b>separately</b></li>\n",
    "    <li> So we need to get the var for each part : we use `tf.trainable_variables()`. This creates a list of all the variables we've defined in our graph. </li>\n",
    "</ul>\n",
    "- The train operations are wrapped in a with tf.control_dependencies block so the batch normalization layers can update their population statistics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_optimizers(d_loss, g_loss, lr_D, lr_G, beta1):\n",
    "    \"\"\"\n",
    "    Get optimization operations\n",
    "    :param d_loss: Discriminator loss Tensor\n",
    "    :param g_loss: Generator loss Tensor\n",
    "    :param learning_rate: Learning Rate Placeholder\n",
    "    :param beta1: The exponential decay rate for the 1st moment in the optimizer\n",
    "    :return: A tuple of (discriminator training operation, generator training operation)\n",
    "    \"\"\"    \n",
    "    # Get the trainable_variables, split into G and D parts\n",
    "    t_vars = tf.trainable_variables()\n",
    "    g_vars = [var for var in t_vars if var.name.startswith(\"generator\")]\n",
    "    d_vars = [var for var in t_vars if var.name.startswith(\"discriminator\")]\n",
    "    \n",
    "    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)\n",
    "    \n",
    "    # Generator update\n",
    "    gen_updates = [op for op in update_ops if op.name.startswith('generator')]\n",
    "    \n",
    "    # Optimizers\n",
    "    with tf.control_dependencies(gen_updates):\n",
    "        d_train_opt = tf.train.AdamOptimizer(learning_rate=lr_D, beta1=beta1).minimize(d_loss, var_list=d_vars)\n",
    "        g_train_opt = tf.train.AdamOptimizer(learning_rate=lr_G, beta1=beta1).minimize(g_loss, var_list=g_vars)\n",
    "        \n",
    "    return d_train_opt, g_train_opt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training 🏃‍♂️\n",
    "### Show output\n",
    "Use this function to show the current output of the generator during training. It will help you determine how well the GANs is training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_generator_output(sess, n_images, input_z, out_channel_dim, image_mode, image_path, save, show):\n",
    "    \"\"\"\n",
    "    Show example output for the generator\n",
    "    :param sess: TensorFlow session\n",
    "    :param n_images: Number of Images to display\n",
    "    :param input_z: Input Z Tensor\n",
    "    :param out_channel_dim: The number of channels in the output image\n",
    "    :param image_mode: The mode to use for images (\"RGB\" or \"L\")\n",
    "    :param image_path: Path to save the image\n",
    "    \"\"\"\n",
    "    cmap = None if image_mode == 'RGB' else 'gray'\n",
    "    z_dim = input_z.get_shape().as_list()[-1]\n",
    "    example_z = np.random.uniform(-1, 1, size=[n_images, z_dim])\n",
    "\n",
    "    samples = sess.run(\n",
    "        generator(input_z, out_channel_dim, False),\n",
    "        feed_dict={input_z: example_z})\n",
    "\n",
    "    images_grid = helper.images_square_grid(samples, image_mode)\n",
    "    \n",
    "    if save == True:\n",
    "        # Save image\n",
    "        images_grid.save(image_path, 'JPEG')\n",
    "    \n",
    "    if show == True:\n",
    "        plt.imshow(images_grid, cmap=cmap)\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epoch_count, batch_size, z_dim, learning_rate_D, learning_rate_G, beta1, get_batches, data_shape, data_image_mode, alpha):\n",
    "    \"\"\"\n",
    "    Train the GAN\n",
    "    :param epoch_count: Number of epochs\n",
    "    :param batch_size: Batch Size\n",
    "    :param z_dim: Z dimension\n",
    "    :param learning_rate: Learning Rate\n",
    "    :param beta1: The exponential decay rate for the 1st moment in the optimizer\n",
    "    :param get_batches: Function to get batches\n",
    "    :param data_shape: Shape of the data\n",
    "    :param data_image_mode: The image mode to use for images (\"RGB\" or \"L\")\n",
    "    \"\"\"\n",
    "    # Create our input placeholders\n",
    "    input_images, input_z, lr_G, lr_D = model_inputs(data_shape[1:], z_dim)\n",
    "        \n",
    "    # Losses\n",
    "    d_loss, g_loss = model_loss(input_images, input_z, data_shape[3], alpha)\n",
    "    \n",
    "    # Optimizers\n",
    "    d_opt, g_opt = model_optimizers(d_loss, g_loss, lr_D, lr_G, beta1)\n",
    "    \n",
    "    i = 0\n",
    "    \n",
    "    version = \"firstTrain\"\n",
    "    with tf.Session() as sess:\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        \n",
    "        # Saver\n",
    "        saver = tf.train.Saver()\n",
    "        \n",
    "        num_epoch = 0\n",
    "        \n",
    "        if from_checkpoint == True:\n",
    "            saver.restore(sess, \"./models/model.ckpt\")\n",
    "            \n",
    "            show_generator_output(sess, 4, input_z, data_shape[3], data_image_mode, image_path, True, False)\n",
    "            \n",
    "        else:\n",
    "            for epoch_i in range(epoch_count):        \n",
    "                num_epoch += 1\n",
    "\n",
    "                if num_epoch % 5 == 0:\n",
    "\n",
    "                    # Save model every 5 epochs\n",
    "                    #if not os.path.exists(\"models/\" + version):\n",
    "                    #    os.makedirs(\"models/\" + version)\n",
    "                    save_path = saver.save(sess, \"./models/model.ckpt\")\n",
    "                    print(\"Model saved\")\n",
    "\n",
    "                for batch_images in get_batches(batch_size):\n",
    "                    # Random noise\n",
    "                    batch_z = np.random.uniform(-1, 1, size=(batch_size, z_dim))\n",
    "\n",
    "                    i += 1\n",
    "\n",
    "                    # Run optimizers\n",
    "                    _ = sess.run(d_opt, feed_dict={input_images: batch_images, input_z: batch_z, lr_D: learning_rate_D})\n",
    "                    _ = sess.run(g_opt, feed_dict={input_images: batch_images, input_z: batch_z, lr_G: learning_rate_G})\n",
    "\n",
    "                    if i % 10 == 0:\n",
    "                        train_loss_d = d_loss.eval({input_z: batch_z, input_images: batch_images})\n",
    "                        train_loss_g = g_loss.eval({input_z: batch_z})\n",
    "\n",
    "                        # Save it\n",
    "                        image_name = str(i) + \".jpg\"\n",
    "                        image_path = \"./images/\" + image_name\n",
    "                        show_generator_output(sess, 4, input_z, data_shape[3], data_image_mode, image_path, True, False) \n",
    "\n",
    "                    # Print every 5 epochs (for stability overwize the jupyter notebook will bug)\n",
    "                    if i % 1500 == 0:\n",
    "\n",
    "                        image_name = str(i) + \".jpg\"\n",
    "                        image_path = \"./images/\" + image_name\n",
    "                        print(\"Epoch {}/{}...\".format(epoch_i+1, epochs),\n",
    "                              \"Discriminator Loss: {:.4f}...\".format(train_loss_d),\n",
    "                              \"Generator Loss: {:.4f}\".format(train_loss_g))\n",
    "                        show_generator_output(sess, 4, input_z, data_shape[3], data_image_mode, image_path, False, True)\n",
    "                \n",
    "            \n",
    "                    \n",
    "    return losses, samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyperparameters\n",
    "Gans are <b> very sensitive to hyperparemeters </b>\n",
    "In general, you want the discriminator loss to be around 0.3, this means it is correctly classifying images as fake or real about 50% of the time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Size input image for discriminator\n",
    "real_size = (128,128,3)\n",
    "\n",
    "# Size of latent vector to generator\n",
    "z_dim = 100\n",
    "learning_rate_D =  .00005 # Thanks to Alexia Jolicoeur Martineau https://ajolicoeur.wordpress.com/cats/\n",
    "learning_rate_G = 2e-4 # Thanks to Alexia Jolicoeur Martineau https://ajolicoeur.wordpress.com/cats/\n",
    "batch_size = 64\n",
    "epochs = 215\n",
    "alpha = 0.2\n",
    "beta1 = 0.5\n",
    "\n",
    "# Create the network\n",
    "#model = DGAN(real_size, z_size, learning_rate, alpha, beta1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the data and train the network here\n",
    "dataset = helper.Dataset(glob(os.path.join(data_resized_dir, '*.jpg')))\n",
    "\n",
    "with tf.Graph().as_default():\n",
    "    losses, samples = train(epochs, batch_size, z_dim, learning_rate_D, learning_rate_G, beta1, dataset.get_batches,\n",
    "          dataset.shape, dataset.image_mode, alpha)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Loss 📈"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "losses = np.array(losses)\n",
    "plt.plot(losses.T[0], label='Discriminator', alpha=0.5)\n",
    "plt.plot(losses.T[1], label='Generator', alpha=0.5)\n",
    "plt.title(\"Training Losses\")\n",
    "plt.legend()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:gameplai]",
   "language": "python",
   "name": "conda-env-gameplai-py"
  },
  "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": 2
}