DCGAN-FACES_AWS.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h2>Deep Convolutional Generative Adversarial Nets (DCGAN) - generate CelebA faces</h2>\n",
    "\n",
    "<p>Original DCGAN paper: <a href='https://arxiv.org/pdf/1511.06434.pdf'>https://arxiv.org/pdf/1511.06434.pdf</a></p>\n",
    "\n",
    "<p>The model structure is borrowed from and modified based on Arthur Juliani's implementation: <a href='https://medium.com/@awjuliani/generative-adversarial-networks-explained-with-a-classic-spongebob-squarepants-episode-54deab2fce39#.4kvc6juwc'>https://medium.com/@awjuliani/generative-adversarial-networks-explained-with-a-classic-spongebob-squarepants-episode-54deab2fce39#.4kvc6juwc</a></p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import tensorflow.contrib.slim as slim\n",
    "import os\n",
    "import time\n",
    "from PIL import Image\n",
    "from images2gif import writeGif"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "np.random.seed(1) \n",
    "tf.set_random_seed(1)\n",
    "\n",
    "num_train_imgs = 10000 # number of training images         \n",
    "original_img_directory = './celeba_align_178_218/'\n",
    "resized_img_directory = 'celeba_resizesd/'\n",
    "\n",
    "pre_train_epochs = 1 # num of epochs to pre-train discriminator\n",
    "training_epochs = 20001 # num of epochs to train discriminator and generator jointly\n",
    "batch_size = 64 # batch size for training the model, must be larger than batch_size_sample\n",
    "\n",
    "# for the generated images\n",
    "batch_size_sample = 36 # number of images to be generated for viewing the training progress\n",
    "# the generated images will be concatenated, the product of the following must equal to batch_size_sample\n",
    "sample_num_rows = 6 # how many rows to have\n",
    "sample_num_columns = 6 # how many columns to have\n",
    "\n",
    "leak = 0.2 # degree of leakiness used in leaky ReLU\n",
    "alpha = 0.0002 # base learning rate\n",
    "beta1 = 0.5 # the fraction factor used in the first momentum term from Adam optimizer\n",
    "k = 1 # this is number of times to update generator for every time discriminator is updated in each epoch\n",
    "logs_path = \"./dcgan_face_log\" # directory to save the training log to\n",
    "train_sample_directory = './dcgan_face/train_sample/' # directory to save the generated images during training\n",
    "model_directory = './dcgan_face/models' # directory to save trained model\n",
    "sample_directory = './dcgan_face/generated_sample/' # directory to save the generated images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Just once - Preprocess training images by resizing them</h3>\n",
    "<p>Face images can be found at: <a href=\"http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html\">http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html</a></p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# get a list of image filenames\n",
    "file_names = [f for f in os.listdir(original_img_directory)[:num_train_imgs+1] if f.endswith('.jpg')] \n",
    "# get a list of images\n",
    "images = []\n",
    "for f in file_names:\n",
    "    img = Image.open(original_img_directory+f)\n",
    "    images.append(img.copy())t\n",
    "    img.close() \n",
    "\n",
    "if not os.path.exists(resized_img_directory):\n",
    "    os.makedirs(resized_img_directory)\n",
    "\n",
    "# save the resized images\n",
    "for i,img in enumerate(images):\n",
    "    img = resizeimage.resize_contain(img, [64, 64]) # resize the image from 178*218 to 64*64\n",
    "    img.save(resized_img_directory+'/'+file_names[i], img.format)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "del images # remove images variable to free memory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read training images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# get a list of image filenames\n",
    "file_names = [f for f in os.listdir(resized_img_directory) if f.endswith('.jpg')] \n",
    "# get a list of images\n",
    "images = []\n",
    "for f in file_names:\n",
    "    img = Image.open(resized_img_directory+f)\n",
    "    images.append(img.copy())\n",
    "    #img.close() \n",
    "    \n",
    "# turn image into ndarray\n",
    "train_images = np.array([np.asarray(img) for img in images])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Declare helpful functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#This function performns a leaky relu activation, which is needed for the discriminator network.\n",
    "def leaky_relu(x, leak=leak, name=\"leaky_relu\"):\n",
    "    with tf.variable_scope(name):\n",
    "        return tf.maximum(leak*x, x)\n",
    "\n",
    "# get a batch of real images, change size from 28*28 to 32*32, return the processed batch\n",
    "def get_x(batch_size):\n",
    "    indices = np.random.randint(num_train_imgs, size=batch_size) # random sample real images\n",
    "    batch = train_images[indices] \n",
    "    batch = 2*(batch/255.)-1 # change range from [0, 255] to [-1, 1]\n",
    "    return batch\n",
    "\n",
    "# wrapper function for real images\n",
    "def save_real_images(images, size, image_path):\n",
    "    concat_img = concat(images, size) # concatenate individual images to a grid of images\n",
    "    concat_img = Image.fromarray(concat_img) # convert ndarray to an image object\n",
    "    concat_img.save(image_path, 'PNG')\n",
    "\n",
    "# wrapper function for generated images\n",
    "def save_generated_images(images, size, image_path):\n",
    "    images = inverse_transform(images)\n",
    "    concat_img = concat(images, size) # concatenate individual images to a grid of images\n",
    "    concat_img = Image.fromarray(concat_img) # convert ndarray to an image object\n",
    "    concat_img.save(image_path, 'PNG')\n",
    "\n",
    "# change values from [-1, 1] to [0, 255]\n",
    "def inverse_transform(images):\n",
    "    return (images+1.)/2. * 255\n",
    "\n",
    "# concatenate individual images to a grid of images\n",
    "def concat(images, size):\n",
    "    # get height and width for a single image generated, e.g. 64 * 64\n",
    "    height, width = images.shape[1], images.shape[2] \n",
    "    \n",
    "    # placeholder for a concatenated img, which have \n",
    "    # img_height = image height * num images vertically, e.g. 64 * 5\n",
    "    # img_width = image width * num images horizontally, e.g. 64 * 5\n",
    "    img = np.zeros((height * size[0], width * size[1],3)) \n",
    "\n",
    "    # loop through each image\n",
    "    for index, image in enumerate(images):\n",
    "        j = index / size[0] # image row index\n",
    "        i = index % size[1] # image column index\n",
    "        img[j*height:j*height+height, i*width:i*width+width,:] = image\n",
    "\n",
    "    return img.astype(np.uint8) # convert ndarray to uint8 type"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save orginal images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "if not os.path.exists(train_sample_directory):\n",
    "    os.makedirs(train_sample_directory)\n",
    "\n",
    "indices = np.random.randint(num_train_imgs, size=batch_size_sample) # random sample real images\n",
    "batch = train_images[indices] # with shape [batch,64,64,3]\n",
    "save_real_images(images=batch, size=[sample_num_rows,sample_num_columns], image_path=train_sample_directory+'/fig_original.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the Generative Adversarial Network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Generator network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# The generator network takes a noise vector z and return a 64 * 64 image. \n",
    "def generator(z, reuse=False, is_train=True):\n",
    "    \n",
    "    # Turn z into the first tensor + batch norm + relu\n",
    "    # Creates a fully connected weight matrix, which is multiplied by the mini-batch of z vectors to \n",
    "    # produce a hidden layer with 4*4*1024 hidden nodes\n",
    "    zP = slim.fully_connected(inputs=z,num_outputs=4*4*1024,normalizer_fn=slim.batch_norm, normalizer_params={\"is_training\":is_train},\\\n",
    "        reuse=reuse,activation_fn=tf.nn.relu,scope='g_init',weights_initializer=initializer)\n",
    "    # Transform the flat 4*4*1024 layer into a tensor, whose kernel size is 4*4 and 512 kernels in total\n",
    "    # -1 stands for the mini-batch size that's to be computed\n",
    "    zCon = tf.reshape(zP,[-1,4,4,1024])\n",
    "    \n",
    "    # Perform fractionally-strided convolution/deconvolution + batch norm + relu\n",
    "    # num_outputs: number of output filters/activation maps\n",
    "    # kernel_size: [kernel_height, kernel_width]\n",
    "    # stride: [stride_height, stride_width], in this case it is generating four pixels out of every pixel, this gives fractionally-strided convolution\n",
    "    # output size batch_size*8*8*512\n",
    "    gen1 = slim.convolution2d_transpose(\\\n",
    "        inputs=zCon,num_outputs=512,kernel_size=[5,5],stride=[2,2],\\\n",
    "        padding=\"SAME\",normalizer_fn=slim.batch_norm, normalizer_params={\"is_training\":is_train},\\\n",
    "        reuse=reuse,activation_fn=tf.nn.relu,scope='g_conv1', weights_initializer=initializer)\n",
    "    \n",
    "    # output size batch_size*16*16*256\n",
    "    gen2 = slim.convolution2d_transpose(\\\n",
    "        inputs=gen1,num_outputs=256,kernel_size=[5,5],stride=[2,2],\\\n",
    "        padding=\"SAME\",normalizer_fn=slim.batch_norm, normalizer_params={\"is_training\":is_train},\\\n",
    "        reuse=reuse,activation_fn=tf.nn.relu,scope='g_conv2', weights_initializer=initializer)\n",
    "    \n",
    "    # output size batch_size*32*32*128\n",
    "    gen3 = slim.convolution2d_transpose(\\\n",
    "        inputs=gen2,num_outputs=128,kernel_size=[5,5],stride=[2,2],\\\n",
    "        padding=\"SAME\",normalizer_fn=slim.batch_norm, normalizer_params={\"is_training\":is_train},\\\n",
    "        reuse=reuse,activation_fn=tf.nn.relu,scope='g_conv3', weights_initializer=initializer)\n",
    "    \n",
    "    # output size batch_size*64*64*3\n",
    "    g_out = slim.convolution2d_transpose(\\\n",
    "        inputs=gen3,num_outputs=3,kernel_size=[5,5],stride=[2,2],padding=\"SAME\",\\\n",
    "        biases_initializer=None,activation_fn=tf.nn.tanh,\\\n",
    "        reuse=reuse,scope='g_output', weights_initializer=initializer)\n",
    "    \n",
    "    return g_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Discriminator network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# The discriminator network takes a 64*64 image and return a probability of whether it is real or generated\n",
    "def discriminator(image, reuse=False, is_train=True):\n",
    "    # Creates 32 4*4 filters to convolve on the mini-batch of 64*64*3 images, also perform batch norm + leaky ReLU activation\n",
    "    # Note that no pooling is performed\n",
    "    # stride here calculates one pixel out of every 2*2 pixels, this gives strided convolution that will shrink the image as a substitute for pooling\n",
    "    # Set reuse=True allows discriminator to evaluate both real samples and generated samples \n",
    "    # Outputs size batch_size*32*32*32 \n",
    "    \n",
    "    dis1 = slim.convolution2d(inputs=image,num_outputs=32,kernel_size=[4,4],stride=[2,2],padding=\"SAME\",\\\n",
    "        biases_initializer=None,activation_fn=leaky_relu,\\\n",
    "        reuse=reuse,scope='d_conv1',weights_initializer=initializer)\n",
    "    \n",
    "    # outputs size batch_size*16*16*64 \n",
    "    dis2 = slim.convolution2d(inputs=dis1,num_outputs=64,kernel_size=[4,4],stride=[2,2],padding=\"SAME\",\\\n",
    "        normalizer_fn=slim.batch_norm, normalizer_params={\"is_training\":is_train}, activation_fn=leaky_relu,\\\n",
    "        reuse=reuse,scope='d_conv2', weights_initializer=initializer)\n",
    "    \n",
    "    # outputs size batch_size*8*8*128 \n",
    "    dis3 = slim.convolution2d(inputs=dis2,num_outputs=128,kernel_size=[4,4],stride=[2,2],padding=\"SAME\",\\\n",
    "        normalizer_fn=slim.batch_norm, normalizer_params={\"is_training\":is_train},activation_fn=leaky_relu,\\\n",
    "        reuse=reuse,scope='d_conv3',weights_initializer=initializer)\n",
    "    \n",
    "    # outputs size batch_size*4*4*256 \n",
    "    dis4 = slim.convolution2d(inputs=dis3,num_outputs=256,kernel_size=[4,4],stride=[2,2],padding=\"SAME\",\\\n",
    "        normalizer_fn=slim.batch_norm, normalizer_params={\"is_training\":is_train},activation_fn=leaky_relu,\\\n",
    "        reuse=reuse,scope='d_conv4',weights_initializer=initializer)\n",
    "\n",
    "    # flatten the tensor to [batch_size, 4*4*256]\n",
    "    dis_flat = slim.flatten(dis4)\n",
    "    \n",
    "    # create a fully connect layer with dis_flat and just one node in the output layer\n",
    "    # note there's no batch normalization at this layer\n",
    "    # outputs size batch_size*1 \n",
    "    d_out = slim.fully_connected(inputs=dis_flat,num_outputs=1,\\\n",
    "        activation_fn=tf.nn.sigmoid, reuse=reuse, scope='d_output', weights_initializer=initializer)\n",
    "    \n",
    "    return d_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connecting two networks together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "z_size = 100 # size of initial noise vector that will be used for generator\n",
    "\n",
    "# initialize all parameters of the networks\n",
    "# weights were initialized from a zero-centered Normal distribution with standard deviation 0.02\n",
    "# tf.truncated_normal returns random values from a normal distribution and made sure no value exceeds 2 std\n",
    "initializer = tf.truncated_normal_initializer(stddev=0.02)\n",
    "\n",
    "# placeholders for inputs into the generator and discriminator, respectively.\n",
    "z_vector = tf.placeholder(shape=[batch_size,z_size],dtype=tf.float32, name='z_vectors') \n",
    "x_vector = tf.placeholder(shape=[batch_size,64,64,3],dtype=tf.float32, name='real_images') \n",
    "whether_is_train = tf.placeholder(tf.bool) # boolean variable indicating if it is in training mode or not\n",
    "\n",
    "# ---- Pre-training ----\n",
    "# the discriminator should output probability=1 for all the training images\n",
    "train_labels=tf.constant(1.0,shape=(batch_size,1), name='pre_train_labels')\n",
    "# feed images to the discriminator and return the predicted probability\n",
    "d_pre_output = discriminator(x_vector)                                              \n",
    "summary_pre_d_x_hist = tf.histogram_summary(\"pre_train_d_prob_x\", d_pre_output)\n",
    "\n",
    "d_pre_loss=tf.reduce_mean(tf.square(d_pre_output-train_labels))\n",
    "summary_pre_d_loss = tf.scalar_summary(\"pre_train_d_loss\", d_pre_loss)\n",
    "# ---- end of Pre-training ----\n",
    "\n",
    "\n",
    "# ---- DCGAN ----\n",
    "g_output = tf.cond(whether_is_train, \n",
    "                   lambda: generator(z_vector,is_train=True), \n",
    "                   lambda: generator(z_vector,reuse=True,is_train=False))# generated mini-batch of images from noisy z vectors \n",
    "#g_output = generator(z_vector,is_train) # generated mini-batch of images from noisy z vectors \n",
    "    \n",
    "d_output_x = discriminator(x_vector,reuse=True) # probabilities for real images\n",
    "d_output_x = tf.maximum(tf.minimum(d_output_x, 0.99), 0.01) # avoid inf and -inf\n",
    "summary_d_x_hist = tf.histogram_summary(\"d_prob_x\", d_output_x)\n",
    "\n",
    "d_output_z = discriminator(g_output,reuse=True) # probabilities for generated images\n",
    "d_output_z = tf.maximum(tf.minimum(d_output_z, 0.99), 0.01) # avoid inf and -inf\n",
    "summary_d_z_hist = tf.histogram_summary(\"d_prob_z\", d_output_z)\n",
    "\n",
    "d_loss = -tf.reduce_mean(tf.log(d_output_x) + tf.log(1-d_output_z)) # loss for discriminator\n",
    "summary_d_loss = tf.scalar_summary(\"d_loss\", d_loss)\n",
    "\n",
    "g_loss = -tf.reduce_mean(tf.log(d_output_z)) # loss for generator\n",
    "summary_g_loss = tf.scalar_summary(\"g_loss\", g_loss)\n",
    "\n",
    "# the following parameter indices may change if the network structure changes\n",
    "para_d = tf.trainable_variables()[:9] # parameters for discriminator\n",
    "para_g = tf.trainable_variables()[9:] # parameters for generator\n",
    "\n",
    "# only update parameters in discriminator during pre-training\n",
    "pre_optimizer = tf.train.AdamOptimizer(learning_rate=alpha,beta1=beta1).minimize(d_pre_loss,var_list=para_d)\n",
    "# only update the weights for the discriminator network\n",
    "optimizer_op_d = tf.train.AdamOptimizer(learning_rate=alpha,beta1=beta1).minimize(d_loss,var_list=para_d)\n",
    "# only update the weights for the generator network\n",
    "optimizer_op_g = tf.train.AdamOptimizer(learning_rate=alpha,beta1=beta1).minimize(g_loss,var_list=para_g)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the DCGAN model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# create a log folder and save the graph structure, do this before training\n",
    "g_writer = tf.train.SummaryWriter(logs_path + '/generator', graph=tf.get_default_graph())\n",
    "d_writer = tf.train.SummaryWriter(logs_path + '/discriminator')\n",
    "\n",
    "# saver saves and loads variables of the model to and from checkpoints, \n",
    "# which are binary files that maps variable names to tensor values\n",
    "saver = tf.train.Saver(max_to_keep=20) \n",
    "\n",
    "\n",
    "with tf.Session(config=config) as sess:  \n",
    "    # variables need to be initialized before we can use them\n",
    "    sess.run(tf.initialize_all_variables())\n",
    "    #print [v.name for v in tf.trainable_variables()] # print all variable names\n",
    "    \n",
    "    \n",
    "    # -------- pre-train discriminator --------\n",
    "    start = time.time()\n",
    "    for epoch in range(pre_train_epochs):\n",
    "        x = get_x(batch_size) # get a batch of real images, with range [-1 ,1]\n",
    "        \n",
    "        d_pre_summary_merge = tf.merge_summary([summary_pre_d_loss, summary_pre_d_x_hist])\n",
    "        summary_pre_d,_=sess.run([d_pre_summary_merge,pre_optimizer], feed_dict={x_vector: x})\n",
    "        \n",
    "        d_writer.add_summary(summary_pre_d, epoch) # add loss summary to tensorboard\n",
    "        time_lapse = time.time()-start\n",
    "        start = time.time()\n",
    "        print \"pre-train epoch: \", epoch,\", time spent: %.2fs\" % time_lapse\n",
    "   \n",
    "    print \"pre-train done.\"\n",
    "    \n",
    "    \n",
    "    \n",
    "    # -------- jointly training discriminator and generator --------\n",
    "    start = time.time()\n",
    "    \n",
    "    # z noise vector that will be used to generate image to check the training progress\n",
    "    z_sample = np.random.uniform(-1.0,1.0,size=[batch_size,z_size]).astype(np.float32)\n",
    "    \n",
    "    for epoch in range(training_epochs):\n",
    "        # get a batch of real images, with range [-1 ,1]\n",
    "        x = get_x(batch_size) \n",
    "        # mini-batch of noise data from [-1, 1]\n",
    "        z = np.random.uniform(-1.0,1.0,size=[batch_size,z_size]).astype(np.float32)\n",
    "        \n",
    "        \n",
    "        if epoch <= 200:\n",
    "            # make a directory for generated images\n",
    "            if not os.path.exists(train_sample_directory):\n",
    "                os.makedirs(train_sample_directory)\n",
    "            # get a generated image, with range [-1, 1]\n",
    "            g_images = sess.run(g_output,feed_dict={z_vector:z_sample}) \n",
    "            # substitute 1/5 of the training images to the generated images, to increase discriminator's difficulty \n",
    "            substitute_indices = np.random.randint(batch_size, size=batch_size/5) \n",
    "            x[substitute_indices] = g_images[substitute_indices]\n",
    "\n",
    "        \n",
    "        # Update the discriminator\n",
    "        d_summary_merge = tf.merge_summary([summary_d_loss, summary_d_x_hist,summary_d_z_hist])\n",
    "        summary_d,_ = sess.run([d_summary_merge,optimizer_op_d],feed_dict={z_vector:z, x_vector:x}) \n",
    "        \n",
    "        \n",
    "        if epoch<20:\n",
    "            # Update the generator for k times in the specified n epoch\n",
    "            for i in range(k):\n",
    "                summary_g,_ = sess.run([summary_g_loss,optimizer_op_g],feed_dict={z_vector:z})\n",
    "        else:\n",
    "            summary_g,_ = sess.run([summary_g_loss,optimizer_op_g],feed_dict={z_vector:z})\n",
    "        \n",
    "        \n",
    "        # add loss summary to tensorboard\n",
    "        if epoch % 1 == 0:\n",
    "            d_writer.add_summary(summary_d, epoch) \n",
    "            g_writer.add_summary(summary_g, epoch)\n",
    "        \n",
    "        # output generated image\n",
    "        if epoch % 200 == 0 or epoch in [5,60,100,150]:\n",
    "            time_lapse = time.time()-start\n",
    "            start = time.time()\n",
    "            \n",
    "            print \"DCGAN epoch: \", epoch,\", time spent: %.2fs\" % time_lapse\n",
    "            g_images = sess.run(g_output,feed_dict={z_vector:z_sample}) # get a generated image, with range [-1, 1]\n",
    "            \n",
    "            # make a directory for generated images\n",
    "            if not os.path.exists(train_sample_directory):\n",
    "                os.makedirs(train_sample_directory)\n",
    "            \n",
    "            #Save sample generator images for viewing training progress.\n",
    "            save_generated_images(images = np.reshape(g_images[0:batch_size_sample],[batch_size_sample,64,64,3]),\\\n",
    "                        size = [sample_num_rows,sample_num_columns], image_path = train_sample_directory+'/'+str(epoch)+'.png')\n",
    "            \n",
    "        if epoch in [0,500,1000,2000,4000,6000,8000,10000,15000]:\n",
    "            # make a directory for trained models\n",
    "            if not os.path.exists(model_directory):\n",
    "                os.makedirs(model_directory)\n",
    "            \n",
    "            # save the trained model at different epoch\n",
    "            saver.save(sess, save_path = model_directory + '/' + str(epoch) + '.cptk')\n",
    "    print \"Done\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Continue training from the last checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "saver = tf.train.Saver(max_to_keep=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ckpt = tf.train.get_checkpoint_state(model_directory)\n",
    "ckpt.model_checkpoint_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print 'Loading models...might take a minute'\n",
    "saver = tf.train.Saver(max_to_keep=20)\n",
    "\n",
    "# create a log folder and save the graph structure, do this before training\n",
    "g_writer = tf.train.SummaryWriter(logs_path + '/generator', graph=tf.get_default_graph())\n",
    "d_writer = tf.train.SummaryWriter(logs_path + '/discriminator')\n",
    "\n",
    "config = tf.ConfigProto()\n",
    "config.gpu_options.allow_growth=True\n",
    "\n",
    "with tf.Session(config=config) as sess:  \n",
    "    start = time.time()\n",
    "    \n",
    "    ckpt = tf.train.get_checkpoint_state(model_directory)\n",
    "    model = ckpt.model_checkpoint_path\n",
    "\n",
    "    # z noise vector that will be used to generate image to check the training progress\n",
    "    z_sample = np.random.uniform(-1.0,1.0,size=[batch_size,z_size]).astype(np.float32)\n",
    "    \n",
    "    saver.restore(sess, save_path=model)\n",
    "    \n",
    "    for epoch in range(14500,training_epochs):\n",
    "        # get a batch of real images, with range [-1 ,1]\n",
    "        x = get_x(batch_size) \n",
    "        # mini-batch of noise data from [-1, 1]\n",
    "        z = np.random.uniform(-1.0,1.0,size=[batch_size,z_size]).astype(np.float32)\n",
    "        \n",
    "        # Update the discriminator\n",
    "        d_summary_merge = tf.merge_summary([summary_d_loss, summary_d_x_hist,summary_d_z_hist])\n",
    "        summary_d,_ = sess.run([d_summary_merge,optimizer_op_d],feed_dict={z_vector:z, x_vector:x}) \n",
    "        \n",
    "        # Update the generator\n",
    "        summary_g,_ = sess.run([summary_g_loss,optimizer_op_g],feed_dict={z_vector:z})\n",
    "        \n",
    "        # add loss summary to tensorboard\n",
    "        if epoch % 20 == 0:\n",
    "            d_writer.add_summary(summary_d, epoch) \n",
    "            g_writer.add_summary(summary_g, epoch)\n",
    "        \n",
    "        # output generated image\n",
    "        if epoch % 200 == 0:\n",
    "            time_lapse = time.time()-start\n",
    "            start = time.time()\n",
    "            \n",
    "            print \"DCGAN epoch: \", epoch,\", time spent: %.2fs\" % time_lapse\n",
    "            g_images = sess.run(g_output,feed_dict={z_vector:z_sample}) # get a generated image, with range [-1, 1]\n",
    "            \n",
    "            # make a directory for generated images\n",
    "            if not os.path.exists(train_sample_directory):\n",
    "                os.makedirs(train_sample_directory)\n",
    "            \n",
    "            #Save sample generator images for viewing training progress.\n",
    "            save_generated_images(images = np.reshape(g_images[0:batch_size_sample],[batch_size_sample,64,64,3]),\\\n",
    "                        size = [sample_num_rows,sample_num_columns], image_path = train_sample_directory+'/'+str(epoch)+'.png')\n",
    "            \n",
    "        if epoch in [16500,17000]:\n",
    "            # make a directory for trained models\n",
    "            if not os.path.exists(model_directory):\n",
    "                os.makedirs(model_directory)\n",
    "            \n",
    "            # save the trained model at different epoch\n",
    "            saver.save(sess, save_path = model_directory + '/' + str(epoch) + '.cptk')\n",
    "    print \"Done\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Combine sample images generated from training and save it as GIF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "file_names = sorted([int(f[:-4]) for f in os.listdir(train_sample_directory) if f.endswith('.png')]) # get image filenames\n",
    "file_names = [str(f)+'.png' for f in file_names] # add the format suffix at the end of the filenames\n",
    "images = [Image.open(train_sample_directory+f) for f in file_names] # turn into image instances\n",
    "filename = \"train_samples.gif\"\n",
    "\n",
    "if not os.path.exists(train_sample_directory):\n",
    "            os.makedirs(train_sample_directory)\n",
    "writeGif(train_sample_directory+filename, images, duration=0.1) # combine images to gif"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Use the trained generator to generate images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "\n",
    "print 'Loading models...might take a minute'\n",
    "saver = tf.train.Saver(max_to_keep=10)\n",
    "\n",
    "with tf.Session() as sess:\n",
    "    #sess.run(tf.initialize_all_variables())\n",
    "    \n",
    "    # reload the trained model.\n",
    "    ckpt = tf.train.get_checkpoint_state(model_directory)\n",
    "    models = ckpt.all_model_checkpoint_paths\n",
    "    #print 'all saved models: ',ckpt.all_model_checkpoint_paths\n",
    "    \n",
    "    # z noise vector that will be used to generate image to check the training progress\n",
    "    z_sample = np.random.uniform(-1.0,1.0,size=[batch_size,z_size]).astype(np.float32)\n",
    "    \n",
    "    # loop through each checkpoint and use models to generate images at different training stages\n",
    "    for index, model in enumerate(models):\n",
    "        saver.restore(sess, save_path=model)\n",
    "\n",
    "        g_images = sess.run(g_output,feed_dict={z_vector:z_sample}) # get a generated image\n",
    "\n",
    "        # make a directory for generated images\n",
    "        if not os.path.exists(sample_directory):\n",
    "            os.makedirs(sample_directory)\n",
    "\n",
    "        #Save generated sample images\n",
    "        save_generated_images(images = np.reshape(g_images[0:batch_size_sample],[batch_size_sample,64,64,3]),\\\n",
    "                    size = [sample_num_rows,sample_num_columns], image_path = sample_directory+'/'+str(index)+'.png')\n",
    "\n",
    "        print \"Model \"+str(index),': ',model\n",
    "        \n",
    "time_lapse = time.time()-start\n",
    "print \" Time spent: %.2fs\" % time_lapse"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Turn generated images from checkpoints into an animated GIF"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p>Restart the kernel if not working..</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "file_names = sorted([int(f[:-4]) for f in os.listdir('.') if f.endswith('.png')]) # get sorted image filenames \n",
    "file_names = [str(f)+'.png' for f in file_names] # add the format suffix at the end of the filenames\n",
    "images = [Image.open(f) for f in file_names] # turn into image instances\n",
    "filename = \"train_samples.gif\"\n",
    "writeGif(filename, images, duration=0.1) # combine images to gif"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Interpolation Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "model_directory = './dcgan_face_15k/models/' # directory to save trained model\n",
    "sample_inter_directory = './dcgan_face_15k/interpolation/' # directory to save the generated images from linear interpolation\n",
    "# make a directory for generated images\n",
    "if not os.path.exists(sample_inter_directory):\n",
    "    os.makedirs(sample_inter_directory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# spherical interpolation between pointA and pointB, each is a ndarray\n",
    "# val is a value between [0,1], where 0 returns pointA, 1 returns pointB\n",
    "def slerp(val, pointA, pointB):\n",
    "    omega = np.arccos(np.clip(np.dot(pointA/np.linalg.norm(pointA), pointB/np.linalg.norm(pointB)), -1, 1))\n",
    "    so = np.sin(omega)\n",
    "    if so == 0:\n",
    "        return (1.0-val) * pointA + val * pointB # L'Hopital's rule/LERP\n",
    "    return np.sin((1.0-val)*omega) / so * pointA + np.sin(val*omega) / so * pointB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading Face models...might take a minute\n",
      "done.\n"
     ]
    }
   ],
   "source": [
    "tf.reset_default_graph()\n",
    "# size of initial noise vector that will be used for generator\n",
    "z_size = 100 \n",
    "\n",
    "# placeholders for inputs into the generator\n",
    "# shape[0] is None means it could take any integer value\n",
    "z_vector = tf.placeholder(shape=[None,z_size],dtype=tf.float32, name='z_vectors')\n",
    "whether_is_train = tf.placeholder(tf.bool) # boolean variable indicating if it is in training mode or not\n",
    "\n",
    "# initialize all parameters of the networks\n",
    "# weights were initialized from a zero-centered Normal distribution with standard deviation 0.02\n",
    "# tf.truncated_normal returns random values from a normal distribution and made sure no value exceeds 2 std\n",
    "initializer = tf.truncated_normal_initializer(stddev=0.02)\n",
    "\n",
    "# generated mini-batch of images from noisy z vectors\n",
    "# when whether_is_train==False, then it is in training mode, otherwise, it is not in training mode\n",
    "# lambda is used to allow injecting function parameters for generator()\n",
    "g_output = tf.cond(whether_is_train, \n",
    "                   lambda: generator(z_vector,is_train=True), \n",
    "                   lambda: generator(z_vector,reuse=True,is_train=False))\n",
    "\n",
    "print 'Loading Face models...might take a minute'\n",
    "saver = tf.train.Saver()\n",
    "print 'done.'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interpolation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "most trained model:  ./dcgan_face_15k/models/16500.cptk\n",
      "spherial interpolation\n",
      "done\n"
     ]
    }
   ],
   "source": [
    "with tf.Session() as sess:\n",
    "\n",
    "    # reload the trained model.\n",
    "    ckpt = tf.train.get_checkpoint_state(model_directory)\n",
    "    model = ckpt.model_checkpoint_path # just read the model at epoch 16500\n",
    "    print 'most trained model: ', model\n",
    "\n",
    "    # z noise vector that will be used to generate image to check the training progress\n",
    "    z_sample = np.random.uniform(-1.0,1.0,size=[batch_size_sample,z_size]).astype(np.float32)\n",
    "\n",
    "    # restore model variables\n",
    "    saver.restore(sess, save_path=model)\n",
    "\n",
    "    # linear interpolation\n",
    "    #z_sample_one = z_sample[:2]\n",
    "    z_sample_one = z_sample[:1]\n",
    "    #print z_sample_one.shape\n",
    "\n",
    "    # set fixed input, interpolate from 7 to 9\n",
    "    sample1='''0.3904759   0.12369447 -0.03130473 -0.77256024  0.54540509 -0.63576019\n",
    "   0.64418173  0.0961884  -0.30221939  0.7397421  -0.41891637  0.23050123\n",
    "  -0.85202283  0.09574939  0.49240941  0.84485519 -0.4560945  -0.73232478\n",
    "   0.88528788 -0.49584585 -0.05611671 -0.94211501  0.80680168 -0.62426782\n",
    "  -0.18791643  0.67589283 -0.06083585 -0.19920608 -0.8094129  -0.3279312\n",
    "   0.99031818  0.72150308 -0.66078311 -0.81424487 -0.60490525 -0.71164876\n",
    "  -0.0634105   0.02162064  0.56584495  0.09616423  0.51447213  0.09199435\n",
    "  -0.5390138   0.02217139 -0.90989834  0.64157981  0.17594992 -0.42963964\n",
    "  -0.59015018  0.83339399  0.2526193  -0.78287238  0.50538272 -0.51408827\n",
    "   0.9690786  -0.90271968  0.93069273 -0.86799151 -0.86437082 -0.39811093\n",
    "   0.06549309 -0.84584093  0.52733558 -0.90270567  0.41961735  0.62828648\n",
    "  -0.81002378  0.53971922  0.14811444 -0.93527913 -0.51290333 -0.48542583\n",
    "   0.74688935  0.4627139   0.54627806 -0.14855324  0.5601418   0.24667268\n",
    "  -0.413488    0.46806765  0.78500813  0.64971685  0.63234931 -0.28815132\n",
    "  -0.42552331  0.40306959 -0.71005774 -0.44297776 -0.76836979  0.24986272\n",
    "  -0.02224439  0.65198964 -0.97852767  0.46914703 -0.8889237  -0.39507776\n",
    "  -0.93041956 -0.88002688 -0.24425523  0.41087016'''\n",
    "    sample2=''' 0.54349154  0.82320267 -0.62965882  0.98068064 -0.56107283 -0.31114981\n",
    "  -0.45138538  0.80631655  0.90001726 -0.9619692   0.93649179  0.67718279\n",
    "   0.13684638 -0.15676339 -0.3471978  -0.83544356  0.24256505 -0.42204261\n",
    "  -0.4512648   0.16144913 -0.82650888  0.07874123 -0.49977717 -0.1294433\n",
    "  -0.8050105   0.44765326  0.39748943  0.84955609 -0.97319156  0.3089633\n",
    "  -0.34072471  0.14846873 -0.08748962  0.14348175  0.03980778  0.03936955\n",
    "   0.11320399 -0.65704525 -0.23477662 -0.94302213  0.45144475  0.05105386\n",
    "  -0.11944827  0.64141428 -0.9440648  -0.7251038  -0.04335017  0.45072994\n",
    "  -0.062515   -0.97800612  0.25549379  0.43639576 -0.00373689 -0.48526773\n",
    "  -0.85005528 -0.71420103  0.16092534 -0.34909213 -0.72042316 -0.85552001\n",
    "  -0.0950563   0.64399147 -0.14060096 -0.29917473  0.39624763 -0.94713563\n",
    "  -0.63286591  0.49637678 -0.2596755  -0.06869141 -0.46832797 -0.86966354\n",
    "  -0.87580574 -0.08289656 -0.35145256 -0.47775638  0.22160572 -0.92910951\n",
    "  -0.05955737 -0.92531306  0.65753967  0.195959   -0.3773948  -0.6927458\n",
    "   0.47357163  0.39965707  0.75121379 -0.83463138  0.92006731  0.96232468\n",
    "   0.14594369 -0.20597722  0.16611969 -0.56408048  0.3470183  -0.48169822\n",
    "  -0.74719411 -0.9993121  -0.33036932  0.7356177 '''\n",
    "    \n",
    "    # turn the string above into an array. sample1 represents a point in a 100-dimensional space\n",
    "    sample1 = np.array([float(i.rstrip()) for i in sample1.split(' ') if i != ''])\n",
    "    #sample1 = np.array(z_sample[0])\n",
    "    sample2 = np.array([float(i.rstrip()) for i in sample2.split(' ') if i != ''])\n",
    "    #sample2 = np.array(z_sample[1])\n",
    "    #print z_sample[:2]\n",
    "    #z_sample_one[1] = np.zeros(100) # keep the second face constant\n",
    "    \n",
    "    num_inter = 20 # number of interpolated images\n",
    "    '''\n",
    "    print 'one-dimension interpolation'\n",
    "    dim = 18 # only change values in this dimension\n",
    "    \n",
    "    \n",
    "    # transition from one face to the face in the middle, then to another face, by changing value in one dimension\n",
    "    for i, value in enumerate(np.linspace(-10,10,num_inter)): \n",
    "    # transition from the face in the middle to another face\n",
    "    #for i, value in enumerate(np.linspace(0,20,num_inter)):  \n",
    "        z_sample_one[0][dim] = value \n",
    "        #print z_sample_one\n",
    "        \n",
    "        # whether_is_train should be false, indicating this is not training mode, the batch norm should have\n",
    "        # parameter: is_training = False. Only with this setting the weights are not normalized when generating images\n",
    "        g_images = sess.run(g_output,feed_dict={z_vector:z_sample_one, whether_is_train:False})  \n",
    "        g_image = g_images[:1] # only retrieve the first returned image\n",
    "\n",
    "        # save interpolated image\n",
    "        save_generated_images(images = g_image,\\\n",
    "            size = [1,1], image_path = sample_inter_directory+'/'+str(i)+'.png')\n",
    "    '''\n",
    "    \n",
    "    '''\n",
    "    print 'linear interpolation'\n",
    "    # linear interpolation between sample1 and sample2            \n",
    "    for i in range(num_inter):\n",
    "        if i == 0: # interpolation start with point A\n",
    "            z_sample_one[0] = sample1\n",
    "        elif i == num_inter-1: # interpolation end with point B\n",
    "            z_sample_one[0] = sample2\n",
    "        else:\n",
    "            z_sample_one[0] += (sample2-sample1)/num_inter # linear interpolate from point A to B\n",
    "        #print z_sample_one,'\\n'\n",
    "\n",
    "        g_images = sess.run(g_output,feed_dict={z_vector:z_sample_one, whether_is_train:False})\n",
    "        g_image = g_images[:1] # only retrieve the first returned image\n",
    "        #print g_image.shape\n",
    "            \n",
    "        # save interpolated image\n",
    "        save_generated_images(images = g_image,\\\n",
    "            size = [1,1], image_path = sample_inter_directory+'/'+str(i)+'.png')\n",
    "    '''\n",
    "    \n",
    "    \n",
    "    print 'spherial interpolation'\n",
    "    # spherical interpolation between pointA and pointB, each is a ndarray\n",
    "    # val is a value between [0,1], where 0 returns pointA, 1 returns pointB\n",
    "    for i, value in enumerate(np.linspace(0,1,num_inter)):\n",
    "        z_sample_one = slerp(val=value, pointA=sample1, pointB=sample2)\n",
    "        z_sample_one = np.array([z_sample_one])\n",
    "        #print z_sample_one,'\\n'\n",
    "\n",
    "        g_images = sess.run(g_output,feed_dict={z_vector:z_sample_one, whether_is_train:False})\n",
    "        g_image = g_images[:1] # only retrieve the first returned image\n",
    "\n",
    "        # save interpolated image\n",
    "        save_generated_images(images = g_image,\\\n",
    "            size = [1,1], image_path = sample_inter_directory+'/'+str(i)+'.png')\n",
    "    \n",
    "    print 'done'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## turn the interpolate images into GIF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "file_names = sorted([int(f[:-4]) for f in os.listdir(sample_inter_directory) if f.endswith('.png')]) # get image filenames\n",
    "file_names = [str(f)+'.png' for f in file_names] # add the format suffix at the end of the filenames\n",
    "images = [Image.open(sample_inter_directory+f) for f in file_names] # turn into image instances\n",
    "filename = \"aa.gif\"\n",
    "writeGif(sample_inter_directory+filename, images, duration=0.1) # combine images to gif"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}