GAN/wgan.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "import tensorflow as tf\n",
    "import sklearn.datasets\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DIM = 512\n",
    "CRITIC_ITERS = 5\n",
    "BATCH_SIZE = 256\n",
    "ITERS = 100000\n",
    "LAMBDA = .1\n",
    "DATASET = '8gaussians'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Generator(n_samples):\n",
    "    with tf.variable_scope('Generator', reuse=tf.AUTO_REUSE):\n",
    "        noise = tf.random.normal([n_samples, 2])\n",
    "        outputs = tf.layers.dense(noise, DIM, activation=tf.nn.relu, name='Generator.1')\n",
    "        outputs = tf.layers.dense(outputs, DIM, activation=tf.nn.relu, name='Generator.2')\n",
    "        outputs = tf.layers.dense(outputs, DIM, activation=tf.nn.relu, name='Generator.3')\n",
    "        outputs = tf.layers.dense(outputs, 2, activation=None, name='Generator.4')\n",
    "        return outputs\n",
    "\n",
    "def Discriminator(inputs):\n",
    "    with tf.variable_scope('Discriminator', reuse=tf.AUTO_REUSE):\n",
    "        outputs = tf.layers.dense(inputs, DIM, activation=tf.nn.relu, name='Discriminator.1')\n",
    "        outputs = tf.layers.dense(outputs, DIM, activation=tf.nn.relu, name='Discriminator.2')\n",
    "        outputs = tf.layers.dense(outputs, DIM, activation=tf.nn.relu, name='Discriminator.3')\n",
    "        outputs = tf.layers.dense(outputs, 1, activation=None, name='Discriminator.4')\n",
    "        return outputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "real_data = tf.placeholder(tf.float32, shape=[None, 2])\n",
    "\n",
    "fake_data = Generator(BATCH_SIZE)\n",
    "disc_real = Discriminator(real_data)\n",
    "disc_fake = Discriminator(fake_data)\n",
    "\n",
    "disc_loss = tf.reduce_mean(disc_fake) - tf.reduce_mean(disc_real)\n",
    "gen_loss = -tf.reduce_mean(disc_fake)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = tf.random_uniform(shape=[BATCH_SIZE, 1], minval=0.0, maxval=1.0)\n",
    "interpolates = alpha*real_data + ((1-alpha)*fake_data)\n",
    "disc_interpolates = Discriminator(interpolates)\n",
    "gradients = tf.gradients(disc_interpolates, [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",
    "\n",
    "disc_loss += LAMBDA * gradient_penalty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "disc_params = tf.trainable_variables('Discriminator')\n",
    "gen_params = tf.trainable_variables('Generator')\n",
    "\n",
    "disc_train_op = tf.train.AdamOptimizer(\n",
    "    learning_rate=1e-4,\n",
    "    beta1=0.5,\n",
    "    beta2=0.9\n",
    ").minimize(disc_loss, var_list=disc_params)\n",
    "\n",
    "gen_train_op = tf.train.AdamOptimizer(\n",
    "    learning_rate=1e-4,\n",
    "    beta1=0.5,\n",
    "    beta2=0.9\n",
    ").minimize(gen_loss, var_list=gen_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_real_data():\n",
    "    if DATASET == '25gaussians':\n",
    "        positions = 2 * np.random.choice(range(-2, 3), size=BATCH_SIZE*2)\n",
    "        noise = 0.05 * np.random.randn(BATCH_SIZE*2)\n",
    "\n",
    "        yield np.reshape(positions + noise, [BATCH_SIZE, 2]) / 2.828\n",
    "    elif DATASET == 'swissroll':\n",
    "        data = sklearn.datasets.make_swiss_roll(\n",
    "            n_samples=BATCH_SIZE, \n",
    "            noise=0.25\n",
    "        )[0]\n",
    "        data = data.astype('float32')[:, [0, 2]]\n",
    "        data /= 7.5\n",
    "        yield data\n",
    "    elif DATASET == '8gaussians':\n",
    "        scale = 2.\n",
    "        centers = [\n",
    "            (1,0),\n",
    "            (-1,0),\n",
    "            (0,1),\n",
    "            (0,-1),\n",
    "            (1./np.sqrt(2), 1./np.sqrt(2)),\n",
    "            (1./np.sqrt(2), -1./np.sqrt(2)),\n",
    "            (-1./np.sqrt(2), 1./np.sqrt(2)),\n",
    "            (-1./np.sqrt(2), -1./np.sqrt(2))\n",
    "        ]\n",
    "        centers = [(scale*x,scale*y) for x,y in centers]\n",
    "        dataset = []\n",
    "        for i in range(BATCH_SIZE):\n",
    "            point = np.random.randn(2)*.02\n",
    "            center = random.choice(centers)\n",
    "            point[0] += center[0]\n",
    "            point[1] += center[1]\n",
    "            dataset.append(point)\n",
    "        dataset = np.array(dataset, dtype='float32')\n",
    "        dataset /= 1.414\n",
    "        yield dataset\n",
    "\n",
    "def generate_image():\n",
    "    N_POINTS = 128\n",
    "    RANGE = 3\n",
    "    points = np.zeros((N_POINTS, N_POINTS, 2), dtype='float32')\n",
    "    points[:,:,0] = np.linspace(-RANGE, RANGE, N_POINTS)[:,None]\n",
    "    points[:,:,1] = np.linspace(-RANGE, RANGE, N_POINTS)[None,:]\n",
    "    points = points.reshape((-1,2))\n",
    "    disc_map = sess.run(disc_real, feed_dict={real_data:points})\n",
    "    x = y = np.linspace(-RANGE, RANGE, N_POINTS)\n",
    "    plt.contour(x, y, disc_map.reshape((len(x), len(y))).transpose())\n",
    "    \n",
    "    fake_data_samples = sess.run(fake_data)\n",
    "    real_data_samples = next(create_real_data())\n",
    "\n",
    "    plt.scatter(real_data_samples[:, 0], real_data_samples[:, 1], c='orange', marker='+')\n",
    "    plt.scatter(fake_data_samples[:, 0], fake_data_samples[:, 1], c='green', marker='+')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess = tf.Session()\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "for i in range(ITERS+1):\n",
    "    for _ in range(CRITIC_ITERS):\n",
    "        fd = {real_data: next(create_real_data())}\n",
    "        _disc_loss, _ = sess.run([disc_loss, disc_train_op], fd)\n",
    "\n",
    "    sess.run(gen_train_op)\n",
    "\n",
    "    if i % 10000 == 0:\n",
    "        print(i, _disc_loss)\n",
    "        generate_image()\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
}