Edward_LSTM.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#!/usr/bin/env python\n",
    "\"\"\"LSTM language model on text8.\n",
    "Default hyperparameters achieve ~78.4 NLL at epoch 50, ~76.1423 NLL at\n",
    "epoch 200; ~13s/epoch on Titan X (Pascal).\n",
    "Samples after 200 epochs:\n",
    "```\n",
    "e the classmaker was cut apart rome the charts sometimes known a\n",
    "hemical place baining examples of equipment accepted manner clas\n",
    "uetean meeting sought to exist as this waiting an excerpt for of\n",
    "erally enjoyed a film writer of unto one two volunteer humphrey\n",
    "y captured by the saughton river goodness where stones were nota\n",
    "```\n",
    "\"\"\"\n",
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import edward as ed\n",
    "import numpy as np\n",
    "import os\n",
    "import string\n",
    "import tensorflow as tf\n",
    "\n",
    "from datetime import datetime\n",
    "from edward.models import Categorical\n",
    "from edward.util import Progbar\n",
    "from observations import text8\n",
    "\n",
    "data_dir = \"/tmp/data\"\n",
    "log_dir = \"/tmp/log\"\n",
    "n_epoch = 200\n",
    "batch_size = 128\n",
    "hidden_size = 512\n",
    "timesteps = 64\n",
    "lr = 5e-3\n",
    "\n",
    "timestamp = datetime.strftime(datetime.utcnow(), \"%Y%m%d_%H%M%S\")\n",
    "hyperparam_str = '_'.join([\n",
    "    var + '_' + str(eval(var)).replace('.', '_')\n",
    "    for var in ['batch_size', 'hidden_size', 'timesteps', 'lr']])\n",
    "log_dir = os.path.join(log_dir, timestamp + '_' + hyperparam_str)\n",
    "if not os.path.exists(log_dir):\n",
    "  os.makedirs(log_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def lstm_cell(x, h, c, name=None, reuse=False):\n",
    "  \"\"\"LSTM returning hidden state and content cell at a specific timestep.\"\"\"\n",
    "  nin = x.shape[-1].value\n",
    "  nout = h.shape[-1].value\n",
    "  with tf.variable_scope(name, default_name=\"lstm\",\n",
    "                         values=[x, h, c], reuse=reuse):\n",
    "    wx = tf.get_variable(\"kernel/input\", [nin, nout * 4],\n",
    "                         dtype=tf.float32,\n",
    "                         initializer=tf.orthogonal_initializer(1.0))\n",
    "    wh = tf.get_variable(\"kernel/hidden\", [nout, nout * 4],\n",
    "                         dtype=tf.float32,\n",
    "                         initializer=tf.orthogonal_initializer(1.0))\n",
    "    b = tf.get_variable(\"bias\", [nout * 4],\n",
    "                        dtype=tf.float32,\n",
    "                        initializer=tf.constant_initializer(0.0))\n",
    "\n",
    "  z = tf.matmul(x, wx) + tf.matmul(h, wh) + b\n",
    "  i, f, o, u = tf.split(z, 4, axis=1)\n",
    "  i = tf.sigmoid(i)\n",
    "  f = tf.sigmoid(f + 1.0)\n",
    "  o = tf.sigmoid(o)\n",
    "  u = tf.tanh(u)\n",
    "  c = f * c + i * u\n",
    "  h = o * tf.tanh(c)\n",
    "  return h, c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def generator(input, batch_size, timesteps, encoder):\n",
    "  \"\"\"Generate batch with respect to input (a list). Encode its\n",
    "  strings to integers, returning an array of shape [batch_size, timesteps].\n",
    "  \"\"\"\n",
    "  while True:\n",
    "    imb = np.random.randint(0, len(input) - timesteps, batch_size)\n",
    "    encoded = np.asarray(\n",
    "        [[encoder[c] for c in input[i:(i + timesteps)]] for i in imb],\n",
    "        dtype=np.int32)\n",
    "    yield encoded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def language_model(input):\n",
    "  \"\"\"Form p(x[0], ..., x[timesteps - 1]),\n",
    "  \\prod_{t=0}^{timesteps - 1} p(x[t] | x[:t]),\n",
    "  To calculate the probability, we call log_prob on\n",
    "  x = [x[0], ..., x[timesteps - 1]] given\n",
    "  `input` = [0, x[0], ..., x[timesteps - 2]].\n",
    "  We implement this separately from the generative model so the\n",
    "  forward pass, e.g., embedding/dense layers, can be parallelized.\n",
    "  [batch_size, timesteps] -> [batch_size, timesteps]\n",
    "  \"\"\"\n",
    "  x = tf.one_hot(input, depth=vocab_size, dtype=tf.float32)\n",
    "  h = tf.fill(tf.stack([tf.shape(x)[0], hidden_size]), 0.0)\n",
    "  c = tf.fill(tf.stack([tf.shape(x)[0], hidden_size]), 0.0)\n",
    "  hs = []\n",
    "  reuse = None\n",
    "  for t in range(timesteps):\n",
    "    if t > 0:\n",
    "      reuse = True\n",
    "    xt = x[:, t, :]\n",
    "    h, c = lstm_cell(xt, h, c, name=\"lstm\", reuse=reuse)\n",
    "    hs.append(h)\n",
    "\n",
    "  h = tf.stack(hs, 1)\n",
    "  logits = tf.layers.dense(h, vocab_size, name=\"dense\")\n",
    "  output = Categorical(logits=logits)\n",
    "  return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def language_model_gen(batch_size):\n",
    "  \"\"\"Generate x ~ prod p(x_t | x_{<t}). Output [batch_size, timesteps].\n",
    "  \"\"\"\n",
    "  # Initialize data input randomly.\n",
    "  x = tf.random_uniform([batch_size], 0, vocab_size, dtype=tf.int32)\n",
    "  h = tf.zeros([batch_size, hidden_size])\n",
    "  c = tf.zeros([batch_size, hidden_size])\n",
    "  xs = []\n",
    "  for _ in range(timesteps):\n",
    "    x = tf.one_hot(x, depth=vocab_size, dtype=tf.float32)\n",
    "    h, c = lstm_cell(x, h, c, name=\"lstm\")\n",
    "    logits = tf.layers.dense(h, vocab_size, name=\"dense\")\n",
    "    x = Categorical(logits=logits).value()\n",
    "    xs.append(x)\n",
    "\n",
    "  xs = tf.cast(tf.stack(xs, 1), tf.int32)\n",
    "  return xs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ed.set_seed(42)\n",
    "\n",
    "# DATA\n",
    "x_train, _, x_test = text8(data_dir)\n",
    "vocab = string.ascii_lowercase + ' '\n",
    "vocab_size = len(vocab)\n",
    "encoder = dict(zip(vocab, range(vocab_size)))\n",
    "decoder = {v: k for k, v in encoder.items()}\n",
    "\n",
    "data = generator(x_train, batch_size, timesteps, encoder)\n",
    "\n",
    "# MODEL\n",
    "x_ph = tf.placeholder(tf.int32, [None, timesteps])\n",
    "with tf.variable_scope(\"language_model\"):\n",
    "  # Shift input sequence to right by 1, [0, x[0], ..., x[timesteps - 2]].\n",
    "  x_ph_shift = tf.pad(x_ph, [[0, 0], [1, 0]])[:, :-1]\n",
    "  x = language_model(x_ph_shift)\n",
    "\n",
    "with tf.variable_scope(\"language_model\", reuse=True):\n",
    "  x_gen = language_model_gen(5)\n",
    "\n",
    "imb = range(0, len(x_test) - timesteps, timesteps)\n",
    "encoded_x_test = np.asarray(\n",
    "    [[encoder[c] for c in x_test[i:(i + timesteps)]] for i in imb],\n",
    "    dtype=np.int32)\n",
    "test_size = encoded_x_test.shape[0]\n",
    "print(\"Test set shape: {}\".format(encoded_x_test.shape))\n",
    "test_nll = -tf.reduce_sum(x.log_prob(x_ph))\n",
    "\n",
    "# INFERENCE\n",
    "inference = ed.MAP({}, {x: x_ph})\n",
    "\n",
    "optimizer = tf.train.AdamOptimizer(learning_rate=lr)\n",
    "inference.initialize(optimizer=optimizer, logdir=log_dir)\n",
    "\n",
    "print(\"Number of sets of parameters: {}\".format(len(tf.trainable_variables())))\n",
    "print(\"Number of parameters: {}\".format(\n",
    "    np.sum([np.prod(v.shape.as_list()) for v in tf.trainable_variables()])))\n",
    "for v in tf.trainable_variables():\n",
    "  print(v)\n",
    "\n",
    "sess = ed.get_session()\n",
    "tf.global_variables_initializer().run()\n",
    "\n",
    "# Double n_epoch and print progress every half an epoch.\n",
    "n_iter_per_epoch = len(x_train) // (batch_size * timesteps * 2)\n",
    "epoch = 0.0\n",
    "for _ in range(n_epoch * 2):\n",
    "  epoch += 0.5\n",
    "  print(\"Epoch: {0}\".format(epoch))\n",
    "  avg_nll = 0.0\n",
    "\n",
    "  pbar = Progbar(n_iter_per_epoch)\n",
    "  for t in range(1, n_iter_per_epoch + 1):\n",
    "    pbar.update(t)\n",
    "    x_batch = next(data)\n",
    "    info_dict = inference.update({x_ph: x_batch})\n",
    "    avg_nll += info_dict['loss']\n",
    "\n",
    "  # Print average bits per character over epoch.\n",
    "  avg_nll /= (n_iter_per_epoch * batch_size * timesteps * np.log(2))\n",
    "  print(\"Train average bits/char: {:0.8f}\".format(avg_nll))\n",
    "\n",
    "  # Print per-data point log-likelihood on test set.\n",
    "  avg_nll = 0.0\n",
    "  for start in range(0, test_size, batch_size):\n",
    "    end = min(test_size, start + batch_size)\n",
    "    x_batch = encoded_x_test[start:end]\n",
    "    avg_nll += sess.run(test_nll, {x_ph: x_batch})\n",
    "\n",
    "  avg_nll /= test_size\n",
    "  print(\"Test average NLL: {:0.8f}\".format(avg_nll))\n",
    "\n",
    "  # Generate samples from model.\n",
    "  samples = sess.run(x_gen)\n",
    "  samples = [''.join([decoder[c] for c in sample]) for sample in samples]\n",
    "  print(\"Samples:\")\n",
    "  for sample in samples:\n",
    "    print(sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}