recursive_seq2seq_v2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Recursive Style Seq2Seq Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow.contrib.rnn import LSTMCell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "NUM_INPUT_UNITS = 1\n",
    "NUM_STATE_INPUT_UNITS = 11\n",
    "NUM_HIDDEN_UNITS = 7\n",
    "MAX_OUPUT_LEN = 3\n",
    "MAX_RESCURISVE_DEPTH = 3\n",
    "NUM_OUTPUTS = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Things I need to track in the recursive while loop\n",
    "- current recurrent depth of \"frame\"\n",
    "- current recurrent accum for the frame\n",
    "- the (frame_idx, time_idx) from where the current frame was spawned from\n",
    "- Can combine the recurrent trackers in the above: elements will look like (BATCH_SIZE, 2) of type float32\n",
    "    - elements will be (current depth, recursive accumulator) \n",
    "    - e.g., [[0.4, 0], [0.7, 1], [0.8, 2]]\n",
    "    \n",
    "- Can combine the frame_idx and time_idx in the above: elements will look like (BATCH_SIZE, 2) of type int32\n",
    "    - elements will be (frame_idx, time_idx) \n",
    "    - e.g., [[0.4, 0], [0.7, 1], [0.8, 2]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 223,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Reset graph\n",
    "tf.reset_default_graph()\n",
    "\n",
    "# Specify inputs\n",
    "inp_hidden = tf.placeholder(tf.float32, shape=(None, NUM_HIDDEN_UNITS))\n",
    "inp_cell = tf.placeholder(tf.float32, shape=(None, NUM_HIDDEN_UNITS))\n",
    "\n",
    "# Constant\n",
    "# decoding initial value\n",
    "start_sentinel = tf.one_hot(\n",
    "    tf.zeros(shape=(tf.shape(inp_hidden)[0],), dtype=tf.int32),\n",
    "    NUM_OUTPUTS+1, dtype=tf.float32\n",
    ")\n",
    "\n",
    "# \"Add 1 to col\" tensor\n",
    "add_1_tensor = tf.ones(shape=(tf.shape(inp_hidden)[0], 1), dtype=tf.float32)\n",
    "add_1_tensor = tf.concat([tf.zeros_like(add_1_tensor, dtype=tf.float32), add_1_tensor], axis=1)\n",
    "\n",
    "\n",
    "# initial values for recurrent accumulators\n",
    "rec_accum = tf.expand_dims(tf.zeros_like(inp_hidden[:, 0]), axis=-1)\n",
    "rec_count = tf.expand_dims(tf.zeros_like(inp_hidden[:, 0], dtype=tf.float32), axis=-1)\n",
    "\n",
    "# Create recurisve tensorarray holders\n",
    "rec_idx = tf.constant(0, dtype=tf.int32)\n",
    "recursive_ta = tf.TensorArray(tf.float32, size=1, dynamic_size=True)\n",
    "frame_ref_ta = tf.TensorArray(tf.int32, size=1, dynamic_size=True)\n",
    "\n",
    "# Initalize the recursive tensorarray\n",
    "recursive_ta = recursive_ta.write(rec_idx, tf.concat([rec_accum, rec_count], axis=-1))\n",
    "frame_ref_ta = frame_ref_ta.write(rec_idx, tf.constant([-1, -1]))\n",
    "\n",
    "# Initialize ouput array\n",
    "final_probs_ta = tf.TensorArray(tf.float32, size=1, dynamic_size=True)\n",
    "final_actions_ta = tf.TensorArray(tf.float32, size=1, dynamic_size=True)\n",
    "\n",
    "# ========================= #\n",
    "# Initialize Recurrent Cell #\n",
    "# ========================= #\n",
    "\n",
    "cell = LSTMCell(NUM_HIDDEN_UNITS)\n",
    "test_init = cell(start_sentinel, (inp_hidden, inp_cell))\n",
    "\n",
    "# =========================== #\n",
    "# Define Inner Function Calls #\n",
    "# =========================== #\n",
    "\n",
    "# full recurrent step including aciton probs\n",
    "def network(prev_output, states):\n",
    "\n",
    "    output, states = cell(prev_output, states)\n",
    "    \n",
    "    action_probs = tf.layers.dense(output, NUM_OUTPUTS, activation=tf.nn.softmax)\n",
    "    \n",
    "    return action_probs, output, states        \n",
    "\n",
    "# ==================== #\n",
    "# Inner While Loop Ops #\n",
    "# ==================== #\n",
    "\n",
    "def cond(time, prev_out, prev_recursive, probs_ta, actions_ta, recursive_ta, frame_ref_ta, *states):\n",
    "        return time <= MAX_OUPUT_LEN\n",
    "    \n",
    "def step(time, prev_out, prev_recursive, probs_ta, actions_ta, recursive_ta, frame_ref_ta, *states):\n",
    "    \n",
    "    # Call the lstm cell\n",
    "    action_probs, output, state_tuple = network(prev_out, states)\n",
    "    states = state_tuple.h, state_tuple.c\n",
    "\n",
    "    # out probs\n",
    "    action_max = tf.argmax(action_probs, axis=1, output_type=tf.float32,)\n",
    "    action_max_one_hot = tf.one_hot(action_max, depth=NUM_OUTPUTS+1)\n",
    "\n",
    "    # write the current action_prob output\n",
    "    probs_ta = probs_ta.write(time, action_probs)\n",
    "    actions_ta = actions_ta.write(time, action_max_one_hot)\n",
    "    \n",
    "    # update recursion metrics\n",
    "    nested_recursive = tf.where(\n",
    "        tf.logical_and(\n",
    "            tf.expand_dims(tf.greater(action_max, 0), 1),\n",
    "            tf.less_equal(prev_recursive[:, 1:], MAX_RESCURISVE_DEPTH)\n",
    "        ),\n",
    "        lambda: prev_recursive + add_1_tensor,\n",
    "        lambda: prev_recursive\n",
    "    )\n",
    "    \n",
    "    # Should we write new frames?\n",
    "    write_new_bool = tf.logical_and(\n",
    "        tf.reduce_any(tf.greater(action_max, 0)),\n",
    "        tf.reduce_any(tf.less(prev_recursive[:, 1], MAX_RESCURISVE_DEPTH))\n",
    "    )\n",
    "    \n",
    "    # For the current time step check to see if we need to spawn new signal trees\n",
    "    recursive_ta = tf.cond(\n",
    "        write_new_bool,\n",
    "        lambda: recursive_ta.write(recursive_ta.size(), nested_recursive),\n",
    "        lambda: recursive_ta\n",
    "    )\n",
    "    \n",
    "    # Write the frame reference\n",
    "    frame_ref_ta = tf.cond(\n",
    "        write_new_bool,\n",
    "        lambda: frame_ref_ta.write(frame_ref_ta.size(), tf.stack([rec_idx, time], axis=0)),\n",
    "        lambda: frame_ref_ta\n",
    "    )\n",
    "\n",
    "    return (\n",
    "        index+1, \n",
    "        action_max_one_hot, \n",
    "        prev_recursive, \n",
    "        probs_ta, \n",
    "        actions_ta,\n",
    "        recursive_ta,\n",
    "        frame_ref_ta\n",
    "    ) + tuple(states)\n",
    "\n",
    "\n",
    "# ==================== #\n",
    "# Outer While Loop Ops #\n",
    "# ==================== #\n",
    "\n",
    "def recursive_cond(rec_idx, recursive_ta, frame_ref_ta, final_probs_ta, final_actions_ta):\n",
    "    \n",
    "#     recursive_ta_idx = recursive_ta.read(rec_idx)\n",
    "    \n",
    "#     below_max_rec_depth = tf.reduce_any(\n",
    "#         tf.less(recursive_ta_idx[:, 0], MAX_RESCURISVE_DEPTH),\n",
    "#         axis=0\n",
    "#     )\n",
    "    \n",
    "#     frames_depleted = tf.greater(recursive_ta.size(), rec_idx)\n",
    "    \n",
    "#     return tf.logical_and(below_max_rec_depth, frames_depleted)\n",
    "    return tf.greater(recursive_ta.size(), rec_idx)\n",
    "\n",
    "def recursive_func(rec_idx, recursive_ta, frame_ref_ta, final_probs_ta, final_actions_ta):\n",
    "    \"\"\"\n",
    "    Takes in the input to a decoder, \n",
    "    \"\"\"\n",
    "    # Instantiate a time index\n",
    "    time = tf.constant(0, dtype=tf.int32)\n",
    "    \n",
    "    # Read in the current recursive state\n",
    "    prev_recursive = recursive_ta.read(rec_idx)\n",
    "    \n",
    "    # Create TensorArrays for the internal times\n",
    "    probs_ta = tf.TensorArray(tf.float32, size=1, dynamic_size=True)\n",
    "    actions_ta = tf.TensorArray(tf.float32, size=1, dynamic_size=True)\n",
    "    \n",
    "    # while loop\n",
    "    (\n",
    "        final_time,\n",
    "        final_action,\n",
    "        final_recurisve,\n",
    "        probs_ta, \n",
    "        actions_ta,\n",
    "        recursive_ta,\n",
    "        frame_ref_ta,\n",
    "        _,\n",
    "        _\n",
    "    ) = tf.while_loop(\n",
    "        cond,\n",
    "        step,\n",
    "        loop_vars=[\n",
    "            time,\n",
    "            start_sentinel, \n",
    "            prev_recursive,\n",
    "            probs_ta, \n",
    "            actions_ta,\n",
    "            recursive_ta,\n",
    "            frame_ref_ta,\n",
    "            inp_hidden, inp_cell\n",
    "        ]\n",
    "    )\n",
    "    \n",
    "    # Stack the probs and actions\n",
    "    probs = probs_ta.stack()\n",
    "    actions = actions_ta.stack()\n",
    "    \n",
    "    # Write the probs and actions \n",
    "    final_probs_ta = final_probs_ta.write(rec_idx, probs)\n",
    "    final_actions_ta = final_actions_ta.write(rec_idx, actions)\n",
    "    \n",
    "    return rec_idx + 1, recursive_ta, frame_ref_ta, final_probs_ta, final_actions_ta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 210,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "rta = recursive_ta.concat()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 211,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "sess = tf.InteractiveSession()\n",
    "\n",
    "tf.global_variables_initializer().run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# ita = tf.TensorArray(tf.float32, size=0, dynamic_size=True)\n",
    "int_ta = tf.TensorArray(tf.int32, size=0, dynamic_size=True)\n",
    "int_ta = int_ta.write(0, tf.constant([[1, 1]], dtype=tf.int32))\n",
    "int_ta = int_ta.write(1, -1*tf.ones(shape=(1, 2), dtype=tf.int32))\n",
    "int_ta = int_ta.write(int_ta.size(), -1*tf.ones(shape=(1, 2), dtype=tf.int32))\n",
    "int_ta = int_ta.write(int_ta.size(), -1*tf.ones(shape=(1, 2), dtype=tf.int32))\n",
    "\n",
    "# ita = ita.write(0, tf.random_normal(shape=(5, 2)))\n",
    "# ita = ita.write(1, tf.random_normal(shape=(5, 2)))\n",
    "# ita = ita.write(2, tf.random_normal(shape=(5, 2)))\n",
    "# ita = ita.write(3, tf.random_normal(shape=(5, 2)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 213,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 10\n",
    "res = sess.run(\n",
    "    [rec_accum, rec_count, rta, add_1_tensor],\n",
    "    feed_dict={\n",
    "        inp_hidden: np.random.rand(BATCH_SIZE, NUM_HIDDEN_UNITS),\n",
    "        inp_cell: np.random.rand(BATCH_SIZE, NUM_HIDDEN_UNITS)\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 215,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.,  1.],\n",
       "       [ 0.,  1.],\n",
       "       [ 0.,  1.],\n",
       "       [ 0.,  1.],\n",
       "       [ 0.,  1.],\n",
       "       [ 0.,  1.],\n",
       "       [ 0.,  1.],\n",
       "       [ 0.,  1.],\n",
       "       [ 0.,  1.],\n",
       "       [ 0.,  1.]], dtype=float32)"
      ]
     },
     "execution_count": 215,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res[-1]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tensorflow_venv",
   "language": "python",
   "name": "tensorflow_venv"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}