DeepSpeech.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook we will reproduce the results of [Deep Speech: Scaling up end-to-end speech recognition](http://arxiv.org/abs/1412.5567). The core of the system is a bidirectional recurrent neural network (BRNN) trained to ingest speech spectrograms and generate English text transcriptions.\n",
    "\n",
    " Let a single utterance $x$ and label $y$ be sampled from a training set $S = \\{(x^{(1)}, y^{(1)}), (x^{(2)}, y^{(2)}), . . .\\}$. Each utterance, $x^{(i)}$ is a time-series of length $T^{(i)}$ where every time-slice is a vector of audio features, $x^{(i)}_t$ where $t=1,\\ldots,T^{(i)}$. We use MFCC as our features; so $x^{(i)}_{t,p}$ denotes the $p$-th MFCC feature in the audio frame at time $t$. The goal of our BRNN is to convert an input sequence $x$ into a sequence of character probabilities for the transcription $y$, with $\\hat{y}_t =\\mathbb{P}(c_t \\mid x)$, where $c_t \\in \\{a,b,c, . . . , z, space, apostrophe, blank\\}$. (The significance of $blank$ will be explained below.)\n",
    "\n",
    "Our BRNN model is composed of $5$ layers of hidden units. For an input $x$, the hidden units at layer $l$ are denoted $h^{(l)}$ with the convention that $h^{(0)}$ is the input. The first three layers are not recurrent. For the first layer, at each time $t$, the output depends on the MFCC frame $x_t$ along with a context of $C$ frames on each side. (We typically use $C \\in \\{5, 7, 9\\}$ for our experiments.) The remaining non-recurrent layers operate on independent data for each time step. Thus, for each time $t$, the first $3$ layers are computed by:\n",
    "\n",
    "$$h^{(l)}_t = g(W^{(l)} h^{(l-1)}_t + b^{(l)})$$\n",
    "\n",
    "where $g(z) = \\min\\{\\max\\{0, z\\}, 20\\}$ is a clipped rectified-linear (ReLu) activation function and $W^{(l)}$, $b^{(l)}$ are the weight matrix and bias parameters for layer $l$. The fourth layer is a bidirectional recurrent layer[[1](http://www.di.ufpe.br/~fnj/RNA/bibliografia/BRNN.pdf)]. This layer includes two sets of hidden units: a set with forward recurrence, $h^{(f)}$, and a set with backward recurrence $h^{(b)}$:\n",
    "\n",
    "$$h^{(f)}_t = g(W^{(4)} h^{(3)}_t + W^{(f)}_r h^{(f)}_{t-1} + b^{(4)})$$\n",
    "$$h^{(b)}_t = g(W^{(4)} h^{(3)}_t + W^{(b)}_r h^{(b)}_{t+1} + b^{(4)})$$\n",
    "\n",
    "Note that $h^{(f)}$ must be computed sequentially from $t = 1$ to $t = T^{(i)}$ for the $i$-th utterance, while\n",
    "the units $h^{(b)}$ must be computed sequentially in reverse from $t = T^{(i)}$ to $t = 1$.\n",
    "\n",
    "The fifth (non-recurrent) layer takes both the forward and backward units as inputs\n",
    "\n",
    "$$h^{(5)} = g(W^{(5)} h^{(4)} + b^{(5)})$$\n",
    "\n",
    "where $h^{(4)} = h^{(f)} + h^{(b)}$. The output layer are standard logits that correspond to the predicted character probabilities for each time slice $t$ and character $k$ in the alphabet:\n",
    "\n",
    "$$h^{(6)}_{t,k} = \\hat{y}_{t,k} = (W^{(6)} h^{(5)}_t)_k + b^{(6)}_k$$\n",
    "\n",
    "Here $b^{(6)}_k$ denotes the $k$-th bias and $(W^{(6)} h^{(5)}_t)_k$ the $k$-th element of the matrix product.\n",
    "\n",
    "Once we have computed a prediction for $\\hat{y}_{t,k}$, we compute the CTC loss[[2]](http://www.cs.toronto.edu/~graves/preprint.pdf) $\\cal{L}(\\hat{y}, y)$ to measure the error in prediction. During training, we can evaluate the gradient $\\nabla \\cal{L}(\\hat{y}, y)$ with respect to the network outputs given the ground-truth character sequence $y$. From this point, computing the gradient with respect to all of the model parameters may be done via back-propagation through the rest of the network. We use the Adam method for training[[3](http://arxiv.org/abs/1412.6980)].\n",
    "\n",
    "The complete BRNN model is illustrated in the figure below.\n",
    "\n",
    "![DeepSpeech BRNN](images/rnn_fig-624x548.png)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preliminaries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we first import all of the packages we require to implement the DeepSpeech BRNN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import time\n",
    "import json\n",
    "import shutil\n",
    "import datetime\n",
    "import tempfile\n",
    "import subprocess\n",
    "import numpy as np\n",
    "from math import ceil\n",
    "from xdg import BaseDirectory as xdg\n",
    "import tensorflow as tf\n",
    "from util.log import merge_logs\n",
    "from util.gpu import get_available_gpus\n",
    "from util.shared_lib import check_cupti\n",
    "from util.text import ctc_label_dense_to_sparse, sparse_tensor_value_to_texts, wers\n",
    "from tensorflow.python.ops import ctc_ops\n",
    "from tensorflow.contrib.session_bundle import exporter\n",
    "\n",
    "ds_importer = os.environ.get('ds_importer', 'ted')\n",
    "ds_dataset_path = os.environ.get('ds_dataset_path', os.path.join('./data', ds_importer))\n",
    "\n",
    "import importlib\n",
    "ds_importer_module = importlib.import_module('util.importers.%s' % ds_importer)\n",
    "\n",
    "from util.website import maybe_publish\n",
    "\n",
    "do_fulltrace = bool(len(os.environ.get('ds_do_fulltrace', '')))\n",
    "\n",
    "if do_fulltrace:\n",
    "    check_cupti()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Global Constants"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we introduce several constants used in the algorithm below.  In particular, we define\n",
    "* `learning_rate` - The learning rate we will employ in Adam optimizer[[3]](http://arxiv.org/abs/1412.6980)\n",
    "* `training_iters` - The number of iterations we will train for\n",
    "* `batch_size` - The number of elements in a batch\n",
    "* `display_step` - The number of epochs we cycle through before displaying progress\n",
    "* `checkpoint_step` - The number of epochs we cycle through before checkpointing the model\n",
    "* `checkpoint_dir` - The directory in which checkpoints are stored\n",
    "* `restore_checkpoint` - Whether to resume from checkpoints when training\n",
    "* `export_dir` - The directory in which exported models are stored\n",
    "* `export_version` - The version number of the exported model\n",
    "* `remove_export` - Whether to remove old exported models\n",
    "* `default_stddev` - The default standard deviation to use when initialising weights and biases\n",
    "* `[bh][12356]_stddev` - Individual standard deviations to use when initialising particular weights and biases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "learning_rate = float(os.environ.get('ds_learning_rate', 0.001)) # TODO: Determine a reasonable value for this\n",
    "beta1 = float(os.environ.get('ds_beta1', 0.9))                   # TODO: Determine a reasonable value for this\n",
    "beta2 = float(os.environ.get('ds_beta2', 0.999))                 # TODO: Determine a reasonable value for this\n",
    "epsilon = float(os.environ.get('ds_epsilon', 1e-8))              # TODO: Determine a reasonable value for this\n",
    "training_iters = int(os.environ.get('ds_training_iters', 15))    # TODO: Determine a reasonable value for this\n",
    "batch_size = int(os.environ.get('ds_batch_size', 64))            # TODO: Determine a reasonable value for this\n",
    "display_step = int(os.environ.get('ds_display_step', 1))         # TODO: Determine a reasonable value for this\n",
    "validation_step = int(os.environ.get('ds_validation_step', 1))   # TODO: Determine a reasonable value for this\n",
    "checkpoint_step = int(os.environ.get('ds_checkpoint_step', 5))   # TODO: Determine a reasonable value for this\n",
    "checkpoint_dir = os.environ.get('ds_checkpoint_dir', xdg.save_data_path('deepspeech'))\n",
    "restore_checkpoint = bool(int(os.environ.get('ds_restore_checkpoint', 0)))\n",
    "export_dir = os.environ.get('ds_export_dir', None)\n",
    "export_version = 1\n",
    "remove_export = bool(int(os.environ.get('ds_remove_export', 0)))\n",
    "use_warpctc = bool(len(os.environ.get('ds_use_warpctc', '')))\n",
    "default_stddev = float(os.environ.get('ds_default_stddev', 0.1))\n",
    "for var in ['b1', 'h1', 'b2', 'h2', 'b3', 'h3', 'b5', 'h5', 'b6', 'h6']:\n",
    "    locals()['%s_stddev' % var] = float(os.environ.get('ds_%s_stddev' % var, default_stddev))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we use the Adam optimizer[[3]](http://arxiv.org/abs/1412.6980) instead of Nesterov’s Accelerated Gradient [[4]](http://www.cs.utoronto.ca/~ilya/pubs/2013/1051_2.pdf) used in the original DeepSpeech paper, as, at the time of writing, TensorFlow does not have an implementation of Nesterov’s Accelerated Gradient [[4]](http://www.cs.utoronto.ca/~ilya/pubs/2013/1051_2.pdf).\n",
    "\n",
    "As we will also employ dropout on the feedforward layers of the network, we need to define a parameter `dropout_rate` that keeps track of the dropout rate for these layers.\n",
    "\n",
    "To avoid graph duplication when performing validation steps, we use a placeholder and an alternative feed_dict during validation to pull data from the corresponding queue."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "dropout_rate = float(os.environ.get('ds_dropout_rate', 0.05))  # TODO: Validate this is a reasonable value\n",
    "\n",
    "# This global placeholder will be used for all dropout definitions\n",
    "dropout_rate_placeholder = tf.placeholder(tf.float32)\n",
    "\n",
    "# This placeholder will be used to select between queues\n",
    "queue_selector_placeholder = tf.placeholder(tf.uint8)\n",
    "\n",
    "# The feed_dict used for training employs the given dropout_rate\n",
    "feed_dict_train = { dropout_rate_placeholder: dropout_rate,\n",
    "                    queue_selector_placeholder: 0 }\n",
    "\n",
    "# The feed dict used for validation employs zero dropout and selects the validation queue\n",
    "feed_dict_validate = { dropout_rate_placeholder: 0.0,\n",
    "                       queue_selector_placeholder: 1 }\n",
    "\n",
    "# While the feed_dict used for test reporting employs zero dropout and selects the test queue\n",
    "feed_dict_test = { dropout_rate_placeholder: 0.0,\n",
    "                   queue_selector_placeholder: 2 }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One more constant required of the non-recurrant layers is the clipping value of the ReLU. We capture that in the value of the variable `relu_clip`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "relu_clip = int(os.environ.get('ds_relu_clip', 20)) # TODO: Validate this is a reasonable value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Geometric Constants"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will introduce several constants related to the geometry of the network.\n",
    "\n",
    "The network views each speech sample as a sequence of time-slices $x^{(i)}_t$ of length $T^{(i)}$. As the speech samples vary in length, we know that $T^{(i)}$ need not equal $T^{(j)}$ for $i \\ne j$. For each batch, BRNN in TensorFlow needs to know `n_steps` which is the maximum $T^{(i)}$ for the batch."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each of the at maximum `n_steps` vectors is a vector of MFCC features of a time-slice of the speech sample. We will make the number of MFCC features dependent upon the sample rate of the data set. Generically, if the sample rate is 8kHz we use 13 features. If the sample rate is 16kHz we use 26 features... We capture the dimension of these vectors, equivalently the number of MFCC features, in the variable `n_input`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_input = 26 # TODO: Determine this programatically from the sample rate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As previously mentioned, the BRNN is not simply fed the MFCC features of a given time-slice. It is fed, in addition, a context of $C \\in \\{5, 7, 9\\}$ frames on either side of the frame in question. The number of frames in this context is captured in the variable `n_context`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_context = 9 # TODO: Determine the optimal value using a validation data set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we will introduce constants that specify the geometry of some of the non-recurrent layers of the network. We do this by simply specifying the number of units in each of the layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_hidden_1 = n_input + 2*n_input*n_context # Note: This value was not specified in the original paper\n",
    "n_hidden_2 = n_input + 2*n_input*n_context # Note: This value was not specified in the original paper\n",
    "n_hidden_5 = n_input + 2*n_input*n_context # Note: This value was not specified in the original paper"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "where `n_hidden_1` is the number of units in the first layer, `n_hidden_2` the number of units in the second, and  `n_hidden_5` the number in the fifth. We haven't forgotten about the third or sixth layer. We will define their unit count below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A LSTM BRNN consists of a pair of LSTM RNN's. One LSTM RNN that works \"forward in time\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"images/LSTM3-chain.png\" alt=\"LSTM\" width=\"800\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and a second LSTM RNN that works \"backwards in time\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"images/LSTM3-chain.png\" alt=\"LSTM\" width=\"800\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dimension of the cell state, the upper line connecting subsequent LSTM units, is independent of the input dimension and the same for both the forward and backward LSTM RNN.\n",
    "\n",
    "Hence, we are free to choose the dimension of this cell state independent of the input dimension. We capture the cell state dimension in the variable `n_cell_dim`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_cell_dim = n_input + 2*n_input*n_context # TODO: Is this a reasonable value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The number of units in the third layer, which feeds in to the LSTM, is determined by `n_cell_dim` as follows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_hidden_3 = 2 * n_cell_dim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we introduce an additional variable `n_character` which holds the number of characters in the target language plus one, for the $blamk$. For English it is the cardinality of the set $\\{a,b,c, . . . , z, space, apostrophe, blank\\}$ we referred to earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_character = 29 # TODO: Determine if this should be extended with other punctuation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The number of units in the sixth layer is determined by `n_character` as follows "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_hidden_6 = n_character"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Graph Creation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we concern ourselves with graph creation.\n",
    "\n",
    "However, before we do so we must introduce a utility function `variable_on_cpu()` used to create a variable in CPU memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def variable_on_cpu(name, shape, initializer):\n",
    "    # Use the /cpu:0 device for scoped operations\n",
    "    with tf.device('/cpu:0'):\n",
    "        # Create or get apropos variable\n",
    "        var = tf.get_variable(name=name, shape=shape, initializer=initializer)\n",
    "    return var"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That done, we will define the learned variables, the weights and biases, within the method `BiRNN()` which also constructs the neural network. The  variables named `hn`, where `n` is an integer, hold the learned weight variables. The variables named `bn`, where `n` is an integer, hold the learned bias variables.\n",
    "\n",
    "In particular, the first variable `h1` holds the learned weight matrix that converts an input vector of dimension `n_input + 2*n_input*n_context`  to a vector of dimension `n_hidden_1`. Similarly, the second variable `h2` holds the weight matrix converting an input vector of dimension `n_hidden_1` to one of dimension `n_hidden_2`. The variables `h3`, `h5`, and `h6` are similar. Likewise, the biases, `b1`, `b2`..., hold the biases for the various layers.\n",
    "\n",
    "That said let us introduce the method `BiRNN()` that takes a batch of data `batch_x` and performs inference upon it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def BiRNN(batch_x, seq_length):\n",
    "    # Input shape: [batch_size, n_steps, n_input + 2*n_input*n_context]\n",
    "    batch_x_shape = tf.shape(batch_x)\n",
    "    # Permute n_steps and batch_size\n",
    "    batch_x = tf.transpose(batch_x, [1, 0, 2])\n",
    "    # Reshape to prepare input for first layer\n",
    "    batch_x = tf.reshape(batch_x, [-1, n_input + 2*n_input*n_context]) # (n_steps*batch_size, n_input + 2*n_input*n_context)\n",
    "    \n",
    "    #Hidden layer with clipped RELU activation and dropout\n",
    "    b1 = variable_on_cpu('b1', [n_hidden_1], tf.random_normal_initializer(stddev=b1_stddev))\n",
    "    h1 = variable_on_cpu('h1', [n_input + 2*n_input*n_context, n_hidden_1], tf.random_normal_initializer(stddev=h1_stddev))\n",
    "    layer_1 = tf.minimum(tf.nn.relu(tf.add(tf.matmul(batch_x, h1), b1)), relu_clip)\n",
    "    layer_1 = tf.nn.dropout(layer_1, (1.0 - dropout_rate_placeholder))\n",
    "    #Hidden layer with clipped RELU activation and dropout\n",
    "    b2 = variable_on_cpu('b2', [n_hidden_2], tf.random_normal_initializer(stddev=b2_stddev))\n",
    "    h2 = variable_on_cpu('h2', [n_hidden_1, n_hidden_2], tf.random_normal_initializer(stddev=h2_stddev))\n",
    "    layer_2 = tf.minimum(tf.nn.relu(tf.add(tf.matmul(layer_1, h2), b2)), relu_clip)\n",
    "    layer_2 = tf.nn.dropout(layer_2, (1.0 - dropout_rate_placeholder))\n",
    "    #Hidden layer with clipped RELU activation and dropout\n",
    "    b3 = variable_on_cpu('b3', [n_hidden_3], tf.random_normal_initializer(stddev=b3_stddev))\n",
    "    h3 = variable_on_cpu('h3', [n_hidden_2, n_hidden_3], tf.random_normal_initializer(stddev=h3_stddev))\n",
    "    layer_3 = tf.minimum(tf.nn.relu(tf.add(tf.matmul(layer_2, h3), b3)), relu_clip)\n",
    "    layer_3 = tf.nn.dropout(layer_3, (1.0 - dropout_rate_placeholder))\n",
    "    \n",
    "    # Define lstm cells with tensorflow\n",
    "    # Forward direction cell\n",
    "    lstm_fw_cell = tf.nn.rnn_cell.BasicLSTMCell(n_cell_dim, forget_bias=1.0)\n",
    "    # Backward direction cell\n",
    "    lstm_bw_cell = tf.nn.rnn_cell.BasicLSTMCell(n_cell_dim, forget_bias=1.0)\n",
    "    \n",
    "    # Reshape data because rnn cell expects shape [max_time, batch_size, input_size]\n",
    "    layer_3 = tf.reshape(layer_3, [-1, batch_x_shape[0], n_hidden_3])\n",
    "\n",
    "    # Get lstm cell output\n",
    "    outputs, output_states = tf.nn.bidirectional_dynamic_rnn(cell_fw=lstm_fw_cell,\n",
    "                                                             cell_bw=lstm_bw_cell,\n",
    "                                                             inputs=layer_3,\n",
    "                                                             dtype=tf.float32,\n",
    "                                                             time_major=True,\n",
    "                                                             sequence_length=seq_length)\n",
    "    \n",
    "    # Reshape outputs from two tensors each of shape [n_steps, batch_size, n_cell_dim]\n",
    "    # to a single tensor of shape [n_steps*batch_size, 2*n_cell_dim]\n",
    "    outputs = tf.concat(2, outputs)\n",
    "    outputs = tf.reshape(outputs, [-1, 2*n_cell_dim])\n",
    "    \n",
    "    #Hidden layer with clipped RELU activation and dropout\n",
    "    b5 = variable_on_cpu('b5', [n_hidden_5], tf.random_normal_initializer(stddev=b5_stddev))\n",
    "    h5 = variable_on_cpu('h5', [(2 * n_cell_dim), n_hidden_5], tf.random_normal_initializer(stddev=h5_stddev))\n",
    "    layer_5 = tf.minimum(tf.nn.relu(tf.add(tf.matmul(outputs, h5), b5)), relu_clip)\n",
    "    layer_5 = tf.nn.dropout(layer_5, (1.0 - dropout_rate_placeholder))\n",
    "    #Hidden layer of logits\n",
    "    b6 = variable_on_cpu('b6', [n_hidden_6], tf.random_normal_initializer(stddev=b6_stddev))\n",
    "    h6 = variable_on_cpu('h6', [n_hidden_5, n_hidden_6], tf.random_normal_initializer(stddev=h6_stddev))\n",
    "    layer_6 = tf.add(tf.matmul(layer_5, h6), b6)\n",
    "    \n",
    "    # Reshape layer_6 from a tensor of shape [n_steps*batch_size, n_hidden_6]\n",
    "    # to a tensor of shape [n_steps, batch_size, n_hidden_6]\n",
    "    layer_6 = tf.reshape(layer_6, [-1, batch_x_shape[0], n_hidden_6])\n",
    "    \n",
    "    # Return layer_6\n",
    "    # Output shape: [n_steps, batch_size, n_input + 2*n_input*n_context]\n",
    "    return layer_6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first few lines of the function `BiRNN`\n",
    "```python\n",
    "def BiRNN(batch_x, seq_length):\n",
    "    # Input shape: [batch_size, n_steps, n_input + 2*n_input*n_context]\n",
    "    batch_x_shape = tf.shape(batch_x)\n",
    "    # Permute n_steps and batch_size\n",
    "    batch_x = tf.transpose(batch_x, [1, 0, 2])\n",
    "    # Reshape to prepare input for first layer\n",
    "    batch_x = tf.reshape(batch_x, [-1, n_input + 2*n_input*n_context])\n",
    "    ...\n",
    "```\n",
    "reshape `batch_x` which has shape `[batch_size, n_steps, n_input + 2*n_input*n_context]` initially, to a tensor with shape `[n_steps*batch_size, n_input + 2*n_input*n_context]`. This is done to prepare the batch for input into the first layer which expects a tensor of rank `2`.\n",
    "\n",
    "The next few lines of  `BiRNN`\n",
    "```python\n",
    "    #Hidden layer with clipped RELU activation and dropout\n",
    "    b1 = variable_on_cpu('b1', [n_hidden_1], tf.random_normal_initializer())\n",
    "    h1 = variable_on_cpu('h1', [n_input + 2*n_input*n_context, n_hidden_1], tf.random_normal_initializer())\n",
    "    layer_1 = tf.minimum(tf.nn.relu(tf.add(tf.matmul(batch_x, h1), b1)), relu_clip)\n",
    "    layer_1 = tf.nn.dropout(layer_1, (1.0 - dropout_rate_placeholder))\n",
    "    ...\n",
    "```\n",
    "pass `batch_x` through the first layer of the non-recurrent neural network, then applies dropout to the result.\n",
    "\n",
    "The next few lines do the same thing, but for the second and third layers\n",
    "```python\n",
    "    #Hidden layer with clipped RELU activation and dropout\n",
    "    b2 = variable_on_cpu('b2', [n_hidden_2], tf.random_normal_initializer())\n",
    "    h2 = variable_on_cpu('h2', [n_hidden_1, n_hidden_2], tf.random_normal_initializer())\n",
    "    layer_2 = tf.minimum(tf.nn.relu(tf.add(tf.matmul(layer_1, h2), b2)), relu_clip)   \n",
    "    layer_2 = tf.nn.dropout(layer_2, (1.0 - dropout_rate_placeholder))\n",
    "    #Hidden layer with clipped RELU activation and dropout\n",
    "    b3 = variable_on_cpu('b3', [n_hidden_3], tf.random_normal_initializer())\n",
    "    h3 = variable_on_cpu('h3', [n_hidden_2, n_hidden_3], tf.random_normal_initializer())\n",
    "    layer_3 = tf.minimum(tf.nn.relu(tf.add(tf.matmul(layer_2, h3), b3)), relu_clip)\n",
    "    layer_3 = tf.nn.dropout(layer_3, (1.0 - dropout_rate_placeholder))\n",
    "```\n",
    "\n",
    "Next we create the forward and backward LSTM units\n",
    "```python\n",
    "    # Define lstm cells with tensorflow\n",
    "    # Forward direction cell\n",
    "    lstm_fw_cell = tf.nn.rnn_cell.BasicLSTMCell(n_cell_dim, forget_bias=1.0)\n",
    "    # Backward direction cell\n",
    "    lstm_bw_cell = tf.nn.rnn_cell.BasicLSTMCell(n_cell_dim, forget_bias=1.0)\n",
    "```\n",
    "both of which have inputs of length `n_cell_dim` and bias `1.0` for the forget gate of the LSTM.\n",
    "\n",
    "The next line of the funtion `BiRNN` does a bit more data preparation.\n",
    "```python\n",
    "    # Reshape data because rnn cell expects shape [max_time, batch_size, input_size]\n",
    "    layer_3 = tf.reshape(layer_3, [-1, batch_x_shape[0], n_hidden_3])\n",
    "```\n",
    "It reshapes `layer_3` in to `[n_steps, batch_size, 2*n_cell_dim]` as the LSTM BRNN expects its input to be of shape `[max_time, batch_size, input_size]`.\n",
    "\n",
    "The next line of `BiRNN`\n",
    "```python\n",
    "    # Get lstm cell output\n",
    "    outputs, output_states = tf.nn.bidirectional_dynamic_rnn(cell_fw=lstm_fw_cell,\n",
    "                                                             cell_bw=lstm_bw_cell,\n",
    "                                                             inputs=layer_3,\n",
    "                                                             dtype=tf.float32,\n",
    "                                                             time_major=True,\n",
    "                                                             sequence_length=seq_length)\n",
    "```\n",
    "feeds `layer_3` to the LSTM BRNN cell and obtains the LSTM BRNN output.\n",
    "\n",
    "The next lines convert `outputs` from two rank two tensors into a single rank two tensor in preparation for passing it to the next neural network layer  \n",
    "```python\n",
    "    # Reshape outputs from two tensors each of shape [n_steps, batch_size, n_cell_dim]\n",
    "    # to a single tensor of shape [n_steps*batch_size, 2*n_cell_dim]\n",
    "    outputs = tf.concat(2, outputs)\n",
    "    outputs = tf.reshape(outputs, [-1, 2*n_cell_dim])\n",
    "```\n",
    "\n",
    "The next couple of lines feed `outputs` to the fifth hidden layer\n",
    "```python\n",
    "    #Hidden layer with clipped RELU activation and dropout\n",
    "    b5 = variable_on_cpu('b5', [n_hidden_5], tf.random_normal_initializer())\n",
    "    h5 = variable_on_cpu('h5', [(2 * n_cell_dim), n_hidden_5], tf.random_normal_initializer())\n",
    "    layer_5 = tf.minimum(tf.nn.relu(tf.add(tf.matmul(outputs, h5), b5)), relu_clip)\n",
    "    layer_5 = tf.nn.dropout(layer_5, (1.0 - dropout_rate_placeholder))\n",
    "```\n",
    "\n",
    "The next line of `BiRNN`\n",
    "```python\n",
    "    #Hidden layer of logits\n",
    "    b6 = variable_on_cpu('b6', [n_hidden_6], tf.random_normal_initializer())\n",
    "    h6 = variable_on_cpu('h6', [n_hidden_5, n_hidden_6], tf.random_normal_initializer())\n",
    "    layer_6 = tf.add(tf.matmul(layer_5, h6), b6)\n",
    "```\n",
    "Applies the weight matrix `h6` and bias `b6` to the output of `layer_5` creating `n_classes` dimensional vectors, the logits.\n",
    "\n",
    "The next lines of `BiRNN`\n",
    "```python\n",
    "    # Reshape layer_6 from a tensor of shape [n_steps*batch_size, n_hidden_6]\n",
    "    # to a tensor of shape [n_steps, batch_size, n_hidden_6]\n",
    "    layer_6 = tf.reshape(layer_6, [-1, batch_x_shape[0], n_hidden_6])\n",
    "```\n",
    "reshapes `layer_6` to the slightly more useful shape `[n_steps, batch_size, n_hidden_6]`. Note, that this differs from the input in that it is time-major.\n",
    "\n",
    "The final line of `BiRNN` returns `layer_6`\n",
    "```python\n",
    "    # Return layer_6\n",
    "    # Output shape: [n_steps, batch_size, n_input + 2*n_input*n_context]\n",
    "    return layer_6\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Accuracy and Loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In accord with [Deep Speech: Scaling up end-to-end speech recognition](http://arxiv.org/abs/1412.5567), the loss function used by our network should be the CTC loss function[[2]](http://www.cs.toronto.edu/~graves/preprint.pdf). Conveniently, this loss function is implemented in TensorFlow. Thus, we can simply make use of this implementation to define our loss.\n",
    "\n",
    "To do so we introduce a utility function `calculate_accuracy_and_loss()` that beam search decodes a mini-batch and calculates the loss and accuracy. Next to total and average loss it returns the accuracy, the decoded result and the batch's original Y."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def calculate_accuracy_and_loss(batch_x, batch_seq_len, batch_y):\n",
    "    # Calculate the logits of the batch using BiRNN\n",
    "    logits = BiRNN(batch_x, tf.to_int64(batch_seq_len))\n",
    "    \n",
    "    # Compute the CTC loss\n",
    "    if use_warpctc:\n",
    "        total_loss = tf.contrib.warpctc.warp_ctc_loss(logits, batch_y, batch_seq_len)\n",
    "    else:\n",
    "        total_loss = ctc_ops.ctc_loss(logits, batch_y, batch_seq_len)\n",
    "    \n",
    "    # Calculate the average loss across the batch\n",
    "    avg_loss = tf.reduce_mean(total_loss)\n",
    "    \n",
    "    # Beam search decode the batch\n",
    "    decoded, _ = ctc_ops.ctc_beam_search_decoder(logits, batch_seq_len, merge_repeated=False)\n",
    "    \n",
    "    # Compute the edit (Levenshtein) distance \n",
    "    distance = tf.edit_distance(tf.cast(decoded[0], tf.int32), batch_y)\n",
    "    \n",
    "    # Compute the accuracy \n",
    "    accuracy = tf.reduce_mean(distance)\n",
    "\n",
    "    # Return results to the caller\n",
    "    return total_loss, avg_loss, accuracy, decoded, batch_y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first line of `calculate_accuracy_and_loss()`\n",
    "```python\n",
    "    # Calculate the logits from the BiRNN\n",
    "    logits = BiRNN(batch_x, batch_seq_len)\n",
    "```\n",
    "calls `BiRNN()` with a batch of data and does inference on the batch.\n",
    "\n",
    "The next few lines\n",
    "```python\n",
    "    # Compute the CTC loss\n",
    "    total_loss = ctc_ops.ctc_loss(logits, batch_y, batch_seq_len)\n",
    "    \n",
    "    # Calculate the average loss across the batch\n",
    "    avg_loss = tf.reduce_mean(total_loss)\n",
    "```\n",
    "calculate the average loss using tensor flow's `ctc_loss` operator. \n",
    "\n",
    "The next lines first beam decode the batch and then compute the accuracy on base of the Levenshtein distance between the decoded batch and the batch's original Y.\n",
    "```python\n",
    "    # Beam search decode the batch\n",
    "    decoded, _ = ctc_ops.ctc_beam_search_decoder(logits, batch_seq_len, merge_repeated=False)\n",
    "    \n",
    "    # Compute the edit (Levenshtein) distance \n",
    "    distance = tf.edit_distance(tf.cast(decoded[0], tf.int32), batch_y)\n",
    "    \n",
    "    # Compute the accuracy \n",
    "    accuracy = tf.reduce_mean(distance)\n",
    "```\n",
    "\n",
    "Finally, the `total_loss`, `avg_loss`, `accuracy`, the `decoded` batch and the original `batch_y` are returned to the caller\n",
    "```python\n",
    "    # Return results to the caller\n",
    "    return total_loss, avg_loss, accuracy, decoded, batch_y\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parallel Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we will implement optimization of the DeepSpeech model across GPU's on a single host. This parallel optimization can take on various forms. For example one can use asynchronous updates of the model, synchronous updates of the model, or some combination of the two."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Asynchronous Parallel Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In asynchronous parallel optimization, for example, one places the model initially in CPU memory. Then each of the $G$ GPU's obtains a mini-batch of data along with the current model parameters. Using this mini-batch each GPU then computes the gradients for all model parameters and sends these gradients back to the CPU when the GPU is done with its mini-batch. The CPU then asynchronously updates the model parameters whenever it recieves a set of gradients from a GPU."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Asynchronous parallel optimization has several advantages and several disadvantages. One large advantage is throughput. No GPU will every be waiting idle. When a GPU is done processing a mini-batch, it can immediately obtain the next mini-batch to process. It never has to wait on other GPU's to finish their mini-batch. However, this means that the model updates will also be asynchronous which can have problems.\n",
    "\n",
    "For example, one may have model parameters $W$ on the CPU and send mini-batch $n$ to GPU 1 and send mini-batch $n+1$ to GPU 2. As processing is asynchronous, GPU 2 may finish before GPU 1 and thus update the CPU's model parameters $W$ with its gradients $\\Delta W_{n+1}(W)$, where the subscript $n+1$ identifies the mini-batch and the argument $W$ the location at which the gradient was evaluated. This results in the new model parameters\n",
    "\n",
    "$$W + \\Delta W_{n+1}(W).$$\n",
    "\n",
    "Next GPU 1 could finish with its mini-batch and update the parameters to\n",
    "\n",
    "$$W + \\Delta W_{n+1}(W) + \\Delta W_{n}(W).$$\n",
    "\n",
    "The problem with this is that $\\Delta W_{n}(W)$ is evaluated at $W$ and not $W + \\Delta W_{n+1}(W)$. Hence, the direction of the gradient $\\Delta W_{n}(W)$ is slightly incorrect as it is evaluated at the wrong location. This can be counteracted through synchronous updates of model, but this is also problematic."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Synchronous Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Synchronous optimization solves the problem we saw above. In synchronous optimization, one places the model initially in CPU memory. Then one of the $G$ GPU's is given a mini-batch of data along with the current model parameters. Using the mini-batch the GPU computes the gradients for all model parameters and sends the gradients back to the CPU. The CPU then updates the model parameters and starts the process of sending out the next mini-batch.\n",
    "\n",
    "As on can readily see, synchronous optimization does not have the problem we found in the last section, that of incorrect gradients. However, synchronous optimization can only make use of a single GPU at a time. So, when we have a multi-GPU setup, $G > 1$, all but one of the GPU's will remain idle, which is unacceptable. However, there is a third alternative which is combines the advantages of asynchronous and synchronous optimization."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hybrid Parallel Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hybrid parallel optimization combines most of the benifits of asynchronous and synchronous optimization. It allows for multiple GPU's to be used, but does not suffer from the incorrect gradient problem exhibited by asynchronous optimization.\n",
    "\n",
    "In hybrid parallel optimization one places the model initially in CPU memory. Then, as in asynchronous optimization, each of the $G$ GPU'S obtains a mini-batch of data along with the current model parameters. Using the mini-batch each of the GPU's then computes the gradients for all model parameters and sends these gradients back to the CPU. Now, in contrast to asynchronous optimization, the CPU waits until each GPU is finished with its mini-batch then takes the mean of all the gradients from the $G$ GPU's and updates the model with this mean gradient."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"images/Parallelism.png\" alt=\"LSTM\" width=\"600\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hybrid parallel optimization has several advantages and few disadvantages. As in asynchronous parallel optimization, hybrid parallel optimization allows for one to use multiple GPU's in parallel. Furthermore, unlike asynchronous parallel optimization, the incorrect gradient problem is not present here. In fact, hybrid parallel optimization performs as if one is working with a single mini-batch which is $G$ times the size of a mini-batch handled by a single GPU. Hoewever, hybrid parallel optimization is not perfect. If one GPU is slower than all the others in completing its mini-batch, all other GPU's will have to sit idle until this straggler finishes with its mini-batch. This hurts throughput. But, if all GPU'S are of the same make and model, this problem should be minimized.\n",
    "\n",
    "So, relatively speaking, hybrid parallel optimization seems the have more advantages and fewer disadvantages as compared to both asynchronous and synchronous optimization. So, we will, for our work, use this hybrid model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adam Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In constrast to [Deep Speech: Scaling up end-to-end speech recognition](http://arxiv.org/abs/1412.5567), in which  [Nesterov’s Accelerated Gradient Descent](www.cs.toronto.edu/~fritz/absps/momentum.pdf) was used, we will use the Adam method for optimization[[3](http://arxiv.org/abs/1412.6980)], because, generally, it requires less fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def create_optimizer():\n",
    "    optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate,\n",
    "                                       beta1=beta1,\n",
    "                                       beta2=beta2,\n",
    "                                       epsilon=epsilon)\n",
    "    return optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Towers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to properly make use of multiple GPU's, one must introduce new abstractions, not present when using a single GPU, that facilitate the multi-GPU use case.\n",
    "\n",
    "In particular, one must introduce a means to isolate the inference and gradient calculations on the various GPU's. The abstraction we intoduce for this purpose is called a 'tower'. A tower is specified by two properties:\n",
    "* **Scope** - A scope, as provided by `tf.name_scope()`, is a means to isolate the operations within a tower. For example, all operations within \"tower 0\" could have their name prefixed with `tower_0/`.\n",
    "*  **Device** - A hardware device, as provided by `tf.device()`, on which all operations within the tower execute. For example, all operations of \"tower 0\" could execute on the first GPU `tf.device('/gpu:0')`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we are introducing one tower for each GPU, first we must determine how many GPU's are available"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Get a list of the available gpu's ['/gpu:0', '/gpu:1'...]\n",
    "available_devices = get_available_gpus()\n",
    "\n",
    "# If there are no GPU's use the CPU\n",
    "if 0 == len(available_devices):\n",
    "    available_devices = ['/cpu:0']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With this preliminary step out of the way, we can for each GPU intoduce a tower for which's batch we calculate \n",
    " \n",
    " * the CTC decodings ```decoded```,\n",
    " * the (total) loss against the outcome (Y) ```total_loss```, \n",
    " * the loss averaged over the whole batch ```avg_loss```,\n",
    " * the optimization gradient (computed on base of the averaged loss) and\n",
    " * the accuracy of the outcome averaged over the whole batch ```accuracy``` \n",
    " \n",
    "and retain the original ```labels``` (Y).\n",
    " \n",
    "```decoded```, ```labels```, the optimization gradient, ```total_loss``` and ```avg_loss``` are collected into the respectful arrays ```tower_decodings, tower_labels, tower_gradients, tower_total_losses, tower_avg_losses``` (dimension 0 being the tower).\n",
    "\n",
    "Finally this new method `get_tower_results()` will return those tower arrays plus (based on ```accuracy```) the averaged accuracy value accross all towers ```avg_accuracy```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_tower_results(batch_sets, optimizer):\n",
    "    # Tower decodings to return\n",
    "    tower_decodings = []\n",
    "    # Tower labels to return\n",
    "    tower_labels = []\n",
    "    # Tower gradients to return\n",
    "    tower_gradients = []\n",
    "    # Tower total batch losses to return\n",
    "    tower_total_losses = []\n",
    "    # Tower avg batch losses to return\n",
    "    tower_avg_losses = []\n",
    "    # Tower accuracies to return\n",
    "    tower_accuracies = []\n",
    "    \n",
    "    # Loop over available_devices\n",
    "    for i in xrange(len(available_devices)):\n",
    "        # Execute operations of tower i on device i\n",
    "        with tf.device(available_devices[i]):\n",
    "            # Create a scope for all operations of tower i\n",
    "            with tf.name_scope('tower_%d' % i) as scope:\n",
    "                # Fetch the next batch of data\n",
    "                batch_x, batch_x_seq_len, batch_y, batch_y_seq_len = \\\n",
    "                tf.cond(tf.less(queue_selector_placeholder, 1),\n",
    "                        lambda: batch_sets[0].next_batch(),\n",
    "                        lambda: tf.cond(tf.less(queue_selector_placeholder, 2),\n",
    "                                        lambda: batch_sets[1].next_batch(),\n",
    "                                        lambda: batch_sets[2].next_batch()))\n",
    "                \n",
    "                # Calculate the avg_loss and accuracy and retrieve the decoded \n",
    "                # batch along with the original batch's labels (Y) of this tower\n",
    "                batch_y = ctc_label_dense_to_sparse(batch_y, batch_y_seq_len, batch_size)\n",
    "                total_loss, avg_loss, accuracy, decoded, labels = calculate_accuracy_and_loss(batch_x,\n",
    "                                                                                              batch_x_seq_len,\n",
    "                                                                                              batch_y)\n",
    "                                \n",
    "                # Allow for variables to be re-used by the next tower\n",
    "                tf.get_variable_scope().reuse_variables()\n",
    "                \n",
    "                # Retain tower's decoded batch\n",
    "                tower_decodings.append(decoded)\n",
    "                \n",
    "                # Retain tower's labels (Y)\n",
    "                tower_labels.append(labels)\n",
    "                \n",
    "                # Compute gradients for model parameters using tower's mini-batch\n",
    "                gradients = optimizer.compute_gradients(avg_loss)\n",
    "\n",
    "                # Retain tower's gradients\n",
    "                tower_gradients.append(gradients)\n",
    "                    \n",
    "                # Retain tower's total losses\n",
    "                tower_total_losses.append(total_loss)\n",
    "                \n",
    "                # Retain tower's avg losses\n",
    "                tower_avg_losses.append(avg_loss)\n",
    "                \n",
    "                # Retain tower's accuracies\n",
    "                tower_accuracies.append(accuracy)\n",
    "                \n",
    "    # Average accuracies over the 'tower' dimension\n",
    "    avg_accuracy = tf.reduce_mean(tower_accuracies, 0)\n",
    "\n",
    "    # Return results to caller\n",
    "    return tower_decodings, tower_labels, tower_gradients, tower_total_losses, tower_avg_losses, avg_accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we want to average the gradients obtained from the GPU's."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We compute the average the gradients obtained from the GPU's for each variable in the function `average_gradients()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def average_gradients(tower_gradients):\n",
    "    # List of average gradients to return to the caller\n",
    "    average_grads = []\n",
    "    \n",
    "    # Loop over gradient/variable pairs from all towers\n",
    "    for grad_and_vars in zip(*tower_gradients):\n",
    "        # Introduce grads to store the gradients for the current variable\n",
    "        grads = []\n",
    "        \n",
    "        # Loop over the gradients for the current variable\n",
    "        for g, _ in grad_and_vars:\n",
    "            # Add 0 dimension to the gradients to represent the tower.\n",
    "            expanded_g = tf.expand_dims(g, 0)\n",
    "            # Append on a 'tower' dimension which we will average over below.\n",
    "            grads.append(expanded_g)\n",
    "            \n",
    "        # Average over the 'tower' dimension\n",
    "        grad = tf.concat(0, grads)\n",
    "        grad = tf.reduce_mean(grad, 0)\n",
    "        \n",
    "        # Create a gradient/variable tuple for the current variable with its average gradient\n",
    "        grad_and_var = (grad, grad_and_vars[0][1])\n",
    "        \n",
    "        # Add the current tuple to average_grads\n",
    "        average_grads.append(grad_and_var)\n",
    "    \n",
    "    #Return result to caller\n",
    "    return average_grads"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note also that this code acts as a syncronization point as it requires all GPU's to be finished with their mini-batch before it can run to completion.\n",
    "\n",
    "Now next we introduce a function to apply the averaged gradients to update the model's paramaters on the CPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def apply_gradients(optimizer, average_grads):\n",
    "    apply_gradient_op = optimizer.apply_gradients(average_grads)\n",
    "    return apply_gradient_op"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Logging"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We introduce a function for logging a tensor variable's current state.\n",
    "It logs scalar values for the mean, standard deviation, minimum and maximum.\n",
    "Furthermore it logs a histogram of its state and (if given) of an optimization gradient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def log_variable(variable, gradient=None):\n",
    "    name = variable.name\n",
    "    mean = tf.reduce_mean(variable)\n",
    "    tf.scalar_summary(name + '/mean', mean)\n",
    "    tf.scalar_summary(name + '/sttdev', tf.sqrt(tf.reduce_mean(tf.square(variable - mean))))\n",
    "    tf.scalar_summary(name + '/max', tf.reduce_max(variable))\n",
    "    tf.scalar_summary(name + '/min', tf.reduce_min(variable))\n",
    "    tf.histogram_summary(name, variable)\n",
    "    if gradient is not None:\n",
    "        if isinstance(gradient, tf.IndexedSlices):\n",
    "            grad_values = gradient.values\n",
    "        else:\n",
    "            grad_values = gradient\n",
    "        if grad_values is not None:\n",
    "            tf.histogram_summary(name + \"/gradients\", grad_values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also introduce a helper function for logging collections of gradient/variable tuples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def log_grads_and_vars(grads_and_vars):\n",
    "    for gradient, variable in grads_and_vars:\n",
    "        log_variable(variable, gradient=gradient)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we define the top directory for all logs and our current log sub-directory of it.\n",
    "We also add some log helpers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "logs_dir = os.environ.get('ds_logs_dir', 'logs')\n",
    "log_dir = '%s/%s' % (logs_dir, time.strftime(\"%Y%m%d-%H%M%S\"))\n",
    "\n",
    "def get_git_revision_hash():\n",
    "    return subprocess.check_output(['git', 'rev-parse', 'HEAD']).strip()\n",
    "\n",
    "def get_git_branch():\n",
    "    return subprocess.check_output(['git', 'rev-parse', '--abbrev-ref', 'HEAD']).strip()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Helpers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's introduce a routine that prints a WER report under a given caption. It'll print the given ```mean``` WER plus summaries of the top ten lowest loss items of the given array of WER result tuples (only items with WER!=0 and ordered by their WER)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def print_wer_report(caption, mean, items=[]):\n",
    "    print\n",
    "    print \"#\" * 80\n",
    "    print \"%s WER: %f\" % (caption, mean)\n",
    "    if len(items) > 0:\n",
    "        # Filter out all items with WER=0\n",
    "        items = [a for a in items if a[2] > 0]\n",
    "        # Order the remaining items by their loss (lowest loss on top)\n",
    "        items.sort(key=lambda a: a[3])\n",
    "        # Take only the first 10 items\n",
    "        items = items[:10]\n",
    "        # Order this top ten items by their WER (lowest WER on top)\n",
    "        items.sort(key=lambda a: a[2])\n",
    "        for a in items:\n",
    "            print \"-\" * 80\n",
    "            print \" - WER:    %f\" % a[2]\n",
    "            print \" - loss:   %f\" % a[3]\n",
    "            print \" - source: \\\"%s\\\"\" % a[0]\n",
    "            print \" - result: \\\"%s\\\"\" % a[1] \n",
    "    print \"#\" * 80\n",
    "    print"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another routine will help collecting partial results for the WER reports. The ```results_tuple``` is composed of an array of the original labels, an array of the corrsponding decodings and an array of the corresponding losses. ```returns``` is built up in a similar way, containing just the unprocessed results of one ```session.run``` call (effectively of one batch). Before splicing them into their corresponding ```results_tuple``` lists, labels and decodings are converted to text. In the case of decodings, for now we just pick the first available path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def collect_results(results_tuple, returns):\n",
    "    # Each of the arrays within results_tuple will get extended by a batch of each available device\n",
    "    for i in xrange(len(available_devices)):\n",
    "        # Collect the labels\n",
    "        results_tuple[0].extend(sparse_tensor_value_to_texts(returns[0][i]))\n",
    "        \n",
    "        # Collect the decodings - at the moment we default to the first one \n",
    "        results_tuple[1].extend(sparse_tensor_value_to_texts(returns[1][i][0]))\n",
    "        \n",
    "        # Collect the losses\n",
    "        results_tuple[2].extend(returns[2][i])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's add a helper that bundles WER calculation and WER report printing. It will return the mean WER."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def calculate_and_print_wer_report(caption, results_tuple):       \n",
    "    # Calculating the individual WERs and the overall mean WER\n",
    "    rates, wer = wers(results_tuple[0], results_tuple[1])\n",
    "    \n",
    "    # Printing the WER report\n",
    "    print_wer_report(caption, wer, zip(results_tuple[0], results_tuple[1], rates, results_tuple[2]))\n",
    "    \n",
    "    # Return the mean WER\n",
    "    return wer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we can assemble a helper routine that facilitates a full inference run plus report, given that the inference graph has already been constructed. It will return the mean WER."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def run_inference(session, caption, total_batches, feed_dict, results_params):\n",
    "    # Initialize the result collection tuple\n",
    "    results = ([],[],[])\n",
    "    \n",
    "    # Loop over the batches\n",
    "    for batch in range(int(ceil(float(total_batches) / len(available_devices)))):\n",
    "        # Run and collect the results\n",
    "        collect_results(results, session.run(results_params, feed_dict))\n",
    "        \n",
    "    # Calculate and print the WER report. Returns the mean WER.\n",
    "    return calculate_and_print_wer_report(caption, results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, as we have prepared all the apropos operators and methods, we can create the method which trains the network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def train(session, data_sets):\n",
    "    # Calculate the total number of batches\n",
    "    total_batches = data_sets.train.total_batches\n",
    "    batches_per_device = float(total_batches) / len(available_devices)\n",
    "        \n",
    "    # Create optimizer\n",
    "    optimizer = create_optimizer()\n",
    "\n",
    "    # Get gradients for each tower (Runs across all GPU's)\n",
    "    tower_decodings, \\\n",
    "    tower_labels, \\\n",
    "    tower_gradients, \\\n",
    "    tower_total_losses, \\\n",
    "    tower_avg_losses, \\\n",
    "    avg_accuracy \\\n",
    "    = get_tower_results([data_sets.train, data_sets.dev, data_sets.test], optimizer)\n",
    "    \n",
    "    # Average tower gradients\n",
    "    avg_tower_gradients = average_gradients(tower_gradients)\n",
    "\n",
    "    # Add logging of averaged gradients\n",
    "    log_grads_and_vars(avg_tower_gradients)\n",
    "\n",
    "    # Apply gradients to modify the model\n",
    "    apply_gradient_op = apply_gradients(optimizer, avg_tower_gradients)\n",
    "\n",
    "    # Create a saver to checkpoint the model\n",
    "    saver = tf.train.Saver(tf.all_variables())\n",
    "\n",
    "    # Init all variables in session\n",
    "    session.run(tf.initialize_all_variables())\n",
    "    \n",
    "    # Start queue runner threads\n",
    "    tf.train.start_queue_runners(sess=session)\n",
    "    \n",
    "    # Start importer's queue threads\n",
    "    data_sets.start_queue_threads(session)\n",
    "    \n",
    "    # Result parameters\n",
    "    results_params = (tower_labels, tower_decodings, tower_total_losses)\n",
    "    \n",
    "    # Prepare tensor board logging\n",
    "    merged = tf.merge_all_summaries()\n",
    "    writer = tf.train.SummaryWriter(log_dir, session.graph)\n",
    "    \n",
    "    # Train parameters\n",
    "    train_params = [apply_gradient_op, merged]\n",
    "\n",
    "    # Init recent word error rate levels\n",
    "    train_wer = 0.0\n",
    "    dev_wer = 0.0\n",
    "    \n",
    "    # Possibly restore checkpoint\n",
    "    start_epoch = 0\n",
    "    if restore_checkpoint:\n",
    "        checkpoint = tf.train.get_checkpoint_state(checkpoint_dir)\n",
    "        if checkpoint and checkpoint.model_checkpoint_path:\n",
    "            saver.restore(session, checkpoint.model_checkpoint_path)\n",
    "            start_epoch = int(checkpoint.model_checkpoint_path.split('-')[-1])\n",
    "            print 'Resuming training from epoch %d' % (start_epoch + 1)\n",
    "    \n",
    "    # Loop over the data set for training_epochs epochs\n",
    "    for epoch in range(start_epoch, training_iters):\n",
    "        params = train_params[:]\n",
    "        \n",
    "        # Determine if we want to display/validate/checkpoint on this iteration\n",
    "        is_display_step = display_step > 0 and (epoch + 1) % display_step == 0\n",
    "        is_validation_step = validation_step > 0 and (epoch + 1) % validation_step == 0\n",
    "        is_checkpoint_step = (epoch > 0 and (epoch + 1) % checkpoint_step == 0) or epoch == training_iters - 1\n",
    "        \n",
    "        # Requirements to display a WER report\n",
    "        if is_display_step:\n",
    "            # Reset accuracy\n",
    "            total_accuracy = 0.0\n",
    "            # Create training results tuple\n",
    "            train_results = ([],[],[])\n",
    "            # Extend the session.run parameters\n",
    "            params.append(results_params)\n",
    "            params.append(avg_accuracy)\n",
    "\n",
    "        # Loop over the batches\n",
    "        for batch in range(int(ceil(batches_per_device))):\n",
    "            extra_params = { }\n",
    "            if do_fulltrace:\n",
    "                loss_run_metadata            = tf.RunMetadata()\n",
    "                extra_params['options']      = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)\n",
    "                extra_params['run_metadata'] = loss_run_metadata\n",
    "            \n",
    "            # Compute the batch\n",
    "            result = session.run(params, feed_dict_train, **extra_params)\n",
    "\n",
    "            # Log all variable states in current step\n",
    "            step = epoch * total_batches + batch * len(available_devices)\n",
    "            writer.add_summary(result[1], step)\n",
    "            if do_fulltrace:\n",
    "                writer.add_run_metadata(loss_run_metadata, 'loss_epoch%d_batch%d'   % (epoch, batch))\n",
    "            writer.flush()\n",
    "            \n",
    "            # Collect individual sample results for WER report\n",
    "            if is_display_step:\n",
    "                collect_results(train_results, result[2])\n",
    "                # Add batch to total_accuracy\n",
    "                total_accuracy += result[3]\n",
    "        \n",
    "        # Print WER report\n",
    "        if is_display_step:\n",
    "            print \"Epoch:\", '%04d' % (epoch), \"avg_cer=\", \"{:.9f}\".format((total_accuracy / ceil(batches_per_device)))\n",
    "            train_wer = calculate_and_print_wer_report(\"Training\", train_results)\n",
    "            print\n",
    "        else:\n",
    "            print \"Epoch: %04d\" % (epoch)\n",
    "            \n",
    "        # Validation step\n",
    "        if is_validation_step:\n",
    "            dev_wer = run_inference(session, \"Validation\", data_sets.dev.total_batches, feed_dict_validate, results_params)\n",
    "\n",
    "        # Checkpoint the model\n",
    "        if is_checkpoint_step:\n",
    "            checkpoint_path = os.path.join(checkpoint_dir, 'model.ckpt')\n",
    "            print \"Checkpointing in directory\", \"%s\" % checkpoint_dir\n",
    "            saver.save(session, checkpoint_path, global_step=epoch)\n",
    "            print\n",
    "        \n",
    "    # Indicate optimization has concluded\n",
    "    print \"Optimization Finished!\"\n",
    "    return train_wer, dev_wer, results_params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As everything is prepared, we are now able to do the training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Define CPU as device on which the muti-gpu training is orchestrated\n",
    "with tf.device('/cpu:0'):\n",
    "    # Obtain all the data, defaulting to TED LIUM\n",
    "    data_sets = ds_importer_module.read_data_sets(ds_dataset_path, batch_size, n_input, n_context)\n",
    "    \n",
    "    # Create session in which to execute\n",
    "    session = tf.Session(config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True))\n",
    "    \n",
    "    # Take start time for time measurement\n",
    "    time_started = datetime.datetime.utcnow()\n",
    "    \n",
    "    # Train the network\n",
    "    last_train_wer, last_dev_wer, results_params = train(session, data_sets)\n",
    "    \n",
    "    # Take final time for time measurement\n",
    "    time_finished = datetime.datetime.utcnow()\n",
    "    \n",
    "    # Calculate duration in seconds\n",
    "    duration = time_finished - time_started\n",
    "    duration = duration.days * 86400 + duration.seconds\n",
    "    \n",
    "    # Finally the model is tested against some unbiased data-set\n",
    "    test_wer = run_inference(session, \"Test\", data_sets.test.total_batches, feed_dict_test, results_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we restore the trained variables into a simpler graph that we can export for serving."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Don't export a model if no export directory has been set\n",
    "if export_dir:\n",
    "    with tf.device('/cpu:0'):\n",
    "        tf.reset_default_graph()\n",
    "        session = tf.Session(config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True))\n",
    "\n",
    "        # Run inference\n",
    "        # Replace the dropout placeholder with a constant\n",
    "        dropout_rate_placeholder = tf.constant(0.0)\n",
    "\n",
    "        # Input tensor will be of shape [batch_size, n_steps, n_input + 2*n_input*n_context]\n",
    "        input_tensor = tf.placeholder(tf.float32, [None, None, n_input + 2*n_input*n_context])\n",
    "\n",
    "        # Calculate input sequence length. This is done by tiling n_steps, batch_size times.\n",
    "        # If there are multiple sequences, it is assumed they are padded with zeros to be of\n",
    "        # the same length.\n",
    "        n_items  = tf.slice(tf.shape(input_tensor), [0], [1])\n",
    "        n_steps = tf.slice(tf.shape(input_tensor), [1], [1])\n",
    "        seq_length = tf.tile(n_steps, n_items)\n",
    "\n",
    "        # Calculate the logits of the batch using BiRNN\n",
    "        logits = BiRNN(input_tensor, tf.to_int64(seq_length))\n",
    "\n",
    "        # Beam search decode the batch\n",
    "        decoded, _ = ctc_ops.ctc_beam_search_decoder(logits, seq_length, merge_repeated=False)\n",
    "        decoded = tf.convert_to_tensor(\n",
    "            [tf.sparse_tensor_to_dense(sparse_tensor) for sparse_tensor in decoded])\n",
    "\n",
    "        # TODO: Transform the decoded output to a string\n",
    "\n",
    "        # Create a saver and exporter using variables from the above newly created graph\n",
    "        saver = tf.train.Saver(tf.all_variables())\n",
    "        model_exporter = exporter.Exporter(saver)\n",
    "        \n",
    "        # Restore variables from training checkpoint\n",
    "        # TODO: This restores the most recent checkpoint, but if we use validation to counterract\n",
    "        #       over-fitting, we may want to restore an earlier checkpoint.\n",
    "        checkpoint = tf.train.get_checkpoint_state(checkpoint_dir)\n",
    "        saver.restore(session, checkpoint.model_checkpoint_path)\n",
    "        print 'Restored checkpoint at training epoch %d' % (int(checkpoint.model_checkpoint_path.split('-')[-1]) + 1)\n",
    "\n",
    "        # Initialise the model exporter and export the model\n",
    "        model_exporter.init(session.graph.as_graph_def(),\n",
    "                            named_graph_signatures = {\n",
    "                                'inputs': exporter.generic_signature(\n",
    "                                    { 'input': input_tensor }),\n",
    "                                'outputs': exporter.generic_signature(\n",
    "                                    { 'outputs': decoded})})\n",
    "        if remove_export:\n",
    "            actual_export_dir = os.path.join(export_dir, '%08d' % export_version)\n",
    "            if os.path.isdir(actual_export_dir):\n",
    "                print 'Removing old export'\n",
    "                shutil.rmtree(actual_export_dir)\n",
    "        try:\n",
    "            model_exporter.export(export_dir, tf.constant(export_version), session)\n",
    "            print 'Model exported at %s' % (export_dir)\n",
    "        except RuntimeError:\n",
    "            print  sys.exc_info()[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Logging Hyper Parameters and Results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, as training and test are done, we persist the results alongside with the involved hyper parameters for further reporting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "with open('%s/%s' % (log_dir, 'hyper.json'), 'w') as dump_file:\n",
    "    json.dump({ \\\n",
    "        'context': { \\\n",
    "            'time_started': time_started.isoformat(), \\\n",
    "            'time_finished': time_finished.isoformat(), \\\n",
    "            'git_hash': get_git_revision_hash(), \\\n",
    "            'git_branch': get_git_branch() \\\n",
    "        }, \\\n",
    "        'parameters': { \\\n",
    "            'learning_rate': learning_rate, \\\n",
    "            'beta1': beta1, \\\n",
    "            'beta2': beta2, \\\n",
    "            'epsilon': epsilon, \\\n",
    "            'training_iters': training_iters, \\\n",
    "            'batch_size': batch_size, \\\n",
    "            'validation_step': validation_step, \\\n",
    "            'dropout_rate': dropout_rate, \\\n",
    "            'relu_clip': relu_clip, \\\n",
    "            'n_input': n_input, \\\n",
    "            'n_context': n_context, \\\n",
    "            'n_hidden_1': n_hidden_1, \\\n",
    "            'n_hidden_2': n_hidden_2, \\\n",
    "            'n_hidden_3': n_hidden_3, \\\n",
    "            'n_hidden_5': n_hidden_5, \\\n",
    "            'n_hidden_6': n_hidden_6, \\\n",
    "            'n_cell_dim': n_cell_dim, \\\n",
    "            'n_character': n_character, \\\n",
    "            'total_batches_train': data_sets.train.total_batches, \\\n",
    "            'total_batches_validation': data_sets.dev.total_batches, \\\n",
    "            'total_batches_test': data_sets.test.total_batches, \\\n",
    "            'data_set': { \\\n",
    "                'name': ds_importer \\\n",
    "            }, \\\n",
    "        }, \\\n",
    "        'results': { \\\n",
    "            'duration': duration, \\\n",
    "            'last_train_wer': last_train_wer, \\\n",
    "            'last_validation_wer': last_dev_wer, \\\n",
    "            'test_wer': test_wer \\\n",
    "        } \\\n",
    "    }, dump_file, sort_keys=True, indent = 4)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also re-populate a central JS file, that contains all the dumps at once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "merge_logs(logs_dir)\n",
    "maybe_publish()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}