GAN_MNIST_modify.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, time, itertools, imageio, pickle\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "from tensorflow.examples.tutorials.mnist import input_data\n",
    "\n",
    "# G(z)\n",
    "def generator(x):\n",
    "    # initializers\n",
    "    w_init = tf.truncated_normal_initializer(mean=0, stddev=0.02)\n",
    "    b_init = tf.constant_initializer(0.)\n",
    "\n",
    "    # 1st hidden layer\n",
    "    w0 = tf.get_variable('G_w0', [x.get_shape()[1], 256], initializer=w_init)\n",
    "    b0 = tf.get_variable('G_b0', [256], initializer=b_init)\n",
    "    h0 = tf.nn.relu(tf.matmul(x, w0) + b0)\n",
    "\n",
    "    # 2nd hidden layer\n",
    "    w1 = tf.get_variable('G_w1', [h0.get_shape()[1], 512], initializer=w_init)\n",
    "    b1 = tf.get_variable('G_b1', [512], initializer=b_init)\n",
    "    h1 = tf.nn.relu(tf.matmul(h0, w1) + b1)\n",
    "\n",
    "    # 3rd hidden layer\n",
    "    w2 = tf.get_variable('G_w2', [h1.get_shape()[1], 1024], initializer=w_init)\n",
    "    b2 = tf.get_variable('G_b2', [1024], initializer=b_init)\n",
    "    h2 = tf.nn.relu(tf.matmul(h1, w2) + b2)\n",
    "\n",
    "    # output hidden layer\n",
    "    w3 = tf.get_variable('G_w3', [h2.get_shape()[1], 784], initializer=w_init)\n",
    "    b3 = tf.get_variable('G_b3', [784], initializer=b_init)\n",
    "    o = tf.nn.tanh(tf.matmul(h2, w3) + b3)\n",
    "\n",
    "    return o\n",
    "\n",
    "# D(x)\n",
    "def discriminator(x, drop_out):\n",
    "\n",
    "    # initializers\n",
    "    w_init = tf.truncated_normal_initializer(mean=0, stddev=0.02)\n",
    "    b_init = tf.constant_initializer(0.)\n",
    "\n",
    "    # 1st hidden layer\n",
    "    w0 = tf.get_variable('D_w0', [x.get_shape()[1], 1024], initializer=w_init)\n",
    "    b0 = tf.get_variable('D_b0', [1024], initializer=b_init)\n",
    "    h0 = tf.nn.relu(tf.matmul(x, w0) + b0)\n",
    "    h0 = tf.nn.dropout(h0, drop_out)\n",
    "\n",
    "    # 2nd hidden layer\n",
    "    w1 = tf.get_variable('D_w1', [h0.get_shape()[1], 512], initializer=w_init)\n",
    "    b1 = tf.get_variable('D_b1', [512], initializer=b_init)\n",
    "    h1 = tf.nn.relu(tf.matmul(h0, w1) + b1)\n",
    "    h1 = tf.nn.dropout(h1, drop_out)\n",
    "\n",
    "    # 3rd hidden layer\n",
    "    w2 = tf.get_variable('D_w2', [h1.get_shape()[1], 256], initializer=w_init)\n",
    "    b2 = tf.get_variable('D_b2', [256], initializer=b_init)\n",
    "    h2 = tf.nn.relu(tf.matmul(h1, w2) + b2)\n",
    "    h2 = tf.nn.dropout(h2, drop_out)\n",
    "\n",
    "    # output layer\n",
    "    w3 = tf.get_variable('D_w3', [h2.get_shape()[1], 1], initializer=w_init)\n",
    "    b3 = tf.get_variable('D_b3', [1], initializer=b_init)\n",
    "    o = tf.sigmoid(tf.matmul(h2, w3) + b3)\n",
    "\n",
    "    return o\n",
    "\n",
    "fixed_z_ = np.random.normal(0, 1, (25, 100))\n",
    "def show_result(num_epoch, show = False, save = False, path = 'result.png', isFix=False):\n",
    "    z_ = np.random.normal(0, 1, (25, 100))\n",
    "\n",
    "    if isFix:\n",
    "        test_images = sess.run(G_z, {z: fixed_z_, drop_out: 0.0})\n",
    "    else:\n",
    "        test_images = sess.run(G_z, {z: z_, drop_out: 0.0})\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(5*5):\n",
    "        i = k // 5\n",
    "        j = k % 5\n",
    "        ax[i, j].cla()\n",
    "        ax[i, j].imshow(np.reshape(test_images[k], (28, 28)), cmap='gray')\n",
    "\n",
    "    label = 'Epoch {0}'.format(num_epoch)\n",
    "    fig.text(0.5, 0.04, label, ha='center')\n",
    "    plt.savefig(path)\n",
    "\n",
    "    if show:\n",
    "        plt.show()\n",
    "    else:\n",
    "        plt.close()\n",
    "\n",
    "def show_train_hist(hist, show = False, save = False, path = 'Train_hist.png'):\n",
    "    x = range(len(hist['D_losses']))\n",
    "\n",
    "    y1 = hist['D_losses']\n",
    "    y2 = hist['G_losses']\n",
    "\n",
    "    plt.plot(x, y1, label='D_loss')\n",
    "    plt.plot(x, y2, label='G_loss')\n",
    "\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Loss')\n",
    "\n",
    "    plt.legend(loc=4)\n",
    "    plt.grid(True)\n",
    "    plt.tight_layout()\n",
    "\n",
    "    if save:\n",
    "        plt.savefig(path)\n",
    "\n",
    "    if show:\n",
    "        plt.show()\n",
    "    else:\n",
    "        plt.close()\n",
    "\n",
    "# training parameters\n",
    "batch_size = 100\n",
    "lr = 0.0002\n",
    "train_epoch = 200\n",
    "\n",
    "# load MNIST\n",
    "mnist = input_data.read_data_sets(\"MNIST_data/\", one_hot=True)\n",
    "train_set = (mnist.train.images - 0.5) / 0.5  # normalization; range: -1 ~ 1\n",
    "\n",
    "# networks : generator\n",
    "with tf.variable_scope('G'):\n",
    "    z = tf.placeholder(tf.float32, shape=(None, 100))\n",
    "    G_z = generator(z)\n",
    "\n",
    "# networks : discriminator\n",
    "with tf.variable_scope('D') as scope:\n",
    "    drop_out = tf.placeholder(dtype=tf.float32, name='drop_out')\n",
    "    x = tf.placeholder(tf.float32, shape=(None, 784))\n",
    "    D_real = discriminator(x, drop_out)\n",
    "    scope.reuse_variables()\n",
    "    D_fake = discriminator(G_z, drop_out)\n",
    "\n",
    "\n",
    "# loss for each network\n",
    "eps = 1e-2\n",
    "D_loss = tf.reduce_mean(-tf.log(D_real + eps) - tf.log(1 - D_fake + eps))\n",
    "G_loss = tf.reduce_mean(-tf.log(D_fake + eps))\n",
    "\n",
    "# trainable variables for each network\n",
    "t_vars = tf.trainable_variables()\n",
    "D_vars = [var for var in t_vars if 'D_' in var.name]\n",
    "G_vars = [var for var in t_vars if 'G_' in var.name]\n",
    "\n",
    "# optimizer for each network\n",
    "D_optim = tf.train.AdamOptimizer(lr).minimize(D_loss, var_list=D_vars)\n",
    "G_optim = tf.train.AdamOptimizer(lr).minimize(G_loss, var_list=G_vars)\n",
    "\n",
    "# open session and initialize all variables\n",
    "sess = tf.InteractiveSession()\n",
    "tf.global_variables_initializer().run()\n",
    "\n",
    "# results save folder\n",
    "if not os.path.isdir('MNIST_GAN_results'):\n",
    "    os.mkdir('MNIST_GAN_results')\n",
    "if not os.path.isdir('MNIST_GAN_results/Random_results'):\n",
    "    os.mkdir('MNIST_GAN_results/Random_results')\n",
    "if not os.path.isdir('MNIST_GAN_results/Fixed_results'):\n",
    "    os.mkdir('MNIST_GAN_results/Fixed_results')\n",
    "\n",
    "train_hist = {}\n",
    "train_hist['D_losses'] = []\n",
    "train_hist['G_losses'] = []\n",
    "train_hist['per_epoch_ptimes'] = []\n",
    "train_hist['total_ptime'] = []\n",
    "\n",
    "# training-loop\n",
    "np.random.seed(int(time.time()))\n",
    "start_time = time.time()\n",
    "for epoch in range(train_epoch):\n",
    "    G_losses = []\n",
    "    D_losses = []\n",
    "    epoch_start_time = time.time()\n",
    "    for iter in range(train_set.shape[0] // batch_size):\n",
    "        # update discriminator\n",
    "        x_ = train_set[iter*batch_size:(iter+1)*batch_size]\n",
    "        z_ = np.random.normal(0, 1, (batch_size, 100))\n",
    "\n",
    "        loss_d_, _ = sess.run([D_loss, D_optim], {x: x_, z: z_, drop_out: 0.3})\n",
    "        D_losses.append(loss_d_)\n",
    "\n",
    "        # update generator\n",
    "        z_ = np.random.normal(0, 1, (batch_size, 100))\n",
    "        loss_g_, _ = sess.run([G_loss, G_optim], {z: z_, drop_out: 0.3})\n",
    "        G_losses.append(loss_g_)\n",
    "\n",
    "    epoch_end_time = time.time()\n",
    "    per_epoch_ptime = epoch_end_time - epoch_start_time\n",
    "    print('[%d/%d] - ptime: %.2f loss_d: %.3f, loss_g: %.3f' % ((epoch + 1), train_epoch, per_epoch_ptime, np.mean(D_losses), np.mean(G_losses)))\n",
    "    p = 'MNIST_GAN_results/Random_results/MNIST_GAN_' + str(epoch + 1) + '.png'\n",
    "    fixed_p = 'MNIST_GAN_results/Fixed_results/MNIST_GAN_' + str(epoch + 1) + '.png'\n",
    "    show_result((epoch + 1), save=True, path=p, isFix=False)\n",
    "    show_result((epoch + 1), save=True, path=fixed_p, isFix=True)\n",
    "    train_hist['D_losses'].append(np.mean(D_losses))\n",
    "    train_hist['G_losses'].append(np.mean(G_losses))\n",
    "    train_hist['per_epoch_ptimes'].append(per_epoch_ptime)\n",
    "\n",
    "end_time = time.time()\n",
    "total_ptime = end_time - start_time\n",
    "train_hist['total_ptime'].append(total_ptime)\n",
    "\n",
    "print('Avg per epoch ptime: %.2f, total %d epochs ptime: %.2f' % (np.mean(train_hist['per_epoch_ptimes']), train_epoch, total_ptime))\n",
    "print(\"Training finish!... save training results\")\n",
    "with open('MNIST_GAN_results/train_hist.pkl', 'wb') as f:\n",
    "    pickle.dump(train_hist, f)\n",
    "show_train_hist(train_hist, save=True, path='MNIST_GAN_results/MNIST_GAN_train_hist.png')\n",
    "\n",
    "images = []\n",
    "for e in range(train_epoch):\n",
    "    img_name = 'MNIST_GAN_results/Fixed_results/MNIST_GAN_' + str(e + 1) + '.png'\n",
    "    images.append(imageio.imread(img_name))\n",
    "imageio.mimsave('MNIST_GAN_results/generation_animation.gif', images, fps=5)\n",
    "\n",
    "sess.close()"
   ]
  },
  {
   "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
}