Turn a GAN into a WGAN and then WGAN GP.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to turn a vanilla GAN into a vanilla WGAN\n",
    "This notebook will help to enhance on the understanding of WGAN algorithm and WGAN_GP by changing a vanilla GAN into a WGAN, and then a WGAN_GP<br/>\n",
    "The GAN code is based on Kojin's GAN workshop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Different Loss Functions \n",
    "### GAN\n",
    "- $D$: $max_D V(D) = E_{x\\sim p_{data}}(x)[logD(x)] + E_{z\\sim p_z(z)}[log(1-D(G(z))]$\n",
    "- $G$: $min_G V(G) = E_{z\\sim p_z(z)}[log(1-D(G(z))]$\n",
    "\n",
    "### WGAN\n",
    "- $D$: $max_D V(D) = E_{x\\sim p_{data}}[D(x)] - E_{z\\sim p_z(z)}[D(G(z)]$\n",
    "- $G$: $max_G V(G) = E_{z\\sim p_z(z)}[D(G(z)]$\n",
    "\n",
    "### WGAN-GP\n",
    "- $D$: $max_D V(D) = E_{x\\sim p_{data}}[D(x)] - E_{z\\sim p_z(z)}[D(G(z)] + \\lambda(\\lVert \\nabla D(\\hat{x}) \\rVert_{2}-1)^2$\n",
    "- $G$: $max_G V(G) = E_{z\\sim p_z(z)}[D(G(z)]$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "|Features                      |  GAN   | WGAN |WGAN-GP\n",
    "| --------------------------- |:------:|:----: |:-----\n",
    "|output layer of Discriminator |Sigmoid | Linear |Linear\n",
    "|optimizer                     | Adam   | RMS  | Adam\n",
    "|weight clipping               | False  | True | False\n",
    "|Batch Normalization           | False  | True | False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### * Batch Normalization (not covered here)\n",
    "batch normalization normalize the output of each activation layer by subtracting the batch mean and dividing by the batch standard deviation to stablize training. <br/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preperation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "matplotlib.use('TkAgg')\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import MNIST data\n",
    "from tensorflow.examples.tutorials.mnist import input_data\n",
    "mnist = input_data.read_data_sets('MNIST_data', one_hot=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Network Params\n",
    "image_dim = 784 # 28*28 pixels\n",
    "gen_hidden_dim = 256\n",
    "disc_hidden_dim = 256\n",
    "noise_dim = 128 # Noise data points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph() # Clearing all tensors before this"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GAN to WGAN "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# <change> add hyperparameter Critic_Iters and c for WGAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# trainning param\n",
    "Batch_Size = 50\n",
    "Critic_Iters = 5 # for WGAN and WGAN-GP, number of critic iters per gen iter\n",
    "c = 0.01 # threshold for weight cliping (-c,c)\n",
    "Iters = 20001 # number of generator iterations to train for"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generator\n",
    "def generator(noises, reuse=False):\n",
    "    with tf.variable_scope('generator') as scope:\n",
    "        if (reuse):\n",
    "            tf.get_variable_scope().reuse_variables()\n",
    "        # hidden layer with name \"g_hidden\"\n",
    "        hidden = tf.layers.dense(noises, gen_hidden_dim, tf.nn.relu, name='g_hidden')\n",
    "        # out layer with name \"g_out\"\n",
    "        out_images = tf.layers.dense(hidden, image_dim, tf.nn.sigmoid, name='g_out')\n",
    "    return out_images\n",
    "\n",
    "# Discriminator\n",
    "def discriminator(images, reuse=False):\n",
    "    with tf.variable_scope('discriminator') as scope:\n",
    "        if (reuse):\n",
    "            tf.get_variable_scope().reuse_variables()            \n",
    "        # hidden layer with name \"d_hidden\"\n",
    "        hidden = tf.layers.dense(images, disc_hidden_dim, tf.nn.relu, name='d_hidden') \n",
    "        # out layer with name \"d_out\"\n",
    "        out = tf.layers.dense(hidden, 1, None, name='d_out') # <change> ReLU output into linear activation\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_input = tf.placeholder(tf.float32, shape=[None, noise_dim], name='input_noise')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fake_data = generator(gen_input)\n",
    "real_data = tf.placeholder(tf.float32, shape=[None, image_dim], name='real_data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "disc_real = discriminator(real_data)\n",
    "disc_fake = discriminator(fake_data, reuse=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# <change> cost function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_cost = -tf.reduce_mean(disc_fake)\n",
    "disc_cost = tf.reduce_mean(disc_fake) - tf.reduce_mean(disc_real)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tvars = tf.trainable_variables()\n",
    "disc_vars = [var for var in tvars if 'd_' in var.name]\n",
    "gen_vars = [var for var in tvars if 'g_' in var.name]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# <change> optimizer to RMS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_gen = tf.train.RMSPropOptimizer(\n",
    "        learning_rate=5e-5, \n",
    "    ).minimize(gen_cost, var_list=gen_vars)\n",
    "\n",
    "train_disc = tf.train.RMSPropOptimizer(\n",
    "        learning_rate=5e-5, \n",
    "    ).minimize(disc_cost, var_list=disc_vars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# <change> add weight clipping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clip_D = [p.assign(tf.clip_by_value(p,-c,c)) for p in disc_vars]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.Session() as sess:\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    for step in range(Iters):\n",
    "\n",
    "        batch_x, _ = mnist.train.next_batch(Batch_Size)\n",
    "        # Generate noise to feed to the generator\n",
    "        z = np.random.uniform(-1., 1., size=[Batch_Size, noise_dim])\n",
    "        \n",
    "        # train discriminator\n",
    "        for i in range(Critic_Iters): #have inner loop for discriminator\n",
    "            _,dl = sess.run([train_disc,disc_cost],\n",
    "                                   feed_dict={real_data:batch_x,gen_input:z})\n",
    "            _ = sess.run(clip_D) #put weight_clipping here\n",
    "        \n",
    "        # train generator\n",
    "        _,gl=sess.run([train_gen,gen_cost],\n",
    "                      feed_dict={gen_input:z})\n",
    "        \n",
    "        if step % 1000 == 0 or step == 1:\n",
    "            print('Step %i: Generator Loss: %f, Discriminator Loss: %f' % (step, gl, dl))\n",
    "            \n",
    "        # Generate images from noise, using the generator network.\n",
    "        if step % 10000 == 0 or step == 1:\n",
    "            f, a = plt.subplots(4, 10, figsize=(10, 4))\n",
    "            for i in range(10):\n",
    "                # Noise input.\n",
    "                z = np.random.uniform(-1., 1., size=[4, noise_dim])\n",
    "                g = sess.run([fake_data], feed_dict={gen_input: z})\n",
    "                g = np.reshape(g, newshape=(4, 28, 28, 1))\n",
    "                # Reverse colours for better display\n",
    "                g = -1 * (g - 1)\n",
    "                img_inventory[\"WGAN_\"+str(step)] = g\n",
    "                \n",
    "                for j in range(4):\n",
    "                    # Generate image from noise. Extend to 3 channels for matplot figure.\n",
    "                    img = np.reshape(np.repeat(g[j][:, :, np.newaxis], 3, axis=2),\n",
    "                                     newshape=(28, 28, 3))\n",
    "                    a[j][i].imshow(img)\n",
    "\n",
    "            plt.draw()\n",
    "            print('wgan'+str(step)+'.png')\n",
    "            plt.savefig('wgan'+str(step)+'.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## WGAN to WGAN-GP "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "WGAN-GP uses the similar loss function as WGAN, but WGAN-GP has extra regularization component. <br/>\n",
    "Besides, they optimize in different ways:<br/>\n",
    "WGAN-GP uses AdamOptimizer, WGAN uses RMSOptimizer <br/>\n",
    "WGAN-GP does not require weight clipping <br/>\n",
    "WGAN-GP needs to hand pick penalty coefficient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph() # Clearing all tensors before this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# <change> add hyperparameter Critic_Iters and Lambda for WGAN-GP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# trainning param\n",
    "Batch_Size = 50\n",
    "Critic_Iters = 5 # for WGAN and WGAN-GP, number of critic iters per gen iter\n",
    "Lambda = 10 # gradient penalty lambda hyperparameter\n",
    "Iters = 200000 # number of generator iterations to train for"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generator\n",
    "def generator(noises, reuse=False):\n",
    "    with tf.variable_scope('generator') as scope:\n",
    "        if (reuse):\n",
    "            tf.get_variable_scope().reuse_variables()\n",
    "        # hidden layer with name \"g_hidden\"\n",
    "        hidden = tf.layers.dense(noises, gen_hidden_dim, tf.nn.relu, name='g_hidden')\n",
    "        # out layer with name \"g_out\"\n",
    "        out_images = tf.layers.dense(hidden, image_dim, tf.nn.sigmoid, name='g_out')\n",
    "    return out_images\n",
    "\n",
    "# Discriminator\n",
    "def discriminator(images, reuse=False):\n",
    "    with tf.variable_scope('discriminator') as scope:\n",
    "        if (reuse):\n",
    "            tf.get_variable_scope().reuse_variables()            \n",
    "        # hidden layer with name \"d_hidden\"\n",
    "        hidden = tf.layers.dense(images, disc_hidden_dim, tf.nn.relu, name='d_hidden')\n",
    "        # out layer with name \"d_out\"\n",
    "        out = tf.layers.dense(hidden, 1, None, name='d_out') # <change> output layer turn into a linear one as WGAN does\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_input = tf.placeholder(tf.float32, shape=[None, noise_dim], name='input_noise')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fake_data = generator(gen_input)\n",
    "real_data = tf.placeholder(tf.float32, shape=[None, image_dim], name='real_data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "disc_real = discriminator(real_data)\n",
    "disc_fake = discriminator(fake_data, reuse=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_cost = -tf.reduce_mean(disc_fake)\n",
    "disc_cost = tf.reduce_mean(disc_fake) - tf.reduce_mean(disc_real)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tvars = tf.trainable_variables()\n",
    "disc_vars = [var for var in tvars if 'd_' in var.name]\n",
    "gen_vars = [var for var in tvars if 'g_' in var.name]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# <change> add regularization component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = tf.random_uniform(shape=[Batch_Size,1],minval=0.,maxval=1.)\n",
    "differences = fake_data-real_data\n",
    "interpolates = real_data + (alpha*differences)\n",
    "gradients = tf.gradients(discriminator(interpolates, reuse=True),[interpolates])[0]\n",
    "slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients),reduction_indices=[1]))\n",
    "gradient_penalty = tf.reduce_mean((slopes-1.)**2)\n",
    "disc_cost += Lambda*gradient_penalty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# keep using AdamOptimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_gen = tf.train.AdamOptimizer(\n",
    "        learning_rate=1e-4, \n",
    "        beta1=0.5,\n",
    "        beta2=0.9\n",
    "    ).minimize(gen_cost, var_list=gen_vars)\n",
    "\n",
    "train_disc = tf.train.AdamOptimizer(\n",
    "        learning_rate=1e-4, \n",
    "        beta1=0.5, \n",
    "        beta2=0.9\n",
    "    ).minimize(disc_cost, var_list=disc_vars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.Session() as sess:\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "\n",
    "    for step in range(Iters):\n",
    "\n",
    "        batch_x, _ = mnist.train.next_batch(Batch_Size)\n",
    "        # Generate noise to feed to the generator\n",
    "        z = np.random.uniform(-1., 1., size=[Batch_Size, noise_dim])\n",
    "        \n",
    "        # train discriminator\n",
    "        for i in range(Critic_Iters):\n",
    "            _,dl = sess.run([train_disc,disc_cost],\n",
    "                                   feed_dict={real_data:batch_x,gen_input:z})\n",
    "        \n",
    "        # train generator\n",
    "        _,gl=sess.run([train_gen,gen_cost],\n",
    "                      feed_dict={gen_input:z})\n",
    "        \n",
    "        if step % 1000 == 0 or step == 1:\n",
    "            print('Step %i: Generator Loss: %f, Discriminator Loss: %f' % (step, gl, dl))\n",
    "    \n",
    "        # Generate images from noise, using the generator network.\n",
    "        if step % 10000 == 0 or step == 1:\n",
    "            f, a = plt.subplots(4, 10, figsize=(10, 4))\n",
    "            for i in range(10):\n",
    "                # Noise input.\n",
    "                z = np.random.uniform(-1., 1., size=[4, noise_dim])\n",
    "                g = sess.run([fake_data], feed_dict={gen_input: z})\n",
    "                g = np.reshape(g, newshape=(4, 28, 28, 1))\n",
    "                # Reverse colours for better display\n",
    "                g = -1 * (g - 1)\n",
    "                for j in range(4):\n",
    "                    # Generate image from noise. Extend to 3 channels for matplot figure.\n",
    "                    img = np.reshape(np.repeat(g[j][:, :, np.newaxis], 3, axis=2),\n",
    "                                     newshape=(28, 28, 3))\n",
    "                    a[j][i].imshow(img)\n",
    "\n",
    "            plt.draw()\n",
    "            print('wgan_gp'+str(step)+'.png')\n",
    "            plt.savefig('wgan_gp'+str(step)+'.png')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:tensorflow_workshop]",
   "language": "python",
   "name": "conda-env-tensorflow_workshop-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
}