notebooks/06_02_custom_rnn.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow.python.keras.datasets import mnist\n",
    "from tensorflow.contrib.eager.python import tfe\n",
    "\n",
    "# Import the BasicLSTM written in TF Eager\n",
    "from utils.basic_lstm import BasicLSTM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# enable eager mode\n",
    "tf.enable_eager_execution()\n",
    "tf.set_random_seed(0)\n",
    "np.random.seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists('weights/'):\n",
    "    os.makedirs('weights/')\n",
    "\n",
    "# constants\n",
    "units = 128\n",
    "batch_size = 100\n",
    "epochs = 2\n",
    "num_classes = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x train (60000, 28, 28)\n",
      "y train (60000, 10)\n",
      "x test (10000, 28, 28)\n",
      "y test (10000, 10)\n"
     ]
    }
   ],
   "source": [
    "# dataset loading\n",
    "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
    "x_train = x_train.astype('float32') / 255.\n",
    "x_test = x_test.astype('float32') / 255.\n",
    "x_train = x_train.reshape((-1, 28, 28))  # 28 timesteps, 28 inputs / timestep\n",
    "x_test = x_test.reshape((-1, 28, 28))  # 28 timesteps, 28 inputs / timeste\n",
    "\n",
    "# one hot encode the labels. convert back to numpy as we cannot use a combination of numpy\n",
    "# and tensors as input to keras\n",
    "y_train_ohe = tf.one_hot(y_train, depth=num_classes).numpy()\n",
    "y_test_ohe = tf.one_hot(y_test, depth=num_classes).numpy()\n",
    "\n",
    "print('x train', x_train.shape)\n",
    "print('y train', y_train_ohe.shape)\n",
    "print('x test', x_test.shape)\n",
    "print('y test', y_test_ohe.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# What is the `BasicLSTM` Model \n",
    "\n",
    "The earlier model uses the LSTM model from Keras and this model uses the `BasicLSTM` Model written in TF Eager style code, which more or less replicates the important parts of the Keras LSTMCell.\n",
    "\n",
    "For some reason, there is a noticeable speed difference between the two models. Perhaps it is due to the usage of K.rnn() internally (inside Keras RNN, the base class of all RNNs) which is causing the slowdown.\n",
    "\n",
    "In comparison, the `BasicLSTM` simply loops over the batch of data in a more pythonic way. It can be found in the utils folder, and is posted as a code snippet here. \n",
    "\n",
    "```python\n",
    "import tensorflow as tf\n",
    "\n",
    "\n",
    "class BasicLSTM(tf.keras.Model):\n",
    "    def __init__(self, units, return_sequence=False, return_states=False, **kwargs):\n",
    "        super(BasicLSTM, self).__init__(**kwargs)\n",
    "        self.units = units\n",
    "        self.return_sequence = return_sequence\n",
    "        self.return_states = return_states\n",
    "\n",
    "        def bias_initializer(_, *args, **kwargs):\n",
    "            # Unit forget bias from the paper\n",
    "            # - [Learning to forget: Continual prediction with LSTM](http://www.mitpressjournals.org/doi/pdf/10.1162/089976600300015015)\n",
    "            return tf.keras.backend.concatenate([\n",
    "                tf.keras.initializers.Zeros()((self.units,), *args, **kwargs),  # input gate\n",
    "                tf.keras.initializers.Ones()((self.units,), *args, **kwargs),  # forget gate\n",
    "                tf.keras.initializers.Zeros()((self.units * 2,), *args, **kwargs),  # context and output gates\n",
    "            ])\n",
    "\n",
    "        self.kernel = tf.keras.layers.Dense(4 * units, use_bias=False)\n",
    "        self.recurrent_kernel = tf.keras.layers.Dense(4 * units, kernel_initializer='glorot_uniform', bias_initializer=bias_initializer)\n",
    "\n",
    "    def call(self, inputs, training=None, mask=None, initial_states=None):\n",
    "        # LSTM Cell in pure TF Eager code\n",
    "        # reset the states initially if not provided, else use those\n",
    "        if initial_states is None:\n",
    "            h_state = tf.zeros((inputs.shape[0], self.units))\n",
    "            c_state = tf.zeros((inputs.shape[0], self.units))\n",
    "        else:\n",
    "            assert len(initial_states) == 2, \"Must pass a list of 2 states when passing 'initial_states'\"\n",
    "            h_state, c_state = initial_states\n",
    "\n",
    "        h_list = []\n",
    "        c_list = []\n",
    "\n",
    "        for t in range(inputs.shape[1]):\n",
    "            # LSTM gate steps\n",
    "            ip = inputs[:, t, :]\n",
    "            z = self.kernel(ip)\n",
    "            z += self.recurrent_kernel(h_state)\n",
    "\n",
    "            z0 = z[:, :self.units]\n",
    "            z1 = z[:, self.units: 2 * self.units]\n",
    "            z2 = z[:, 2 * self.units: 3 * self.units]\n",
    "            z3 = z[:, 3 * self.units:]\n",
    "\n",
    "            # gate updates\n",
    "            i = tf.keras.activations.sigmoid(z0)\n",
    "            f = tf.keras.activations.sigmoid(z1)\n",
    "            c = f * c_state + i * tf.nn.tanh(z2)\n",
    "\n",
    "            # state updates\n",
    "            o = tf.keras.activations.sigmoid(z3)\n",
    "            h = o * tf.nn.tanh(c)\n",
    "\n",
    "            h_state = h\n",
    "            c_state = c\n",
    "\n",
    "            h_list.append(h_state)\n",
    "            c_list.append(c_state)\n",
    "\n",
    "        hidden_outputs = tf.stack(h_list, axis=1)\n",
    "        hidden_states = tf.stack(c_list, axis=1)\n",
    "\n",
    "        if self.return_states and self.return_sequence:\n",
    "            return hidden_outputs, [hidden_outputs, hidden_states]\n",
    "        elif self.return_states and not self.return_sequence:\n",
    "            return hidden_outputs[:, -1, :], [h_state, c_state]\n",
    "        elif self.return_sequence and not self.return_states:\n",
    "            return hidden_outputs\n",
    "        else:\n",
    "            return hidden_outputs[:, -1, :]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BasicLSTMModel(tf.keras.Model):\n",
    "    def __init__(self, units, num_classes):\n",
    "        super(BasicLSTMModel, self).__init__()\n",
    "        self.units = units\n",
    "        self.lstm = BasicLSTM(units)\n",
    "        self.classifier = tf.keras.layers.Dense(num_classes)\n",
    "\n",
    "    def call(self, inputs, training=None, mask=None):\n",
    "        h = self.lstm(inputs)\n",
    "        output = self.classifier(h)\n",
    "\n",
    "        # softmax op does not exist on the gpu, so always use cpu\n",
    "        with tf.device('/cpu:0'):\n",
    "            output = tf.nn.softmax(output)\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 60000 samples, validate on 10000 samples\n",
      "Epoch 1/2\n",
      "60000/60000 [==============================] - 59s 989us/step - loss: 0.2877 - acc: 0.9087 - val_loss: 0.1071 - val_acc: 0.9685\n",
      "Epoch 2/2\n",
      "60000/60000 [==============================] - 59s 988us/step - loss: 0.0884 - acc: 0.9741 - val_loss: 0.0933 - val_acc: 0.9722\n",
      "10000/10000 [==============================] - 4s 368us/step\n",
      "Final test loss and accuracy : [0.09331267344765365, 0.9722000068426132]\n"
     ]
    }
   ],
   "source": [
    "device = '/cpu:0' if tfe.num_gpus() == 0 else '/gpu:0'\n",
    "\n",
    "with tf.device(device):\n",
    "    # build model and optimizer\n",
    "    model = BasicLSTMModel(units, num_classes)\n",
    "    model.compile(optimizer=tf.train.AdamOptimizer(0.01), loss='categorical_crossentropy',\n",
    "                  metrics=['accuracy'])\n",
    "\n",
    "    # TF Keras tries to use entire dataset to determine shape without this step when using .fit()\n",
    "    # Fix = Use exactly one sample from the provided input dataset to determine input/output shape/s for the model\n",
    "    dummy_x = tf.zeros((1, 28, 28))\n",
    "    model._set_inputs(dummy_x)\n",
    "\n",
    "    # train\n",
    "    model.fit(x_train, y_train_ohe, batch_size=batch_size, epochs=epochs,\n",
    "              validation_data=(x_test, y_test_ohe), verbose=1)\n",
    "\n",
    "    # evaluate on test set\n",
    "    scores = model.evaluate(x_test, y_test_ohe, batch_size, verbose=1)\n",
    "    print(\"Final test loss and accuracy :\", scores)\n",
    "\n",
    "    saver = tfe.Saver(model.variables)\n",
    "    saver.save('weights/06_02_rnn/weights.ckpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:Anaconda3]",
   "language": "python",
   "name": "conda-env-Anaconda3-py"
  },
  "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}