{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hackathon 8\n",
    "\n",
    "Topics:\n",
    "- TensorFlow RNNs and Cells\n",
    "- LSTM\n",
    "\n",
    "In today's demo, we'll teach an RNN how to speak English.\n",
    "\n",
    "This is all setup in a IPython notebook so you can run any code you want to experiment with. Feel free to edit any cell, or add some to run your own code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# we'll start with our library imports...\n",
    "# tensorflow to specify and run computation graphs\n",
    "# numpy to run any numerical operations that need to take place outside of the TF graph\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "# these ones let us draw images in our notebook\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### RNN/LSTM theory recap\n",
    "\n",
    "Recurrent neural networks (RNNs) are computation graphs with loops (i.e., not directed acyclic graphs). Because the backpropagation algorithm only works with DAGs, we have to unroll the RNN through time. Tensorflow provides code that handles this automatically.\n",
    "\n",
    "<img src=\"http://colah.github.io/posts/2015-08-Understanding-LSTMs/img/RNN-unrolled.png\">\n",
    "\n",
    "\n",
    "The most common RNN unit is the LSTM, depicted below:\n",
    "\n",
    "<img src=\"http://colah.github.io/posts/2015-08-Understanding-LSTMs/img/LSTM3-chain.png\">\n",
    "\n",
    "We can see that each unit takes 3 inputs and produces 3 outputs, two which are forwarded to the same unit at the next timestep and one true output, $h_t$ depicted coming out of the top of the cell.\n",
    "\n",
    "The upper output going to the next timestep is the cell state. It carries long-term information between cells, and is calculated as: \n",
    "\n",
    "<img src=http://colah.github.io/posts/2015-08-Understanding-LSTMs/img/LSTM3-focus-C.png>\n",
    "\n",
    "where the first term uses the forget gate $f_t$ to decide to scale the previous state (potentially making it smaller to \"forget\" it), and the second term is the product of the update gate $i_t$ and the state update $\\tilde{C}_t$. Each of the forget and update gates are activated with sigmoid, so their range is (0,1).\n",
    "\n",
    "The true output and the second, lower output on the diagram are calculated by the output gate:\n",
    "\n",
    "<img src=http://colah.github.io/posts/2015-08-Understanding-LSTMs/img/LSTM3-focus-o.png>\n",
    "\n",
    "First, $o_t$ is calculated from the output of the previous timestep concatenated with the current input, but then it's mixed with the cell state to get the true output. Passing on this output to the next timestep as the hidden state gives the unit a kind of short term memory.\n",
    "\n",
    "(Images sourced from [Colah's Blog](http://colah.github.io/posts/2015-08-Understanding-LSTMs/))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're going to use some Google code to load the [PTB dataset](http://www.fit.vutbr.cz/~imikolov/rnnlm/), a text dataset to teach the RNN to speak English."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ptb_reader\n",
    "\n",
    "TIME_STEPS = 20\n",
    "BATCH_SIZE = 20\n",
    "DATA_DIR = '/work/cse496dl/shared/hackathon/08/ptbdata'\n",
    "\n",
    "class PTBInput(object):\n",
    "  \"\"\"The input data.\n",
    "  \n",
    "  Code sourced from https://github.com/tensorflow/models/blob/master/tutorials/rnn/ptb/ptb_word_lm.py\n",
    "  \"\"\"\n",
    "\n",
    "  def __init__(self, data, batch_size, num_steps, name=None):\n",
    "    self.batch_size = batch_size\n",
    "    self.num_steps = num_steps\n",
    "    self.epoch_size = ((len(data) // batch_size) - 1) // num_steps\n",
    "    self.input_data, self.targets = ptb_reader.ptb_producer(\n",
    "        data, batch_size, num_steps, name=name)\n",
    "\n",
    "raw_data = ptb_reader.ptb_raw_data(DATA_DIR)\n",
    "train_data, valid_data, test_data, _ = raw_data\n",
    "train_input = PTBInput(train_data, BATCH_SIZE, TIME_STEPS, name=\"TrainInput\")\n",
    "print(\"The time distributed training data: \" + str(train_input.input_data))\n",
    "print(\"The similarly distributed targets: \" + str(train_input.targets))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each datum is a string of 20 successive words from the corpus, and the target is a similar window, but shifted forward by one word. This is setup to train the model to, given a few preceding words, predict what the next word in the sequence will be.\n",
    "\n",
    "Initially, in the data each word in the sequence is represented as an integer (notice the shape). This discrete representation fails to capture any semantic relationships between words. I.e., the model wouldn't know that \"crimson\" and \"scarlet\" are more similar than \"red\" and \"blue\". The solution is to learn an word embedding as the first part of the model to transform each integer into a relatively small, dense vector (as compared to a one-hot). Then, similar words will train to have similar embeddings.\n",
    "\n",
    "We'll use [tf.nn.embedding_lookup](https://www.tensorflow.org/api_docs/python/tf/nn/embedding_lookup) to do this which we provide a (usually trainable) VOCAB_SIZE x EMBEDDING_SIZE matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "VOCAB_SIZE = 10000\n",
    "EMBEDDING_SIZE = 100\n",
    "\n",
    "# setup input and embedding\n",
    "embedding_matrix = tf.get_variable('embedding_matrix', dtype=tf.float32, shape=[VOCAB_SIZE, EMBEDDING_SIZE], trainable=True)\n",
    "word_embeddings = tf.nn.embedding_lookup(embedding_matrix, train_input.input_data)\n",
    "print(\"The output of the word embedding: \" + str(word_embeddings))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TensorFlow separates the declaration of [RNNCells](https://www.tensorflow.org/api_docs/python/tf/contrib/rnn/RNNCell) from the [RNNs](https://www.tensorflow.org/api_docs/python/tf/nn/dynamic_rnn) that run them. In the code below, we declare an [LSTM cell](https://www.tensorflow.org/api_docs/python/tf/contrib/rnn/BasicLSTMCell), and create tensors for the inputs to the first unit. We use zeros for the initial hidden state and current state, but it's also possible to declare trainable variables for these as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LSTM_SIZE = 200 # number of units in the LSTM layer, this number taken from a \"small\" language model\n",
    "\n",
    "lstm_cell = tf.contrib.rnn.BasicLSTMCell(LSTM_SIZE)\n",
    "\n",
    "# Initial state of the LSTM memory.\n",
    "initial_state = lstm_cell.zero_state(BATCH_SIZE, tf.float32)\n",
    "print(\"Initial state of the LSTM: \" + str(initial_state))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we'll pass the newly declared cell and the training sequence of word embeddings to [tf.nn.dynamic_rnn](https://www.tensorflow.org/api_docs/python/tf/nn/dynamic_rnn) as the inputs over time to the LSTM. `dynamic_rnn` runs an `RNNCell` using an internal `while` loop, and returns the sequence of outputs from the LSTM at each timestep and the final state of the LSTM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup RNN\n",
    "outputs, state = tf.nn.dynamic_rnn(lstm_cell, word_embeddings,\n",
    "                                   initial_state=initial_state,\n",
    "                                   dtype=tf.float32)\n",
    "print(\"The outputs over all timesteps: \"+ str(outputs))\n",
    "print(\"The final state of the LSTM layer: \" + str(state))\n",
    "logits = tf.layers.dense(outputs, VOCAB_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And to calculate the loss between two sequences, we'll import a function from [tf.contrib.seq2seq](https://www.tensorflow.org/api_docs/python/tf/contrib/seq2seq) called [sequence_loss](https://www.tensorflow.org/api_docs/python/tf/contrib/seq2seq/sequence_loss). It calculates the weighted cross-entropy loss between the first two arguments, and the third argument provides weights for averaging. We weight uniformly here, but weights could also be calculated based on where in the sequence the target is (e.g., penalize less earlier in the sequence, but more later) or based on the content of the target (e.g., low weight on guessing articles correctly and larger weight on getting nouns and verbs correct).\n",
    "\n",
    "We'll optimize using TensorFlow's [RMSProp](https://www.tensorflow.org/api_docs/python/tf/train/RMSPropOptimizer) optimizer, which requires an explicit learning rate, but otherwise as usual. We switch from the Adam optmizer because we don't want the adaptive learning rate feature, which can interact badly with the recurrent gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "LEARNING_RATE = 1e-4\n",
    "\n",
    "loss = tf.contrib.seq2seq.sequence_loss(\n",
    "    logits,\n",
    "    train_input.targets,\n",
    "    tf.ones([BATCH_SIZE, TIME_STEPS], dtype=tf.float32),\n",
    "    average_across_timesteps=True,\n",
    "    average_across_batch=True)\n",
    "\n",
    "optimizer = tf.train.RMSPropOptimizer(LEARNING_RATE)\n",
    "train_op = optimizer.minimize(loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, before trying to run any operations, we need code to start the queue runners. The Google code we're using internally uses TensorFlow queues to load the data rather than input placeholders. This means that, rather than passing a `feed_dict` to each `run` call, we have to start the queue runners which spin up CPU threads to enqueue data. We declare a [Coordinator](https://www.tensorflow.org/api_docs/python/tf/train/Coordinator) and use it to call [tf.train.start_queue_runners](https://www.tensorflow.org/api_docs/python/tf/train/start_queue_runners) which does the job. If you're using someone's TF code and notice that the program hangs at the first `run` call, chances are that no queue_runner has been started (I've done this a few times in my own code...)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session = tf.Session()\n",
    "session.run(tf.global_variables_initializer())\n",
    "\n",
    "# start queue runners\n",
    "coord = tf.train.Coordinator()\n",
    "threads = tf.train.start_queue_runners(sess=session, coord=coord)\n",
    "\n",
    "# retrieve some data to look at\n",
    "examples = session.run([train_input.input_data, train_input.targets])\n",
    "# we can run the train op as usual\n",
    "_ = session.run(train_op)\n",
    "\n",
    "print(\"Example input data:\\n\" + str(examples[0][1]))\n",
    "print(\"Example target:\\n\" + str(examples[1][1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hackathon 8 Exercise\n",
    "\n",
    "Use a TF [Bidirectional RNN](https://www.tensorflow.org/api_docs/python/tf/contrib/rnn/stack_bidirectional_dynamic_rnn) to create a 2 layer [LSTMCell](https://www.tensorflow.org/api_docs/python/tf/contrib/rnn/BasicLSTMCell) with an [Attention Wrapper](https://www.tensorflow.org/api_docs/python/tf/contrib/rnn/AttentionCellWrapper), just so you can say you did it once. Your code should use `train_input` inputs as above, use different cells going forward and backward, and your code should finish with the `loss` tensor. This should be pretty easy with the TensorFlow documentation.\n",
    "\n",
    "This model is very large and trains for a long time, so please don't try to optimize it in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Your code here"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "TensorFlow 1.4 (py36)",
   "language": "python",
   "name": "tensorflow-1.4.0-py36"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
