CGAN-CIFAR10.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true,
    "pycharm": {}
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "# tfe = tf.contrib.eager\n",
    "# tf.enable_eager_execution()\n",
    "\n",
    "import os, time, itertools, imageio, pickle\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from tensorflow.examples.tutorials.mnist import input_data\n",
    "import numpy as np\n",
    "import matplotlib.gridspec as gridspec\n",
    "import platform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#load dataset\n",
    "def load_pickle(f):\n",
    "    version = platform.python_version_tuple()\n",
    "    if version[0] == '2':\n",
    "        return  pickle.load(f)\n",
    "    elif version[0] == '3':\n",
    "        return  pickle.load(f, encoding='latin1')\n",
    "    raise ValueError(\"invalid python version: {}\".format(version))\n",
    "    \n",
    "def load_CIFAR_batch(filename):\n",
    "    \"\"\" load single batch of cifar \"\"\"\n",
    "    with open(filename, 'rb') as f:\n",
    "        datadict = load_pickle(f)\n",
    "        X = datadict['data']\n",
    "        Y = datadict['labels']\n",
    "        X = X.reshape(10000, 3, 32, 32).transpose(0,2,3,1).astype(\"float32\")\n",
    "        Y = np.array(Y)\n",
    "        return X, Y\n",
    "\n",
    "def load_CIFAR10(ROOT):\n",
    "    \"\"\" load all of cifar \"\"\"\n",
    "    xs = []\n",
    "    ys = []\n",
    "    for b in range(1,6):\n",
    "        f = os.path.join(ROOT, 'data_batch_%d' % (b, ))\n",
    "        X, Y = load_CIFAR_batch(f)\n",
    "        xs.append(X)\n",
    "        ys.append(Y)    \n",
    "    Xtr = np.concatenate(xs)\n",
    "    Ytr = np.concatenate(ys)\n",
    "    del X, Y\n",
    "    Xte, Yte = load_CIFAR_batch(os.path.join(ROOT, 'test_batch'))\n",
    "    return Xtr, Ytr, Xte, Yte\n",
    "#dataset\n",
    "cifar10_dir = 'cifar-10-batches-py'\n",
    "train_image, train_label, test_image, _ = load_CIFAR10(cifar10_dir)\n",
    "#Normalizing\n",
    "train_image /= 255\n",
    "test_image /= 255"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def show_result(num_epoch, show = False, save = False, path = 'result.png'):\n",
    "    test_images = sess.run(G_sample, {z: fixed_z_, y: fixed_y_})\n",
    "\n",
    "    size_figure_grid = 10\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], (32, 32, 3)))\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()\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",
    "def preprocess_img(x):\n",
    "    return 2 * x - 1.0\n",
    "\n",
    "def leaky_relu(x, alpha=0.2, name=\"lrelu\"):\n",
    "    condition = tf.less(x, 0)\n",
    "    return tf.where(condition, alpha * x, x)\n",
    "    \n",
    "def sample_noise(batch_size, dim):\n",
    "    z = np.random.randn(batch_size, dim)\n",
    "    return z\n",
    "\n",
    "def get_session():\n",
    "    config = tf.ConfigProto()\n",
    "    config.gpu_options.allow_growth = True\n",
    "    session = tf.Session(config=config)\n",
    "    return session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# label preprocess\n",
    "onehot = np.eye(10)\n",
    "temp_z_ = np.random.normal(0, 1, (10, 100))\n",
    "fixed_z_ = temp_z_\n",
    "fixed_y_ = np.zeros((10, 1))\n",
    "\n",
    "for i in range(9):\n",
    "    fixed_z_ = np.concatenate([fixed_z_, temp_z_], 0)\n",
    "    temp = np.ones((10,1)) + i\n",
    "    fixed_y_ = np.concatenate([fixed_y_, temp], 0)\n",
    "\n",
    "fixed_y_ = onehot[fixed_y_.astype(np.int32)].squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#dataset\n",
    "# class MNIST(object):\n",
    "#     def __init__(self, batch_size, shuffle=False):\n",
    "#         train, _ = tf.keras.datasets.mnist.load_data()\n",
    "#         X, y = train\n",
    "#         X = X.astype(np.float32)/255.\n",
    "#         X = X.reshape((X.shape[0], -1))\n",
    "#         self.X, self.y = X, y\n",
    "#         self.batch_size, self.shuffle = batch_size, shuffle\n",
    "    \n",
    "#     def __iter__(self):\n",
    "#         N, B = self.X.shape[0], self.batch_size\n",
    "#         idxs = np.range(N)\n",
    "#         if self.shuffle:\n",
    "#             np.random.shuffle(idxs)\n",
    "#         return iter((self.X[i:i+B], self.y[i:i+B]) for i in range(0, N, B))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# generator and discriminator\n",
    "def generator(z, y):\n",
    "    with tf.variable_scope('generator'):\n",
    "        w_init = tf.contrib.layers.xavier_initializer()\n",
    "        combine = tf.concat([z, y], 1)\n",
    "        dense1 = tf.layers.dense(combine, 4*4*1024, activation=tf.nn.relu, kernel_initializer=w_init)\n",
    "        dense2 = tf.reshape(dense1, (-1, 4, 4, 1024))\n",
    "        conv1 = tf.layers.conv2d_transpose(dense2, filters=512, kernel_size=3, strides=2, padding='SAME', activation=tf.nn.relu, kernel_initializer=w_init)\n",
    "        conv2 = tf.layers.conv2d_transpose(conv1, filters=256, kernel_size=3, strides=2, padding='SAME', activation=tf.nn.relu, kernel_initializer=w_init)\n",
    "        output = tf.layers.conv2d_transpose(conv2, filters=3, kernel_size=3, strides=2, padding='SAME', activation=tf.nn.tanh, kernel_initializer=w_init)\n",
    "        return output\n",
    "\n",
    "def discriminator(x, y):\n",
    "    with tf.variable_scope('discriminator'):\n",
    "        w_init = tf.contrib.layers.xavier_initializer()\n",
    "        x = tf.reshape(x, (100, 32, 32, 3))\n",
    "        y = tf.reshape(y,(100, 1, 1, 10))\n",
    "        combine = tf.concat([x, y * tf.ones([x.shape[0], x.shape[1], x.shape[2], y.shape[3]])], 3) # combine: (100, 32, 32, 13)\n",
    "        conv1 = tf.layers.conv2d(combine, filters=64, kernel_size=3, strides=1, padding='SAME', kernel_initializer=w_init) #(100,32,32,64)\n",
    "        lrelu1 = leaky_relu(conv1)\n",
    "        conv2 = tf.layers.conv2d(lrelu1, filters=128, kernel_size=4, strides=2, padding='SAME', kernel_initializer=w_init)#(16,16,128)\n",
    "        lrelu2 = leaky_relu(conv2)\n",
    "        conv3 = tf.layers.conv2d(lrelu2, filters=256, kernel_size=4, strides=2, padding='SAME', kernel_initializer=w_init)#(8,8,256)\n",
    "        lrelu3 = leaky_relu(conv3)\n",
    "        conv4 = tf.layers.conv2d(lrelu3, filters=512, kernel_size=4, strides=2, padding='SAME', kernel_initializer=w_init)#(4,4,512)\n",
    "        flatten = tf.layers.flatten(conv4)\n",
    "        output = tf.layers.dense(flatten, 1, kernel_initializer=w_init)\n",
    "        return output\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def gan_loss(logits_real, logits_fake):\n",
    "    label_real = tf.ones_like(logits_real)\n",
    "    label_fake = tf.zeros_like(logits_fake)\n",
    "    \n",
    "    D_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=label_real, logits=logits_real))\n",
    "    D_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=label_fake, logits=logits_fake))\n",
    "    D_loss = D_loss_real + D_loss_fake\n",
    "    \n",
    "    G_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.ones_like(logits_fake), logits=logits_fake))\n",
    "    return D_loss, G_loss\n",
    "\n",
    "def get_solver(learning_rate=0.0002, beta1=0.5):\n",
    "    D_solver = tf.train.AdamOptimizer(learning_rate=learning_rate, beta1=beta1)\n",
    "    G_solver = tf.train.AdamOptimizer(learning_rate=learning_rate, beta1=beta1)\n",
    "    return D_solver, G_solver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-11-3f6dbb24b1c1>:6: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.dense instead.\n",
      "WARNING:tensorflow:From C:\\Users\\THINKPAD\\Anaconda5.3.0\\envs\\cs231n\\lib\\site-packages\\tensorflow\\python\\framework\\op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Colocations handled automatically by placer.\n",
      "WARNING:tensorflow:From <ipython-input-11-3f6dbb24b1c1>:8: conv2d_transpose (from tensorflow.python.layers.convolutional) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.conv2d_transpose instead.\n",
      "WARNING:tensorflow:From <ipython-input-11-3f6dbb24b1c1>:19: conv2d (from tensorflow.python.layers.convolutional) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.conv2d instead.\n",
      "WARNING:tensorflow:From <ipython-input-11-3f6dbb24b1c1>:26: flatten (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.flatten instead.\n"
     ]
    }
   ],
   "source": [
    "#put it together\n",
    "tf.reset_default_graph()\n",
    "\n",
    "# number of images for each batch\n",
    "batch_size = 100\n",
    "\n",
    "\n",
    "x = tf.placeholder(tf.float32, [None, 32, 32, 3])\n",
    "y = tf.placeholder(tf.float32, [None, 10])\n",
    "z = tf.placeholder(tf.float32, [None, 100])\n",
    "\n",
    "G_sample = generator(z, y)\n",
    "\n",
    "with tf.variable_scope(\"\") as scope:\n",
    "    logits_real = discriminator(x, y)\n",
    "    scope.reuse_variables()\n",
    "    logits_fake = discriminator(G_sample, y)\n",
    "\n",
    "D_var = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,'discriminator')\n",
    "G_var = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,'generator')\n",
    "\n",
    "D_solver, G_solver = get_solver()\n",
    "D_loss, G_loss = gan_loss(logits_real, logits_fake)\n",
    "\n",
    "D_train_step = D_solver.minimize(D_loss, var_list=D_var)\n",
    "G_train_step = G_solver.minimize(G_loss, var_list=G_var)\n",
    "D_extra_step = tf.get_collection(tf.GraphKeys.UPDATE_OPS,'discriminator')\n",
    "G_extra_step = tf.get_collection(tf.GraphKeys.UPDATE_OPS,'generator')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#run a gan\n",
    "def run_a_gan(sess, D_train_step, D_loss, G_train_step, G_loss, D_extra_step,G_extra_step,\\\n",
    "              batch_size=100, num_epoch=100):\n",
    "    #(x_train, y_train), (x_test, _) = tf.keras.datasets.mnist.load_data()\n",
    "    x_train = train_image\n",
    "    y_train = train_label\n",
    "    \n",
    "    train_hist = {}\n",
    "    train_hist['D_losses'] = []\n",
    "    train_hist['G_losses'] = []\n",
    "    \n",
    "    root = 'cifar_CGAN_result/'\n",
    "    model = 'conv/'\n",
    "    if not os.path.isdir(root):\n",
    "        os.mkdir(root)\n",
    "    if not os.path.isdir(root + model):\n",
    "        os.mkdir(root + model)\n",
    "    \n",
    "    for epoch in range(num_epoch):\n",
    "        for iter in range(len(x_train) // batch_size):\n",
    "            D_losses = []\n",
    "            G_losses = []\n",
    "            #update discriminator\n",
    "            x_ = x_train[iter * batch_size:(iter + 1) * batch_size]\n",
    "            y_ = y_train[iter * batch_size:(iter + 1) * batch_size]\n",
    "            y_ = y_[:,np.newaxis]\n",
    "            y_ = onehot[y_.astype(np.int32)].squeeze()\n",
    "            \n",
    "            z_ = sample_noise(batch_size, 100)\n",
    "            _, D_loss_curr = sess.run([D_train_step, D_loss], feed_dict={x:x_, y:y_, z:z_})\n",
    "            D_losses.append(D_loss_curr)\n",
    "            #update generator\n",
    "            y_ = np.random.randint(0, 10, (batch_size, 1))\n",
    "            y_ = onehot[y_.astype(np.int32)].squeeze()\n",
    "            z_ = sample_noise(batch_size, 100)\n",
    "            _, G_loss_curr = sess.run([G_train_step, G_loss], feed_dict={z:z_, x:x_, y:y_})\n",
    "            G_losses.append(G_loss_curr)\n",
    "        print(\"Epoch:{}, D:{:.4}, G:{:.4}\".format(epoch+1, D_loss_curr, G_loss_curr))\n",
    "  \n",
    "        PATH = root + model + str(epoch+1) + '.png' \n",
    "\n",
    "        show_result(epoch+1, save=True, path=PATH)\n",
    "        train_hist['D_losses'].append(np.mean(D_losses))\n",
    "        train_hist['G_losses'].append(np.mean(G_losses))\n",
    "        \n",
    "    with open(root + model + 'train_hist.pkl', 'wb') as f:\n",
    "        pickle.dump(train_hist, f)\n",
    "\n",
    "    show_train_hist(train_hist, save=True, path=root + model + 'train_hist.png')\n",
    "\n",
    "    images = []\n",
    "    for e in range(num_epoch):\n",
    "        img_name = root + model + str(e + 1) + '.png'\n",
    "        images.append(imageio.imread(img_name))\n",
    "    imageio.mimsave(root + model + 'generation_animation.gif', images, fps=5)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with get_session() as sess:\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    run_a_gan(sess,D_train_step,D_loss,G_train_step,G_loss,D_extra_step,G_extra_step)\n",
    "    sess.close()"
   ]
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}