05. Memory Augmented Networks/5.7 Building MANN in Tensorflow .ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building MANN in Tensorflow "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we will see how to implement MANN in tensorflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let us import all the required libraries,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "import tensorflow as tf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class MANNCell():\n",
    "    def __init__(self, rnn_size, memory_size, memory_vector_dim, head_num, gamma=0.95,\n",
    "                 reuse=False):\n",
    "        \n",
    "        #initialize all the variables\n",
    "        self.rnn_size = rnn_size\n",
    "        self.memory_size = memory_size\n",
    "        self.memory_vector_dim = memory_vector_dim\n",
    "        self.head_num = head_num                                   \n",
    "        self.reuse = reuse\n",
    "        self.step = 0\n",
    "        self.gamma = gamma\n",
    "      \n",
    "        \n",
    "        #initialize controller as the basic rnn cell\n",
    "        self.controller = tf.nn.rnn_cell.BasicRNNCell(self.rnn_size)\n",
    "\n",
    "        \n",
    "\n",
    "    def __call__(self, x, prev_state):\n",
    "        prev_read_vector_list = prev_state['read_vector_list']     \n",
    "        controller_input = tf.concat([x] + prev_read_vector_list, axis=1)\n",
    "        \n",
    "        #next we pass the controller, which is the RNN cell, the controller_input and prev_controller_state\n",
    "        with tf.variable_scope('controller', reuse=self.reuse):\n",
    "            controller_output, controller_state = self.controller(controller_input, prev_controller_state)\n",
    "\n",
    "            \n",
    "        num_parameters_per_head = self.memory_vector_dim + 1\n",
    "        total_parameter_num = num_parameters_per_head * self.head_num\n",
    "        \n",
    "        #Initiliaze weight matrix and bias and compute the parameters\n",
    "        \n",
    "        with tf.variable_scope(\"o2p\", reuse=(self.step > 0) or self.reuse):\n",
    "            o2p_w = tf.get_variable('o2p_w', [controller_output.get_shape()[1], total_parameter_num],\n",
    "                                    initializer=tf.random_uniform_initializer(minval=-0.1, maxval=0.1))\n",
    "            o2p_b = tf.get_variable('o2p_b', [total_parameter_num],\n",
    "                                    initializer=tf.random_uniform_initializer(minval=-0.1, maxval=0.1))\n",
    "            parameters = tf.nn.xw_plus_b(controller_output, o2p_w, o2p_b)\n",
    "        \n",
    "        head_parameter_list = tf.split(parameters, self.head_num, axis=1)\n",
    "        \n",
    "        \n",
    "        #previous read weight vector\n",
    "        prev_w_r_list = prev_state['w_r_list']   \n",
    "        \n",
    "        #previous memory\n",
    "        prev_M = prev_state['M']\n",
    "        \n",
    "        #previous usage weight vector\n",
    "        prev_w_u = prev_state['w_u']\n",
    "        \n",
    "        #previous index and least used weight vector\n",
    "        prev_indices, prev_w_lu = self.least_used(prev_w_u)\n",
    "        \n",
    "        #read weight vector\n",
    "        w_r_list = []\n",
    "        \n",
    "        #write weight vector\n",
    "        w_w_list = []\n",
    "        \n",
    "        #key vector\n",
    "        k_list = []\n",
    "    \n",
    "        #now, we will initialize some of the important parameters that we use for addressing.     \n",
    "        for i, head_parameter in enumerate(head_parameter_list):\n",
    "            with tf.variable_scope('addressing_head_%d' % i):\n",
    "                \n",
    "                #key vector\n",
    "                k = tf.tanh(head_parameter[:, 0:self.memory_vector_dim], name='k')\n",
    "                \n",
    "                #sig_alpha\n",
    "                sig_alpha = tf.sigmoid(head_parameter[:, -1:], name='sig_alpha')\n",
    "                \n",
    "                #read weights\n",
    "                w_r = self.read_head_addressing(k, prev_M)\n",
    "                \n",
    "                #write weights\n",
    "                w_w = self.write_head_addressing(sig_alpha, prev_w_r_list[i], prev_w_lu)\n",
    "           \n",
    "            w_r_list.append(w_r)\n",
    "            w_w_list.append(w_w)\n",
    "            k_list.append(k)\n",
    "            \n",
    "\n",
    "        #usage weight vector \n",
    "        w_u = self.gamma * prev_w_u + tf.add_n(w_r_list) + tf.add_n(w_w_list)   \n",
    "\n",
    "        #update the memory\n",
    "        M_ = prev_M * tf.expand_dims(1. - tf.one_hot(prev_indices[:, -1], self.memory_size), dim=2)\n",
    "        \n",
    "        \n",
    "        #write operation\n",
    "        M = M_\n",
    "        with tf.variable_scope('writing'):\n",
    "            for i in range(self.head_num):\n",
    "                \n",
    "                w = tf.expand_dims(w_w_list[i], axis=2)\n",
    "                k = tf.expand_dims(k_list[i], axis=1)\n",
    "                M = M + tf.matmul(w, k)\n",
    "\n",
    "        #read opearion\n",
    "        read_vector_list = []\n",
    "        with tf.variable_scope('reading'):\n",
    "            for i in range(self.head_num):\n",
    "                read_vector = tf.reduce_sum(tf.expand_dims(w_r_list[i], dim=2) * M, axis=1)\n",
    "                read_vector_list.append(read_vector)\n",
    "\n",
    "        \n",
    "        #controller output\n",
    "        NTM_output = tf.concat([controller_output] + read_vector_list, axis=1)\n",
    "\n",
    "        state = {\n",
    "            'controller_state': controller_state,\n",
    "            'read_vector_list': read_vector_list,\n",
    "            'w_r_list': w_r_list,\n",
    "            'w_w_list': w_w_list,\n",
    "            'w_u': w_u,\n",
    "            'M': M,\n",
    "        }\n",
    "\n",
    "        self.step += 1\n",
    "        return NTM_output, state\n",
    "\n",
    "    #weight vector for read operation   \n",
    "    def read_head_addressing(self, k, prev_M):\n",
    "        \n",
    "        \"content based cosine similarity\"\n",
    "        \n",
    "        k = tf.expand_dims(k, axis=2)\n",
    "        inner_product = tf.matmul(prev_M, k)\n",
    "        k_norm = tf.sqrt(tf.reduce_sum(tf.square(k), axis=1, keep_dims=True))\n",
    "        M_norm = tf.sqrt(tf.reduce_sum(tf.square(prev_M), axis=2, keep_dims=True))\n",
    "        norm_product = M_norm * k_norm\n",
    "        K = tf.squeeze(inner_product / (norm_product + 1e-8))                  \n",
    "        K_exp = tf.exp(K)\n",
    "        w = K_exp / tf.reduce_sum(K_exp, axis=1, keep_dims=True)               \n",
    "        \n",
    "        return w\n",
    "    \n",
    "    #weight vector for write operation\n",
    "    def write_head_addressing(sig_alpha, prev_w_r_list, prev_w_lu):\n",
    "        return sig_alpha * prev_w_r + (1. - sig_alpha) * prev_w_lu     \n",
    "    \n",
    "    #least used weight vector\n",
    "    def least_used(w_u):\n",
    "        _, indices = tf.nn.top_k(w_u, k=self.memory_size)\n",
    "        w_lu = tf.reduce_sum(tf.one_hot(indices[:, -self.head_num:], depth=self.memory_size), axis=1)\n",
    "        return indices, w_lu\n",
    "\n",
    "    \n",
    "    #next we define the function called zero state for initializing all the states - \n",
    "    #controller state, read vector, weights and memory\n",
    "    def zero_state(self, batch_size, dtype):\n",
    "        one_hot_weight_vector = np.zeros([batch_size, self.memory_size])\n",
    "        one_hot_weight_vector[..., 0] = 1\n",
    "        one_hot_weight_vector = tf.constant(one_hot_weight_vector, dtype=tf.float32)\n",
    "        with tf.variable_scope('init', reuse=self.reuse):\n",
    "            state = {\n",
    "                'controller_state': self.controller.zero_state(batch_size, dtype),\n",
    "                'read_vector_list': [tf.zeros([batch_size, self.memory_vector_dim])\n",
    "                                     for _ in range(self.head_num)],\n",
    "                'w_r_list': [one_hot_weight_vector for _ in range(self.head_num)],\n",
    "                'w_u': one_hot_weight_vector,\n",
    "                'M': tf.constant(np.ones([batch_size, self.memory_size, self.memory_vector_dim]) * 1e-6, dtype=tf.float32)\n",
    "            }\n",
    "            return state"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [default]",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}