WGAN-GP_MNIST.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow.examples.tutorials.mnist import input_data\n",
    "import os\n",
    "import time\n",
    "import matplotlib.pyplot as plt\n",
    "import itertools, imageio, pickle\n",
    "tf.logging.set_verbosity(tf.logging.ERROR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist = input_data.read_data_sets('./MNIST_data', one_hot=True, reshape=[])\n",
    "\n",
    "class DataSampler(object):\n",
    "    def __init__(self, dataset):\n",
    "        self.dataset = dataset\n",
    "        self.num_sample = dataset.shape[0]\n",
    "        self.index_in_epoch = 0\n",
    "        self.epochs_completed = 0\n",
    "        \n",
    "    def __call__(self, batch_size,shuffle=True):\n",
    "        start = self.index_in_epoch\n",
    "        if self.epochs_completed == 0 and start == 0 and shuffle:\n",
    "            perm = np.arange(self.num_sample)\n",
    "            np.random.shuffle(perm)\n",
    "            self.dataset = self.dataset[perm]\n",
    "        \n",
    "        if start + batch_size > self.num_sample:\n",
    "            self.epochs_completed += 1\n",
    "            rest_num_examples = self.num_sample - start\n",
    "            data_rest_part = self.dataset[start:self.num_sample]\n",
    "            \n",
    "            if shuffle:\n",
    "                prem = np.arange(self.num_sample)\n",
    "                np.random.shuffle(prem)\n",
    "                self.dataset = self.dataset[prem]\n",
    "                \n",
    "            start = 0\n",
    "            self.index_in_epoch = batch_size - rest_num_examples\n",
    "            end = self.index_in_epoch\n",
    "            data_new_part = self.dataset[start:end]\n",
    "            return np.concatenate((data_rest_part,data_new_part),axis=0 )\n",
    "        else:\n",
    "            self.index_in_epoch += batch_size\n",
    "            end = self.index_in_epoch\n",
    "            return self.dataset[start:end]\n",
    "        \n",
    "class NoiseSampler(object):\n",
    "    def __call__(self, batch_size, z_dim):\n",
    "        return np.random.normal(0., 1.0, [batch_size,1,1,z_dim])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lrelu(x, th=0.2):\n",
    "    return tf.maximum(th * x, x)\n",
    "\n",
    "def layer_norm(x, trainable=True, name='layer_norm'):\n",
    "    with tf.variable_scope(name):\n",
    "        return tf.contrib.layers.layer_norm(x, trainable=trainable)\n",
    "\n",
    "# G(z)\n",
    "def generator(x, isTrain=True, reuse=False):\n",
    "    with tf.variable_scope('generator', reuse=reuse):\n",
    "\n",
    "        # 1st hidden layer\n",
    "        conv1 = tf.layers.conv2d_transpose(x, 1024, [4, 4], strides=(1, 1), padding='valid')\n",
    "        lrelu1 = lrelu(tf.layers.batch_normalization(conv1, training=isTrain), 0.2)\n",
    "\n",
    "        # 2nd hidden layer\n",
    "        conv2 = tf.layers.conv2d_transpose(lrelu1, 512, [4, 4], strides=(2, 2), padding='same')\n",
    "        lrelu2 = lrelu(tf.layers.batch_normalization(conv2, training=isTrain), 0.2)\n",
    "\n",
    "        # 3rd hidden layer\n",
    "        conv3 = tf.layers.conv2d_transpose(lrelu2, 256, [4, 4], strides=(2, 2), padding='same')\n",
    "        lrelu3 = lrelu(tf.layers.batch_normalization(conv3, training=isTrain), 0.2)\n",
    "\n",
    "        # 4th hidden layer\n",
    "        conv4 = tf.layers.conv2d_transpose(lrelu3, 128, [4, 4], strides=(2, 2), padding='same')\n",
    "        lrelu4 = lrelu(tf.layers.batch_normalization(conv4, training=isTrain), 0.2)\n",
    "\n",
    "        # output layer\n",
    "        conv5 = tf.layers.conv2d_transpose(lrelu4, 1, [4, 4], strides=(2, 2), padding='same')\n",
    "        o = tf.nn.tanh(conv5)\n",
    "\n",
    "        return o\n",
    "    \n",
    "def discriminator(x, reuse=False):\n",
    "    with tf.variable_scope('discriminator', reuse=reuse):\n",
    "        # 1st hidden layer\n",
    "        conv1 = tf.layers.conv2d(x, 128, [4, 4], strides=(2, 2), padding='same')\n",
    "        lrelu1 = lrelu(conv1, 0.2)\n",
    "\n",
    "        # 2nd hidden layer\n",
    "        conv2 = tf.layers.conv2d(lrelu1, 256, [4, 4], strides=(2, 2), padding='same')\n",
    "        lrelu2 = lrelu(layer_norm(conv2,name='ln1'), 0.2)\n",
    "\n",
    "        # 3rd hidden layer\n",
    "        conv3 = tf.layers.conv2d(lrelu2, 512, [4, 4], strides=(2, 2), padding='same')\n",
    "        lrelu3 = lrelu(layer_norm(conv3,name='ln2'), 0.2)\n",
    "\n",
    "        # 4th hidden layer\n",
    "        conv4 = tf.layers.conv2d(lrelu3, 1024, [4, 4], strides=(2, 2), padding='same')\n",
    "        lrelu4 = lrelu(layer_norm(conv4,name='ln3'), 0.2)\n",
    "\n",
    "        # output layer\n",
    "        conv5 = tf.layers.conv2d(lrelu4, 1, [4, 4], strides=(1, 1), padding='valid')\n",
    "        \n",
    "        return conv5    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fixed_z_ = np.random.normal(0, 1, (25, 1, 1, 100))\n",
    "def show_result(num_epoch, show = False, save = False, path = 'result.png'):\n",
    "    test_images = sess.run(Gz, {z: fixed_z_, isTrain: False})\n",
    "\n",
    "    size_figure_grid = 5\n",
    "    fig, ax = plt.subplots(size_figure_grid, size_figure_grid, figsize=(5, 5))\n",
    "    for i, j in itertools.product(range(size_figure_grid), range(size_figure_grid)):\n",
    "        ax[i, j].get_xaxis().set_visible(False)\n",
    "        ax[i, j].get_yaxis().set_visible(False)\n",
    "\n",
    "    for k in range(size_figure_grid*size_figure_grid):\n",
    "        i = k // size_figure_grid\n",
    "        j = k % size_figure_grid\n",
    "        ax[i, j].cla()\n",
    "        ax[i, j].imshow(np.reshape(test_images[k], (64, 64)), cmap='gray')\n",
    "\n",
    "    label = 'Epoch {0}'.format(num_epoch)\n",
    "    fig.text(0.5, 0.04, label, ha='center')\n",
    "\n",
    "    if save:\n",
    "        plt.savefig(path)\n",
    "\n",
    "    if show:\n",
    "        plt.show()\n",
    "    else:\n",
    "        plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tf.placeholder(tf.float32,[None, 64,64,1])\n",
    "z = tf.placeholder(tf.float32, shape=(None, 1, 1, 100))\n",
    "isTrain = tf.placeholder(dtype=tf.bool)\n",
    "\n",
    "Gz = generator(z, isTrain)\n",
    "\n",
    "D_real = discriminator(x)\n",
    "D_fake = discriminator(Gz,reuse=True)\n",
    "\n",
    "g_loss = tf.reduce_mean(D_fake)\n",
    "d_loss = tf.reduce_mean(D_real) - tf.reduce_mean(D_fake)\n",
    "\n",
    "epsilon = tf.random_uniform([],0.,1.0)\n",
    "x_hat = epsilon*x + (1 - epsilon)*Gz\n",
    "\n",
    "d_hat = discriminator(x_hat, reuse=True)\n",
    "\n",
    "ddx = tf.gradients(d_hat,x_hat)[0]\n",
    "ddx = tf.sqrt(tf.reduce_sum(tf.square(ddx), axis=1))\n",
    "ddx = tf.reduce_mean(tf.square(ddx - 1.0)*10.)\n",
    "\n",
    "d_loss = d_loss + ddx\n",
    "\n",
    "# trainable variables for each network\n",
    "T_vars = tf.trainable_variables()\n",
    "D_vars = [var for var in T_vars if var.name.startswith('discriminator')]\n",
    "print(D_vars)\n",
    "print(\"=========================================\")\n",
    "G_vars = [var for var in T_vars if var.name.startswith('generator')]\n",
    "print(G_vars)\n",
    "\n",
    "# 此处应该加tf.control_dependencies来同步bacth_norm,后续研究\n",
    "with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS,scope='discriminator')):\n",
    "    D_opt = tf.train.AdamOptimizer(learning_rate=1e-4,beta1=0.,beta2=0.9).minimize(d_loss,var_list=D_vars)\n",
    "with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS,scope='generator')):\n",
    "    G_opt = tf.train.AdamOptimizer(learning_rate=1e-4,beta1=0.,beta2=0.9).minimize(g_loss,var_list=G_vars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpu_options = tf.GPUOptions(allow_growth=True)\n",
    "sess = tf.InteractiveSession(config=tf.ConfigProto(gpu_options=gpu_options))\n",
    "\n",
    "tf.global_variables_initializer().run()\n",
    "\n",
    "x_train= tf.image.resize_images(mnist.train.images, [64, 64]).eval()\n",
    "x_train =(x_train - 0.5)/0.5\n",
    "\n",
    "x_sampler = DataSampler(x_train)\n",
    "z_sampler = NoiseSampler()\n",
    "\n",
    "\n",
    "# results save folder\n",
    "root = 'MNIST_WDCGAN-GP_results/'\n",
    "model = 'MNIST_WDCGAN-GP_'\n",
    "if not os.path.isdir(root):\n",
    "    os.mkdir(root)\n",
    "if not os.path.isdir(root + 'Fixed_results'):\n",
    "    os.mkdir(root + 'Fixed_results')\n",
    "    \n",
    "# training-loop\n",
    "np.random.seed(int(time.time()))\n",
    "print(\"training start!\")\n",
    "start_time = time.time()\n",
    "\n",
    "batch_size = 64\n",
    "num_iters = 100000\n",
    "for iter in range(0,num_iters):\n",
    "    ncritic = 5\n",
    "    if iter % 500 == 0 or iter < 25:\n",
    "        ncritic = 100\n",
    "    for _ in range(0, ncritic):\n",
    "        x_batch = x_sampler(batch_size)\n",
    "        z_batch = z_sampler(batch_size,100)\n",
    "        sess.run(D_opt, feed_dict={x:x_batch,z:z_batch, isTrain:True})\n",
    "        \n",
    "        \n",
    "    z_batch = z_sampler(batch_size, 100)\n",
    "    sess.run(G_opt, feed_dict = {x:x_batch, z:z_batch, isTrain:True})\n",
    "    \n",
    "    if iter % 100 ==0:\n",
    "        D_loss = sess.run(d_loss,feed_dict={x:x_batch,z:z_batch,isTrain:True})\n",
    "        G_loss = sess.run(g_loss,feed_dict={x:x_batch,z:z_batch,isTrain:True})\n",
    "        print(\"Iter: {},d_loss: {}, g_loss: {}\".format(iter,D_loss,G_loss))\n",
    "    if iter % 500 == 0:\n",
    "        fixed_p = root + 'Fixed_results/' + model + str(iter) + '.png'\n",
    "        show_result((iter + 1), save=True, path=fixed_p)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}