demo_mnist.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A simple demo for PGVGrow based on PGGAN implementation\n",
    "### Please run this demo after preprocessing mnist dataset with data_tool.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import imageio\n",
    "import os\n",
    "from collections import OrderedDict\n",
    "import matplotlib.pyplot as plt\n",
    "import data_tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Hyper-parameters\n",
    "\"\"\"\n",
    "# different minibatch size for different resolution\n",
    "minibatch_dict = {4: 256, 8: 256, 16: 128, 32: 64}\n",
    "# number of channels, 1 for gray images, 3 for RGB images\n",
    "num_channels = 1\n",
    "# number of feature maps\n",
    "nf = 128\n",
    "# resolution of raw images\n",
    "resolution = 32\n",
    "resolution_log2 = int(np.log2(resolution))\n",
    "init_resolution = 4\n",
    "# dimension of latent space\n",
    "z_dim = 128\n",
    "# number of images used in a phase\n",
    "dur_nimg = 300000\n",
    "# number of images used in total\n",
    "total_nimg = 1800000\n",
    "\n",
    "# number of images for visualizing\n",
    "num_row = 10\n",
    "num_line = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# a function for visualizing\n",
    "def montage(images, grid):\n",
    "\n",
    "    s = np.shape(images)\n",
    "    assert s[0] == np.prod(grid) and np.shape(s)[0] == 4\n",
    "    bigimg = np.zeros((s[1]*grid[0], s[1]*grid[1], s[3]), dtype=np.float32)\n",
    "\n",
    "    for i in range(grid[0]):\n",
    "        for j in range(grid[1]):\n",
    "            bigimg[s[1] * i : s[1] * i + s[1], s[1] * j : s[1] * j + s[1]] += images[grid[1] * i + j]\n",
    "\n",
    "    return np.rint(bigimg*255).clip(0, 255).astype(np.uint8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate coefficients with given divergence\n",
    "def coef_div(d_score, div):\n",
    "\n",
    "    if div == 'KL':\n",
    "        s = np.ones_like(d_score)\n",
    "    elif div == 'LogD':\n",
    "        s = 1 / (1 + np.exp(d_score))\n",
    "    elif div == 'JS':\n",
    "        s = 1 / (1 + np.exp(-d_score))\n",
    "    elif div == 'Jef':\n",
    "        s = 1 + np.exp(d_score)\n",
    "\n",
    "    return np.reshape(s, [-1, 1, 1, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lod(num_img):\n",
    "\n",
    "    ph_num = num_img // (2*dur_nimg)\n",
    "    remain_num = num_img - ph_num * (2 * dur_nimg)\n",
    "\n",
    "    if np.log2(resolution / init_resolution) <= ph_num:\n",
    "        return 0.\n",
    "    elif remain_num <= dur_nimg:\n",
    "        return np.log2(resolution / init_resolution) - ph_num\n",
    "    else: \n",
    "        return np.log2(resolution / init_resolution) - ph_num - (remain_num - dur_nimg) / dur_nimg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lerp(a, b, t): return a + (b - a) * t\n",
    "\n",
    "def cset(cur_lambda, new_cond, new_lambda): return lambda: tf.cond(new_cond, new_lambda, cur_lambda)\n",
    "\n",
    "def upscale2d(x, factor=2):\n",
    "    s = x.shape\n",
    "    x = tf.reshape(x, [-1, s[1], 1, s[2], 1, s[3]])\n",
    "    x = tf.tile(x, [1, 1, factor, 1, factor, 1])\n",
    "    x = tf.reshape(x, [-1, s[1] * factor, s[2] * factor, s[3]])\n",
    "    return x\n",
    "\n",
    "def downscale2d(x, factor=2):\n",
    "    return tf.nn.avg_pool(x, ksize=[1, factor, factor, 1], strides=[1, factor, factor, 1], padding='VALID')\n",
    "\n",
    "def pixel_norm(x, axis=3):\n",
    "    return x / tf.sqrt(tf.reduce_mean(tf.square(x), axis=axis, keepdims=True) + 1e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generator(latents_in, lod_in, num_channels=num_channels, resolution=resolution, latent_size=z_dim, num_features=nf, reuse=None):\n",
    "\n",
    "    def G_block(x, res):\n",
    "        with tf.variable_scope('%dx%d' % (2**res, 2**res), reuse=tf.AUTO_REUSE):\n",
    "            if res == 2:\n",
    "                x = pixel_norm(x, axis=1)\n",
    "                with tf.variable_scope('Dense'):\n",
    "                    w = tf.get_variable('Weight', dtype=tf.float32, shape=[latent_size, num_features*16], initializer=tf.initializers.random_normal())\n",
    "                    std = tf.constant(np.sqrt(2/latent_size), dtype=tf.float32, name='std')\n",
    "                    x = tf.reshape(tf.matmul(x, w*std), [-1, 4, 4, num_features])\n",
    "                    b = tf.get_variable('Bias', dtype=tf.float32, shape=[num_features], initializer=tf.initializers.zeros())\n",
    "                    x = pixel_norm(tf.nn.leaky_relu(x + b, alpha=0.2), axis=1)\n",
    "                with tf.variable_scope('Conv'):\n",
    "                    w = tf.get_variable('Kernel', dtype=tf.float32, shape=[3, 3, num_features, num_features], initializer=tf.initializers.random_normal())\n",
    "                    std = tf.constant(np.sqrt(2/(3*3*num_features)), dtype=tf.float32, name='std')\n",
    "                    b = tf.get_variable('Bias', dtype=tf.float32, shape=[num_features], initializer=tf.initializers.zeros())\n",
    "                    x = pixel_norm(tf.nn.leaky_relu(tf.nn.conv2d(x, w*std, strides=[1, 1, 1, 1], padding='SAME') + b, alpha=0.2))\n",
    "            else:\n",
    "                x = upscale2d(x)\n",
    "                with tf.variable_scope('Conv1'):\n",
    "                    w = tf.get_variable('Kernel', dtype=tf.float32, shape=[3, 3, num_features, num_features], initializer=tf.initializers.random_normal())\n",
    "                    std = tf.constant(np.sqrt(2/(3*3*num_features)), dtype=tf.float32, name='std')\n",
    "                    b = tf.get_variable('Bias', dtype=tf.float32, shape=[num_features], initializer=tf.initializers.zeros())\n",
    "                    x = pixel_norm(tf.nn.leaky_relu(tf.nn.conv2d(x, w*std, strides=[1, 1, 1, 1], padding='SAME') + b, alpha=0.2))\n",
    "                with tf.variable_scope('Conv2'):\n",
    "                    w = tf.get_variable('Kernel', dtype=tf.float32, shape=[3, 3, num_features, num_features], initializer=tf.initializers.random_normal())\n",
    "                    std = tf.constant(np.sqrt(2/(3*3*num_features)), dtype=tf.float32, name='std')\n",
    "                    b = tf.get_variable('Bias', dtype=tf.float32, shape=[num_features], initializer=tf.initializers.zeros())\n",
    "                    x = pixel_norm(tf.nn.leaky_relu(tf.nn.conv2d(x, w*std, strides=[1, 1, 1, 1], padding='SAME') + b, alpha=0.2))\n",
    "            return x\n",
    "\n",
    "    def FtoI(x, res): # Feature maps to Image\n",
    "        with tf.variable_scope('FtoI_%dx%d' % (2**res, 2**res), reuse=tf.AUTO_REUSE):\n",
    "            w = tf.get_variable('Kernel', dtype=tf.float32, shape=[1, 1, num_features, num_channels], initializer=tf.initializers.random_normal())\n",
    "            std = tf.constant(np.sqrt(1/num_features), dtype=tf.float32, name='std')\n",
    "            b = tf.get_variable('Bias', dtype=tf.float32, shape=[num_channels], initializer=tf.initializers.zeros())\n",
    "            x = tf.nn.conv2d(x, w*std, strides=[1, 1, 1, 1], padding='SAME') + b\n",
    "            return x\n",
    "\n",
    "    def grow(x, res, lod):\n",
    "        y = G_block(x, res)\n",
    "        img = lambda: upscale2d(FtoI(y, res), 2**lod)\n",
    "        if res > 2: img = cset(img, (lod_in > lod), lambda: upscale2d(lerp(FtoI(y, res), upscale2d(FtoI(x, res - 1)), lod_in - lod), 2**lod))\n",
    "        if lod > 0: img = cset(img, (lod_in < lod), lambda: grow(y, res + 1, lod - 1))\n",
    "        return img()\n",
    "\n",
    "    with tf.variable_scope('generator', reuse=reuse):\n",
    "        resolution_log2 = int(np.log2(resolution))\n",
    "        images_out = grow(latents_in, 2, resolution_log2 - 2)\n",
    "        return images_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def discriminator(images_in, lod_in, num_channels=num_channels, resolution=resolution, num_features=nf, reuse=None):\n",
    "\n",
    "    def D_block(x, res):\n",
    "        with tf.variable_scope('%dx%d' % (2**res, 2**res), reuse=tf.AUTO_REUSE):\n",
    "            if res >= 3:\n",
    "                with tf.variable_scope('Conv1'):\n",
    "                    w = tf.get_variable('Kernel', dtype=tf.float32, shape=[3, 3, num_features, num_features], initializer=tf.initializers.random_normal())\n",
    "                    std = tf.constant(np.sqrt(2/(3*3*num_features)), dtype=tf.float32, name='std')\n",
    "                    b = tf.get_variable('Bias', dtype=tf.float32, shape=[num_features], initializer=tf.initializers.zeros())\n",
    "                    x = tf.nn.leaky_relu(tf.nn.conv2d(x, w*std, strides=[1, 1, 1, 1], padding='SAME') + b, alpha=0.2)\n",
    "                with tf.variable_scope('Conv2'):\n",
    "                    w = tf.get_variable('Kernel', dtype=tf.float32, shape=[3, 3, num_features, num_features], initializer=tf.initializers.random_normal())\n",
    "                    std = tf.constant(np.sqrt(2/(3*3*num_features)), dtype=tf.float32, name='std')\n",
    "                    b = tf.get_variable('Bias', dtype=tf.float32, shape=[num_features], initializer=tf.initializers.zeros())\n",
    "                    x = tf.nn.leaky_relu(tf.nn.conv2d(x, w*std, strides=[1, 1, 1, 1], padding='SAME') + b, alpha=0.2)\n",
    "                x = downscale2d(x)\n",
    "            else:\n",
    "                with tf.variable_scope('Conv'):\n",
    "                    w = tf.get_variable('Kernel', dtype=tf.float32, shape=[3, 3, num_features, num_features], initializer=tf.initializers.random_normal())\n",
    "                    std = tf.constant(np.sqrt(2/(3*3*num_features)), dtype=tf.float32, name='std')\n",
    "                    b = tf.get_variable('Bias', dtype=tf.float32, shape=[num_features], initializer=tf.initializers.zeros())\n",
    "                    x = tf.nn.leaky_relu(tf.nn.conv2d(x, w*std, strides=[1, 1, 1, 1], padding='SAME') + b, alpha=0.2)\n",
    "                with tf.variable_scope('Dense1'):\n",
    "                    x = tf.layers.flatten(x)\n",
    "                    w = tf.get_variable('Weight', dtype=tf.float32, shape=[16*num_features, num_features], initializer=tf.initializers.random_normal())\n",
    "                    std = tf.constant(np.sqrt(2/(16*num_features)), dtype=tf.float32, name='std')\n",
    "                    b = tf.get_variable('Bias', dtype=tf.float32, shape=num_features, initializer=tf.initializers.zeros())\n",
    "                    x = tf.nn.leaky_relu(tf.nn.xw_plus_b(x, w*std, b), alpha=0.2)\n",
    "                with tf.variable_scope('Dense2'):\n",
    "                    w = tf.get_variable('Weight', dtype=tf.float32, shape=[num_features, 1], initializer=tf.initializers.random_normal())\n",
    "                    std = tf.constant(np.sqrt(1/num_features), dtype=tf.float32, name='std')\n",
    "                    b = tf.get_variable('Bias', dtype=tf.float32, shape=1, initializer=tf.initializers.zeros())\n",
    "                    x = tf.nn.xw_plus_b(x, w*std, b)\n",
    "            return x\n",
    "\n",
    "    def ItoF(x, res): # Image to Feature maps\n",
    "        with tf.variable_scope('ItoF_%dx%d' % (2**res, 2**res), reuse=tf.AUTO_REUSE):\n",
    "            w = tf.get_variable('Kernel', dtype=tf.float32, shape=[1, 1, num_channels, num_features], initializer=tf.initializers.random_normal())\n",
    "            std = tf.constant(np.sqrt(2/num_channels), dtype=tf.float32, name='std')\n",
    "            b = tf.get_variable('Bias', dtype=tf.float32, shape=[num_features], initializer=tf.initializers.zeros())\n",
    "            x = tf.nn.leaky_relu(tf.nn.conv2d(x, w*std, strides=[1, 1, 1, 1], padding='SAME') + b, alpha=0.2)\n",
    "            return x\n",
    "\n",
    "    def grow(res, lod):\n",
    "        x = lambda: ItoF(downscale2d(images_in, 2**lod), res)\n",
    "        if lod > 0: x = cset(x, (lod_in < lod), lambda: grow(res + 1, lod - 1))\n",
    "        x = D_block(x(), res); y = lambda: x\n",
    "        if res > 2: y = cset(y, (lod_in > lod), lambda: lerp(x, ItoF(downscale2d(images_in, 2**(lod+1)), res - 1), lod_in - lod))\n",
    "        return y()\n",
    "\n",
    "    with tf.variable_scope('discriminator', reuse=reuse):\n",
    "        resolution_log2 = int(np.log2(resolution))\n",
    "        scores_out = grow(2, resolution_log2 - 2)\n",
    "        return scores_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "pre-process real images\n",
    "\"\"\"\n",
    "\n",
    "def downscale(img):\n",
    "    s = img.shape\n",
    "    out = np.reshape(img, [-1, s[1]//2, 2, s[2]//2, 2, s[3]])\n",
    "    return np.mean(out, axis=(2, 4))\n",
    "\n",
    "def upscale(img):\n",
    "    return np.repeat(np.repeat(img, 2, axis=1), 2, axis=2)\n",
    "\n",
    "def process_real(x, lod_in):\n",
    "    y = x / 127.5 - 1\n",
    "    alpha = lod_in - np.floor(lod_in)\n",
    "    y = (1 - alpha)*y + alpha*upscale(downscale(y))\n",
    "    for i in range(int(np.floor(lod_in))):\n",
    "        y = upscale(y)\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "bulid graph\n",
    "\"\"\"\n",
    "\n",
    "x_p    = tf.placeholder(dtype=tf.float32, shape=[None, resolution, resolution, num_channels], name='images')\n",
    "z_p    = tf.placeholder(dtype=tf.float32, shape=[None, z_dim], name='latents')\n",
    "G_z_p  = tf.placeholder(dtype=tf.float32, shape=[None, resolution, resolution, num_channels], name='particles')\n",
    "lod_in = tf.placeholder(dtype=tf.float32, shape=[], name='level_of_details')\n",
    "\n",
    "G_z   = generator(z_p, lod_in)\n",
    "\n",
    "# discriminator loss:\n",
    "d_real_logits = discriminator(x_p, lod_in)\n",
    "d_fake_logits = discriminator(G_z_p, lod_in, reuse=True)\n",
    "loss_d_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_real_logits, labels=tf.ones_like(d_real_logits)))\n",
    "loss_d_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_fake_logits, labels=tf.zeros_like(d_fake_logits)))\n",
    "loss_d = loss_d_real + loss_d_fake\n",
    "\n",
    "# generator loss:\n",
    "loss_g = 0.5 * tf.reduce_mean(tf.reduce_sum((G_z - G_z_p)**2))\n",
    "\n",
    "# computing gradient:\n",
    "d_grad = tf.gradients(d_fake_logits, G_z_p)[0]\n",
    "\n",
    "# optimizers:\n",
    "vars_g = [var for var in tf.trainable_variables() if var.name.startswith('generator')]\n",
    "vars_d = [var for var in tf.trainable_variables() if var.name.startswith('discriminator')]\n",
    "\n",
    "optimizer_d = tf.train.AdamOptimizer(learning_rate=0.001, beta1=0.0, beta2=0.99, epsilon=1e-8, name='opt_d')\n",
    "update_d    = optimizer_d.minimize(loss_d, var_list=vars_d)\n",
    "\n",
    "optimizer_g = tf.train.AdamOptimizer(learning_rate=0.001, beta1=0.0, beta2=0.99, epsilon=1e-8, name='opt_g')\n",
    "update_g    = optimizer_g.minimize(loss_g, var_list=vars_g)\n",
    "\n",
    "reset_optimizer_d = tf.variables_initializer(optimizer_d.variables())\n",
    "reset_optimizer_g = tf.variables_initializer(optimizer_g.variables())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# KL divergence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with tf.Session() as sess:\n",
    "\n",
    "    iterators = [data_tool.data_iterator('mnist', lod_in=lod, batch_size=minibatch_dict[2**(resolution_log2-lod)], resolution_log2=resolution_log2) for lod in range(int(np.log2(resolution/init_resolution))+1)]\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    num_img = 0\n",
    "    tick_kimg = 0\n",
    "    z_fixed = np.random.randn(num_row*num_line, z_dim)\n",
    "    prev_lod = -1.0\n",
    "\n",
    "    while num_img <= total_nimg:\n",
    "        cur_lod = lod(num_img)\n",
    "        # reset Adam optimizers states when increasing resolution:\n",
    "        if np.floor(cur_lod) != np.floor(prev_lod) or np.ceil(cur_lod) != np.ceil(prev_lod):\n",
    "            sess.run([reset_optimizer_d, reset_optimizer_g])\n",
    "        prev_lod = cur_lod\n",
    "        # get mini-batch size:\n",
    "        batch_size = minibatch_dict[2**int(resolution_log2-np.floor(cur_lod))]\n",
    "\n",
    "        # sample a latent pool and get particles:\n",
    "        z = np.random.randn(batch_size, z_dim)\n",
    "        P = sess.run(G_z, feed_dict={z_p: z, lod_in: cur_lod})\n",
    "\n",
    "        # optimize discriminator:\n",
    "        x = next(iterators[int(np.floor(cur_lod))])\n",
    "        x = process_real(x, cur_lod)\n",
    "        num_img += batch_size # count\n",
    "        sess.run(update_d, feed_dict={x_p: x, G_z_p: P, lod_in: cur_lod})\n",
    "\n",
    "        # move particles\n",
    "        d_score = sess.run(d_fake_logits, feed_dict={G_z_p: P, lod_in: cur_lod})\n",
    "        grad = sess.run(d_grad, feed_dict={G_z_p: P, lod_in: cur_lod})\n",
    "        P += coef_div(d_score, div='KL') * grad\n",
    "\n",
    "        # optimize generator:\n",
    "        sess.run(update_g, feed_dict={z_p: z, G_z_p: P, lod_in: cur_lod})\n",
    "\n",
    "        if (num_img // 1000) >= tick_kimg + 150:\n",
    "            tick_kimg = (num_img // 1000)\n",
    "            cur_lod = lod(num_img)\n",
    "            real_loss, fake_loss = sess.run([loss_d_real, loss_d_fake], feed_dict={x_p: x, G_z_p: P, lod_in: cur_lod})\n",
    "            G_loss = sess.run(loss_g, feed_dict={z_p: z, G_z_p: P, lod_in: cur_lod})\n",
    "            print('num_img: %d ' % num_img, '  |  lod_in: %.2f' % cur_lod, '  |  D real loss: %.6f' % real_loss, '  |  D fake loss: %.6f' % fake_loss, '  |  Projection loss: %.6f' % G_loss)\n",
    "            gen_imgs = sess.run(G_z, feed_dict={z_p: z_fixed[:num_row*num_line], lod_in: cur_lod})\n",
    "            gen_imgs = (gen_imgs + 1) / 2\n",
    "            plt.imshow(montage(gen_imgs, grid=[num_row, num_line])[:, :, 0], cmap ='gray')\n",
    "            plt.axis('off')\n",
    "            plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LogD divergence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.Session() as sess:\n",
    "\n",
    "    iterators = [data_tool.data_iterator('mnist', lod_in=lod, batch_size=minibatch_dict[2**(resolution_log2-lod)], resolution_log2=resolution_log2) for lod in range(int(np.log2(resolution/init_resolution))+1)]\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    num_img = 0\n",
    "    tick_kimg = 0\n",
    "    z_fixed = np.random.randn(num_row*num_line, z_dim)\n",
    "    prev_lod = -1.0\n",
    "\n",
    "    while num_img <= total_nimg:\n",
    "        cur_lod = lod(num_img)\n",
    "        # reset Adam optimizers states when increasing resolution:\n",
    "        if np.floor(cur_lod) != np.floor(prev_lod) or np.ceil(cur_lod) != np.ceil(prev_lod):\n",
    "            sess.run([reset_optimizer_d, reset_optimizer_g])\n",
    "        prev_lod = cur_lod\n",
    "        # get mini-batch size:\n",
    "        batch_size = minibatch_dict[2**int(resolution_log2-np.floor(cur_lod))]\n",
    "\n",
    "        # sample a latent pool and get particles:\n",
    "        z = np.random.randn(batch_size, z_dim)\n",
    "        P = sess.run(G_z, feed_dict={z_p: z, lod_in: cur_lod})\n",
    "\n",
    "        # optimize discriminator:\n",
    "        x = next(iterators[int(np.floor(cur_lod))])\n",
    "        x = process_real(x, cur_lod)\n",
    "        num_img += batch_size # count\n",
    "        sess.run(update_d, feed_dict={x_p: x, G_z_p: P, lod_in: cur_lod})\n",
    "\n",
    "        # move particles\n",
    "        d_score = sess.run(d_fake_logits, feed_dict={G_z_p: P, lod_in: cur_lod})\n",
    "        grad = sess.run(d_grad, feed_dict={G_z_p: P, lod_in: cur_lod})\n",
    "        P += coef_div(d_score, div='LogD') * grad\n",
    "\n",
    "        # optimize generator:\n",
    "        sess.run(update_g, feed_dict={z_p: z, G_z_p: P, lod_in: cur_lod})\n",
    "\n",
    "        if (num_img // 1000) >= tick_kimg + 150:\n",
    "            tick_kimg = (num_img // 1000)\n",
    "            cur_lod = lod(num_img)\n",
    "            real_loss, fake_loss = sess.run([loss_d_real, loss_d_fake], feed_dict={x_p: x, G_z_p: P, lod_in: cur_lod})\n",
    "            G_loss = sess.run(loss_g, feed_dict={z_p: z, G_z_p: P, lod_in: cur_lod})\n",
    "            print('num_img: %d ' % num_img, '  |  lod_in: %.2f' % cur_lod, '  |  D real loss: %.6f' % real_loss, '  |  D fake loss: %.6f' % fake_loss, '  |  Projection loss: %.6f' % G_loss)\n",
    "            gen_imgs = sess.run(G_z, feed_dict={z_p: z_fixed[:num_row*num_line], lod_in: cur_lod})\n",
    "            gen_imgs = (gen_imgs + 1) / 2\n",
    "            plt.imshow(montage(gen_imgs, grid=[num_row, num_line])[:, :, 0], cmap ='gray')\n",
    "            plt.axis('off')\n",
    "            plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# JS divergence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.Session() as sess:\n",
    "\n",
    "    iterators = [data_tool.data_iterator('mnist', lod_in=lod, batch_size=minibatch_dict[2**(resolution_log2-lod)], resolution_log2=resolution_log2) for lod in range(int(np.log2(resolution/init_resolution))+1)]\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    num_img = 0\n",
    "    tick_kimg = 0\n",
    "    z_fixed = np.random.randn(num_row*num_line, z_dim)\n",
    "    prev_lod = -1.0\n",
    "\n",
    "    while num_img <= total_nimg:\n",
    "        cur_lod = lod(num_img)\n",
    "        # reset Adam optimizers states when increasing resolution:\n",
    "        if np.floor(cur_lod) != np.floor(prev_lod) or np.ceil(cur_lod) != np.ceil(prev_lod):\n",
    "            sess.run([reset_optimizer_d, reset_optimizer_g])\n",
    "        prev_lod = cur_lod\n",
    "        # get mini-batch size:\n",
    "        batch_size = minibatch_dict[2**int(resolution_log2-np.floor(cur_lod))]\n",
    "\n",
    "        # sample a latent pool and get particles:\n",
    "        z = np.random.randn(batch_size, z_dim)\n",
    "        P = sess.run(G_z, feed_dict={z_p: z, lod_in: cur_lod})\n",
    "\n",
    "        # optimize discriminator:\n",
    "        x = next(iterators[int(np.floor(cur_lod))])\n",
    "        x = process_real(x, cur_lod)\n",
    "        num_img += batch_size # count\n",
    "        sess.run(update_d, feed_dict={x_p: x, G_z_p: P, lod_in: cur_lod})\n",
    "\n",
    "        # move particles\n",
    "        d_score = sess.run(d_fake_logits, feed_dict={G_z_p: P, lod_in: cur_lod})\n",
    "        grad = sess.run(d_grad, feed_dict={G_z_p: P, lod_in: cur_lod})\n",
    "        P += coef_div(d_score, div='JS') * grad\n",
    "\n",
    "        # optimize generator:\n",
    "        sess.run(update_g, feed_dict={z_p: z, G_z_p: P, lod_in: cur_lod})\n",
    "\n",
    "        if (num_img // 1000) >= tick_kimg + 150:\n",
    "            tick_kimg = (num_img // 1000)\n",
    "            cur_lod = lod(num_img)\n",
    "            real_loss, fake_loss = sess.run([loss_d_real, loss_d_fake], feed_dict={x_p: x, G_z_p: P, lod_in: cur_lod})\n",
    "            G_loss = sess.run(loss_g, feed_dict={z_p: z, G_z_p: P, lod_in: cur_lod})\n",
    "            print('num_img: %d ' % num_img, '  |  lod_in: %.2f' % cur_lod, '  |  D real loss: %.6f' % real_loss, '  |  D fake loss: %.6f' % fake_loss, '  |  Projection loss: %.6f' % G_loss)\n",
    "            gen_imgs = sess.run(G_z, feed_dict={z_p: z_fixed[:num_row*num_line], lod_in: cur_lod})\n",
    "            gen_imgs = (gen_imgs + 1) / 2\n",
    "            plt.imshow(montage(gen_imgs, grid=[num_row, num_line])[:, :, 0], cmap ='gray')\n",
    "            plt.axis('off')\n",
    "            plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Jef divergence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with tf.Session() as sess:\n",
    "\n",
    "    iterators = [data_tool.data_iterator('mnist', lod_in=lod, batch_size=minibatch_dict[2**(resolution_log2-lod)], resolution_log2=resolution_log2) for lod in range(int(np.log2(resolution/init_resolution))+1)]\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    num_img = 0\n",
    "    tick_kimg = 0\n",
    "    z_fixed = np.random.randn(num_row*num_line, z_dim)\n",
    "    prev_lod = -1.0\n",
    "\n",
    "    while num_img <= total_nimg:\n",
    "        cur_lod = lod(num_img)\n",
    "        # reset Adam optimizers states when increasing resolution:\n",
    "        if np.floor(cur_lod) != np.floor(prev_lod) or np.ceil(cur_lod) != np.ceil(prev_lod):\n",
    "            sess.run([reset_optimizer_d, reset_optimizer_g])\n",
    "        prev_lod = cur_lod\n",
    "        # get mini-batch size:\n",
    "        batch_size = minibatch_dict[2**int(resolution_log2-np.floor(cur_lod))]\n",
    "\n",
    "        # sample a latent pool and get particles:\n",
    "        z = np.random.randn(batch_size, z_dim)\n",
    "        P = sess.run(G_z, feed_dict={z_p: z, lod_in: cur_lod})\n",
    "\n",
    "        # optimize discriminator:\n",
    "        x = next(iterators[int(np.floor(cur_lod))])\n",
    "        x = process_real(x, cur_lod)\n",
    "        num_img += batch_size # count\n",
    "        sess.run(update_d, feed_dict={x_p: x, G_z_p: P, lod_in: cur_lod})\n",
    "\n",
    "        # move particles\n",
    "        d_score = sess.run(d_fake_logits, feed_dict={G_z_p: P, lod_in: cur_lod})\n",
    "        grad = sess.run(d_grad, feed_dict={G_z_p: P, lod_in: cur_lod})\n",
    "        P += coef_div(d_score, div='Jef') * grad\n",
    "\n",
    "        # optimize generator:\n",
    "        sess.run(update_g, feed_dict={z_p: z, G_z_p: P, lod_in: cur_lod})\n",
    "\n",
    "        if (num_img // 1000) >= tick_kimg + 150:\n",
    "            tick_kimg = (num_img // 1000)\n",
    "            cur_lod = lod(num_img)\n",
    "            real_loss, fake_loss = sess.run([loss_d_real, loss_d_fake], feed_dict={x_p: x, G_z_p: P, lod_in: cur_lod})\n",
    "            G_loss = sess.run(loss_g, feed_dict={z_p: z, G_z_p: P, lod_in: cur_lod})\n",
    "            print('num_img: %d ' % num_img, '  |  lod_in: %.2f' % cur_lod, '  |  D real loss: %.6f' % real_loss, '  |  D fake loss: %.6f' % fake_loss, '  |  Projection loss: %.6f' % G_loss)\n",
    "            gen_imgs = sess.run(G_z, feed_dict={z_p: z_fixed[:num_row*num_line], lod_in: cur_lod})\n",
    "            gen_imgs = (gen_imgs + 1) / 2\n",
    "            plt.imshow(montage(gen_imgs, grid=[num_row, num_line])[:, :, 0], cmap ='gray')\n",
    "            plt.axis('off')\n",
    "            plt.show()"
   ]
  }
 ],
 "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
}