Draw-Gumbol-softmax.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import argparse\n",
    "import sys\n",
    "\n",
    "from tensorflow.examples.tutorials.mnist import input_data\n",
    "import tensorflow as tf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "batch_size = 100\n",
    "img_sz = 28\n",
    "img_len = img_sz*img_sz\n",
    "n_itr = 10000\n",
    "n_class = 10\n",
    "lstm_size = 256\n",
    "T = 10\n",
    "A = img_sz\n",
    "B = img_sz\n",
    "eps = 1e-7\n",
    "N = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# calculate F_x and F_y\n",
    "def get_Fx_Fy(h):\n",
    "    params = tf.matmul(h, params_W) + params_b\n",
    "    g_x_, g_y_, log_var_Fxy, log_delta_, log_gamma = tf.split(params, num_or_size_splits=5, axis=1)\n",
    "    var_Fxy = tf.exp(log_var_Fxy)\n",
    "    gamma = tf.exp(log_gamma)\n",
    "    \n",
    "    g_x = ((A + 1) * 0.5) * (g_x_ + 1)\n",
    "    g_y = ((B + 1) * 0.5) * (g_y_ + 1)\n",
    "    delta = ((max(A, B) - 1) / (N - 1)) * tf.exp(log_delta_)\n",
    "\n",
    "    filter_i = tf.reshape(tf.cast(tf.range(N), tf.float32), [1, -1])\n",
    "    mu_x = g_x + (filter_i - N / 2 - 0.5) * delta\n",
    "    mu_y = g_y + (filter_i - N / 2 - 0.5) * delta\n",
    "\n",
    "    a = tf.reshape(tf.cast(tf.range(A), tf.float32), [1, 1, -1])\n",
    "    b = tf.reshape(tf.cast(tf.range(B), tf.float32), [1, 1, -1])\n",
    "    mu_x = tf.reshape(mu_x, [-1, N, 1])\n",
    "    mu_y = tf.reshape(mu_y, [-1, N, 1])\n",
    "    var_Fxy = tf.reshape(var_Fxy, [-1, 1, 1])    \n",
    "    \n",
    "    F_x = tf.exp(-tf.square((a - mu_x) / (2 * var_Fxy)))\n",
    "    F_y = tf.exp(-tf.square((b - mu_y) / (2 * var_Fxy))) \n",
    "    \n",
    "    F_x = F_x / tf.maximum(tf.reduce_sum(F_x, 2, keep_dims=True), eps)\n",
    "    F_y = F_y / tf.maximum(tf.reduce_sum(F_y, 2, keep_dims=True), eps)\n",
    "    return (F_x, F_y, gamma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def read(r, F_x, F_y, gamma):\n",
    "    # (F_x.T r F_y) * gamma\n",
    "    r = tf.reshape(r, [-1, B, A])\n",
    "    F_x_transpose = tf.transpose(F_x, perm=[0,2,1])\n",
    "    r = tf.matmul(F_y, tf.matmul(r, F_x_transpose))\n",
    "    r = tf.reshape(r, [-1, N*N])\n",
    "    r = r * tf.reshape(gamma,[-1,1])\n",
    "    return r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def write(w, F_x, F_y, gamma):\n",
    "    # (F_y.T w_ F_x) / gamma\n",
    "    w = tf.reshape(w, [batch_size, N, N])\n",
    "    F_y_transpose = tf.transpose(F_y, perm=[0,2,1])\n",
    "    w = tf.matmul(F_y_transpose, tf.matmul(w, F_x))\n",
    "    w = tf.reshape(w, [batch_size, B*A])\n",
    "    w = w * tf.reshape(1.0/gamma,[-1,1])\n",
    "    return w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = tf.placeholder(tf.float32, [batch_size, img_len])\n",
    "y_ = tf.placeholder(tf.float32, [batch_size, n_class])\n",
    "\n",
    "rnn_enc = tf.contrib.rnn.LSTMCell(lstm_size)\n",
    "rnn_dec = tf.contrib.rnn.LSTMCell(lstm_size)\n",
    "\n",
    "write_W = tf.Variable(tf.zeros([lstm_size, N*N]))\n",
    "write_b = tf.Variable(tf.zeros([N*N]))\n",
    "\n",
    "mu_W = tf.Variable(tf.zeros([lstm_size, n_class]))\n",
    "mu_b = tf.Variable(tf.zeros([n_class]))\n",
    "\n",
    "\n",
    "# learnable params for tau\n",
    "#log_tau_W = tf.Variable(tf.zeros([lstm_size, n_class]))\n",
    "#log_tau_b = tf.Variable(tf.zeros([n_class]))\n",
    "\n",
    "params_W = tf.Variable(tf.zeros([lstm_size, 5]))\n",
    "params_b = tf.Variable(tf.zeros([5]))\n",
    "\n",
    "lr = tf.placeholder(tf.float32, shape=[])\n",
    "tau = tf.placeholder(tf.float32, shape=[])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "c0 = tf.zeros([batch_size, img_len])\n",
    "#h0_enc = tf.zeros([batch_size, lstm_size])\n",
    "h0_dec = tf.zeros([batch_size, lstm_size])\n",
    "\n",
    "c = [0]*T\n",
    "#h_enc = h0_enc\n",
    "h_dec = h0_dec\n",
    "state_enc = rnn_enc.zero_state(batch_size, tf.float32) \n",
    "state_dec = rnn_dec.zero_state(batch_size, tf.float32) \n",
    "\n",
    "Lz_list = [0]*T\n",
    "\n",
    "F_x, F_y, gamma = get_Fx_Fy(h_dec)\n",
    "# read\n",
    "for t in range(T):\n",
    "    c_prev = c0 if t==0 else c[t-1]\n",
    "    #read\n",
    "    x_hat = x - tf.sigmoid(c_prev)\n",
    "    # no atten\n",
    "    #r = tf.concat(1,[x, x_hat])\n",
    "    r_x = read(x, F_x, F_y, gamma)\n",
    "    r_x_hat = read(x_hat, F_x, F_y, gamma)\n",
    "    r = tf.concat([r_x, r_x_hat], axis=1)\n",
    "    \n",
    "    #encoder\n",
    "    x_enc = tf.concat([r, h_dec], axis=1)\n",
    "    with tf.variable_scope(\"encoder\", reuse=(t != 0)):\n",
    "        h_enc, state_enc = rnn_enc(x_enc, state_enc)\n",
    "\n",
    "    # sample Q\n",
    "    mu_Q = tf.matmul(h_enc, mu_W) + mu_b\n",
    "    #log_tau_Q = tf.matmul(h_enc, log_tau_W) + log_tau_b\n",
    "    #tau_Q = tf.exp(log_tau_Q)\n",
    "    # Gumbol softmax distribution\n",
    "    U = tf.random_uniform([batch_size, n_class],minval=0,maxval=1)\n",
    "    # get gumble(0, 1) sample\n",
    "    # gumble has CDF as -ln(-ln(x))\n",
    "    # so indexing on CDF using a \n",
    "    # sample from uniform distribution \n",
    "    # will give the sample form gumble\n",
    "    G = -tf.log(-tf.log(U + eps) + eps)\n",
    "    # gubmle(logits, 1) sample\n",
    "    z_sample = mu_Q + G\n",
    "    # softmax approximation of one hot (categorical) distrubution\n",
    "    #z_sfmx = tf.nn.softmax(z_sample / tau_Q)\n",
    "    z_sfmx = tf.nn.softmax(z_sample / tau)\n",
    "    \n",
    "    # get the one hot encoding\n",
    "    z_one_hot = tf.cast(tf.equal(z_sfmx, tf.reduce_max(z_sfmx, 1, keep_dims=True)), z_sfmx.dtype)\n",
    "    \n",
    "    # staright through cat. distr.\n",
    "    # we need one hot encoding as output (z_one_hot)\n",
    "    # but gradient on z_sfmx\n",
    "    diff = tf.stop_gradient(z_one_hot - z_sfmx) \n",
    "    # zero gradient flow through this diff varibale \n",
    "    # affectively making variable diff a constant\n",
    "    # so gradient flow through z is same as z_sfmx\n",
    "    z = diff + z_sfmx\n",
    "    \n",
    "    # variational loss\n",
    "    # KL(Q(z|h_enc)||P(z)) = sum_over(Q(z|h_enc) * log (Q(z|h_enc)/P(z)))\n",
    "    # prior => P(z) = (1.0 / n_class)\n",
    "    # Gumbel-Softmax posterior condition on image => Q(z|h_enc) => approximated by z_sfmx\n",
    "    Lz_list[t] = z_sfmx * (tf.log(z_sfmx + eps) - tf.log(1.0 / n_class))\n",
    "                 \n",
    "    # decoder\n",
    "    with tf.variable_scope(\"decoder\", reuse=(t != 0)):\n",
    "        h_dec, state_dec = rnn_dec(z, state_dec) \n",
    "    \n",
    "    # no atten\n",
    "    w = tf.matmul(h_dec, write_W) + write_b\n",
    "        \n",
    "    #write       \n",
    "    F_x, F_y, gamma = get_Fx_Fy(h_dec)\n",
    "    w = write(w, F_x, F_y, gamma)\n",
    "                 \n",
    "    c[t] = c_prev + w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# reconstruction loss\n",
    "x_recons = tf.sigmoid(c[-1])\n",
    "# binary cross entropy\n",
    "Lx = tf.reduce_sum((x * tf.log(x_recons + eps) + (1.0 - x) * tf.log(1.0 - x_recons + eps)), 1)\n",
    "Lx = -tf.reduce_mean(Lx)\n",
    "\n",
    "#KL divergence\n",
    "Lz = tf.add_n(Lz_list)\n",
    "Lz = tf.reduce_mean(Lz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "loss = Lx + Lz\n",
    "optimizer = tf.train.AdamOptimizer(learning_rate=lr, beta1=0.5)\n",
    "grads_and_vars = optimizer.compute_gradients(loss)\n",
    "capped_grads_and_vars = [(tf.clip_by_norm(grad, 5.), var) \n",
    "                         if grad is not None else (grad, var) \n",
    "                         for grad, var in grads_and_vars]\n",
    "train_step = optimizer.apply_gradients(capped_grads_and_vars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "sess = tf.InteractiveSession()\n",
    "\n",
    "saver = tf.train.Saver()\n",
    "tf.global_variables_initializer().run()\n",
    "recon_loss = np.zeros(n_itr)\n",
    "variation_loss = np.zeros(n_itr)\n",
    "\n",
    "mnist = input_data.read_data_sets(\"MNIST_data/\", one_hot=True)\n",
    "tau0 = 1.0\n",
    "tau_val = tau0\n",
    "anneal_rate = 0.00003\n",
    "min_temp = 0.5\n",
    "lr_val = 1e-3 \n",
    "for i in range(n_itr):\n",
    "    batch_xs, _ = mnist.train.next_batch(batch_size)\n",
    "    recon_loss[i], variation_loss[i], _= sess.run([Lx, Lz, train_step], \n",
    "                                                  feed_dict={x: batch_xs, \n",
    "                                                             tau:tau_val, \n",
    "                                                             lr:lr_val})\n",
    "    \n",
    "    if i%100==0:\n",
    "        print(\"iter=%d : Lx: %f Lz: %f\" % (i,recon_loss[i], variation_loss[i]))\n",
    "    if i % 1000 == 1:\n",
    "        tau_val = np.maximum(tau0 * np.exp(-anneal_rate*i),min_temp)\n",
    "        tau_val *= 0.9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "batch_xs, _ = mnist.train.next_batch(batch_size)\n",
    "output_images = sess.run(c, feed_dict={x: batch_xs, tau:tau_val, lr:lr_val}) \n",
    "output_images = np.array(output_images) \n",
    "\n",
    "out_file=\"draw_data_gambol_sfmx.npy\"\n",
    "np.save(out_file,[output_images, recon_loss, variation_loss])\n",
    "print(\"Outputs saved in file: %s\" % out_file)\n",
    "\n",
    "ckpt_file= \"saved_models/draw_model_gambol_sfmx.ckpt\"\n",
    "print(\"Model saved in file: %s\" % saver.save(sess, ckpt_file))\n",
    "\n",
    "sess.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "import sys\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "def xrecons_grid(X,B,A):\n",
    "    padsize=1\n",
    "    padval=.5\n",
    "    ph=B+2*padsize\n",
    "    pw=A+2*padsize\n",
    "    batch_size=X.shape[0]\n",
    "    N=int(np.sqrt(batch_size))\n",
    "    X=X.reshape((N,N,B,A))\n",
    "    img=np.ones((N*ph,N*pw))*padval\n",
    "    for i in range(N):\n",
    "        for j in range(N):\n",
    "            startr=i*ph+padsize\n",
    "            endr=startr+B\n",
    "            startc=j*pw+padsize\n",
    "            endc=startc+A\n",
    "            img[startr:endr,startc:endc]=X[i,j,:,:]\n",
    "    return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "out_file=\"draw_data_gambol_sfmx.npy\"\n",
    "[C,Lxs,Lzs]=np.load(out_file)\n",
    "T,batch_size,img_size=C.shape\n",
    "X=1.0/(1.0+np.exp(-C)) \n",
    "B=A=int(np.sqrt(img_size))\n",
    "\n",
    "plt.plot(Lxs,label='Reconstruction Loss Lx')\n",
    "plt.plot(Lzs,label='Latent Loss Lz')\n",
    "plt.xlabel('iterations')\n",
    "plt.legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "for t in range(T):\n",
    "    img=xrecons_grid(X[t,:,:],B,A)\n",
    "    plt.matshow(img,cmap=plt.cm.gray)\n",
    "    ## convert -delay 10 -loop 0 *.png mnist.gif\n",
    "    #imgname='img_%d.png' % ('images', t) \n",
    "    #plt.savefig(imgname)\n",
    "    #print(imgname)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [conda root]",
   "language": "python",
   "name": "conda-root-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}