GAN/sngan_ttur.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 = 1\n",
    "BATCH_SIZE = 256\n",
    "ITERS = 100000\n",
    "DATASET = '8gaussians'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def spectral_norm(w, iteration=1):\n",
    "    w_shape = w.shape.as_list()\n",
    "    w = tf.reshape(w, [-1, w_shape[-1]])\n",
    "\n",
    "    u = tf.get_variable('u', [1, w_shape[-1]], initializer=tf.random_normal_initializer(), trainable=False)\n",
    "\n",
    "    u_hat = u\n",
    "    v_hat = None\n",
    "    for i in range(iteration):\n",
    "        v_ = tf.matmul(u_hat, tf.transpose(w))\n",
    "        v_hat = tf.nn.l2_normalize(v_)\n",
    "\n",
    "        u_ = tf.matmul(v_hat, w)\n",
    "        u_hat = tf.nn.l2_normalize(u_)\n",
    "\n",
    "    u_hat = tf.stop_gradient(u_hat)\n",
    "    v_hat = tf.stop_gradient(v_hat)\n",
    "\n",
    "    sigma = tf.matmul(tf.matmul(v_hat, w), tf.transpose(u_hat))\n",
    "\n",
    "    with tf.control_dependencies([u.assign(u_hat)]):\n",
    "        w_norm = w / sigma\n",
    "        w_norm = tf.reshape(w_norm, w_shape)\n",
    "\n",
    "    return w_norm\n",
    "\n",
    "def snlinear(inputs, input_dim, output_dim, activation, name):\n",
    "    with tf.variable_scope(name):\n",
    "        weights = tf.get_variable(\n",
    "            'kernel',\n",
    "            [input_dim, output_dim],\n",
    "            initializer=tf.random_normal_initializer(0., output_dim ** -0.5)\n",
    "        )\n",
    "        biases = tf.get_variable('bias', [output_dim])\n",
    "        outputs = tf.matmul(inputs, spectral_norm(weights)) + biases\n",
    "        if activation:\n",
    "            outputs = activation(outputs)\n",
    "\n",
    "        return outputs\n",
    "\n",
    "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 = snlinear(inputs, 2, DIM, activation=tf.nn.relu, name='Discriminator.1')\n",
    "        outputs = snlinear(outputs, DIM, DIM, activation=tf.nn.relu, name='Discriminator.2')\n",
    "        outputs = snlinear(outputs, DIM, DIM, activation=tf.nn.relu, name='Discriminator.3')\n",
    "        outputs = snlinear(outputs, DIM, 1, activation=None, name='Discriminator.4')\n",
    "        return outputs"
   ]
  },
  {
   "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)"
   ]
  },
  {
   "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=3e-4,\n",
    "    beta1=0.0,\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.0,\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",
    "        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
}