lstmTreeNotebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# Combining Symbolic Expressions and Black-box Function Evaluations in Neural Programs\n",
    "\n",
    "*written by Forough Arabshahi*\n",
    "\n",
    "This notebook presents the code for [this](https://arxiv.org/abs/1801.04342) paper.\n",
    "\n",
    "The focus of the paper is on neural programing. Traditional methods to neural programming either rely on black-box function evaluation data or rely on program execution traces for training the neteworks. Both of these methods lack generalizability. Black-box function evaluations do not contain any infomation about the structure of the problem. Porgram execution traces, On the other hand, are expensive to collect resulting in a lack of domain coverage.\n",
    "\n",
    "In many problems, one has access to symbolic representation of the problem that encodes the relationships between the given variables and functions in a succinct manner. For example, declarative programs greatly simplify parallel programs through the generation of symbolic computation graphs. As another example, the properties of mathematical functions are encoded through symbolic expressions. E.g. symbolic expression $x+y = y+x$ encodes the commutative property of the addition function. Therefore, symbolic expressions efficiently represent the properties of the problem, preserve structure and are readily accessible. Thus, they are a great alternative to black-box function evaluations and program execution traces. However, by themselves, they do not enable the desired task of function evaluation. \n",
    "\n",
    "The main contribution of this paper is combining symbolic expressions with black-box function evaluation data for training neural programmers. This results in generalizable models that are also capable of function evaluation. The paper studies modeling mathematical equations and it shows that this combination allows one to model up to 28 mathematical functions that scales up the domain by about $3.5\\times$ while increasing the complexity of the mathematical equations compared to the state-of-the-art. The authors propose a dataset generation strategy that generates a balanced dataset of symbolic and function evaluation data with a good coverage of the domain under study. They propose using Tree LSTM's that mirror the parse-tree of the symbolic and function evaluation expressions for modeling mathematical equations. The paper finally evaluates the model on tasks such as equation verification and equation completion.\n",
    "\n",
    "***\n",
    "\n",
    "## Implementation Details\n",
    "\n",
    "In this notebook, we present the code for training the tree LSTM model for the task of equation verification. There are also another notebooks attached, that covers the dataset generation.\n",
    "The code is implemented in Pythin 2.7 and uses [MxNet](https://mxnet.incubator.apache.org) as the underlying deep learning platform.\n",
    "\n",
    "### 1. Importing Modules\n",
    "<a id=\"sec:import\"></a>\n",
    "\n",
    "Let us start with importing the relevant modules.\n",
    "\n",
    "Our designed neuralAlgonometry module is a module containing the tree class *EquationTree* and several other useful functions such as functions *readBlockJsonEquations* for reading the input data and *saveResults* for saving the results. Bellow is an example of equations that are represented using the EquationTree class.\n",
    "\n",
    "$\\sin^2(\\theta) + \\cos^2(\\theta) = 1$ | $\\sin(-2.5) = -0.6$ | Decimal expression tree for $2.5$\n",
    "- | - | -\n",
    "<img src=\"figs/eTree.png\", width=\"300\", height=\"300\"/>  | <img src=\"figs/numTree.png\", width=\"300\", height=\"300\"/> | <img src=\"figs/num_tree.png\", width=\"300\", height=\"300\"/>\n",
    "\n",
    "nnTreeMain is a module that contains the neural network tree classes. We use MxNet's [bucketingModule](https://mxnet.incubator.apache.org/how_to/bucketing.html) for implementing dynamic networks. Class *lstmTreeInpOutMSE* implements treeLSTMs for the combination of symbolic and black-box function evaluation data. The implementation of the baseline models used in the paper are also present in nnTreeMain and are called *nnTreeInpOutMSE*, *LSTMtree* and *nnTree* for treeNNs with a combination of symbolic and function evaluation data, treeLSTMs for symbolic data and treeNNs for symbolic data, respectively. Replacing lstmTreeInpOutMSE with any of these calsses perform training and equation verification for these models.\n",
    "\n",
    "*BucketEqIteratorInpOut* is the data iterator class used by the bucketing module and *bucketIndex* is the class that is passed to the *sym_gen* function of the bukcetingModule. precision, recall and accuracy are subclasses of mx.metric.EvalMetric.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/forougharabshahi/anaconda2/lib/python2.7/site-packages/urllib3/contrib/pyopenssl.py:46: DeprecationWarning: OpenSSL.rand is deprecated - you should use os.urandom instead\n",
      "  import OpenSSL.SSL\n",
      "neuralAlgonometry.py:13: DeprecationWarning: The compiler package is deprecated and removed in Python 3.x.\n",
      "  import compiler\n"
     ]
    }
   ],
   "source": [
    "# importing utilities\n",
    "import mxnet as mx\n",
    "import numpy as np\n",
    "import random\n",
    "import re\n",
    "import copy\n",
    "from neuralAlgonometry import buildNNTree, encode_equations, EquationTree, readBlockJsonEquations,\\\n",
    "                              saveResults, updateWorksheet, dumpParams, writeJson, putNodeInTree, encodeTerminals, \\\n",
    "                              genMSEeq\n",
    "from nnTreeMain import lstmTreeInpOutMSE, BucketEqIteratorInpOut, bucketIndex, precision, recall, Accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. One-hot encoding of the terminals\n",
    "<a id=\"sec:one-hot\"></a>\n",
    "\n",
    "As stated in Sectin 2 of [the paper](https://openreview.net/forum?id=Hksj2WWAW&noteId=Hksj2WWAW), the terminals in the grammar are the leaves of the expression tree. In the neural network, these terminals are represented using the one-hot encoding. This function creates a dictionary containing the key-value pairs (terminal:index), wehre terminal, is one of the terminals, e.g. symbol $x, y$ or integers $1,2,3,\\dots$ and index is the unique one-hot index. The terminals used in the paper are listed below. It is worth noting that these are the terminals for symbolic experssions only. The terminals for function evaluations are floating numbers of precision $2$ in the range $[-3.14, 3.14]$ and are represented using their decimal tree expanssions. This means that they can all be represented using the integers listed below. More explanation about how these floating point numbers are inputed to the neural network will be explained in [this section](#sec:sym_gen)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "functionDictionary: {'Rational_-1/2': 26, 'Exp1': 8, 'Pi': 9, 'Symbol_var_3': 3, 'Symbol_var_2': 2, 'Symbol_var_1': 1, 'Symbol_var_0': 0, 'NegativeOne': 5, 'Symbol_var_4': 4, 'NaN': 6, 'Half': 11, 'Integer_10': 22, 'Integer_6': 18, 'Integer_7': 19, 'Integer_4': 16, 'Integer_5': 17, 'Integer_2': 14, 'Integer_3': 15, 'Integer_0': 12, 'Integer_1': 13, 'Infinity': 7, 'Integer_8': 20, 'Integer_9': 21, 'Float': 28, 'One': 10, 'Rational_2/5': 25, 'Rational_0': 27, 'Integer_-2': 23, 'Integer_-3': 24}\n"
     ]
    }
   ],
   "source": [
    "# terminals:\n",
    "variables = ['Symbol']\n",
    "consts = ['NegativeOne', 'NaN', 'Infinity', 'Exp1', 'Pi', 'One',\n",
    "          'Half', 'Integer', 'Rational', 'Float']\n",
    "terminals = []\n",
    "terminals.extend(variables) \n",
    "terminals.extend(consts)\n",
    "\n",
    "intList = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, -2, -3]\n",
    "ratList = ['2/5', '-1/2', '0']\n",
    "floatList = [0.7]\n",
    "varList = ['var_%d'%(d) for d in range(5)]\n",
    "\n",
    "functionDictionary = encodeTerminals(terminals, intList, ratList, floatList, varList)\n",
    "print \"functionDictionary:\", functionDictionary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Model Parameters\n",
    "<a id=\"sec:params\"></a>\n",
    "\n",
    "The model hyper-parameters are given in this block. You can change these hyper-parameters to tune the neural network.\n",
    "\n",
    "As stated in the paper, the depth of an equation can be indicative of the equation's complexity. *trainDepth* and *testDepth* refers to the depth of equations included in training and testing. These are used for generating the results given in Tables 2 and 3 of the paper to assess the generalizability of the model. When trainDepth = testDepth = [1,2,3,4], train and test sets include all the equations of depths 1 through 4 and the performance of the model is assessed on unseen equations in the test set. These reults are shown in Table 2. In order to reproduce the results of Table 3 set:\n",
    "\n",
    "trainDepth = [1,2,3], testDepth = [4]\n",
    "\n",
    "and\n",
    "\n",
    "trainDepth = [1,3,4], testDepth = [2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Model parameters and training setup \n",
    "params = None\n",
    "contexts = mx.cpu(0)\n",
    "num_hidden = 50\n",
    "vocabSize = len(functionDictionary)\n",
    "emb_dimension = 50\n",
    "out_dimension = 1\n",
    "batch_size = 1\n",
    "splitRatio = 0.8 # proportion of equations in the train set\n",
    "devSplit = 1\n",
    "excludeFunc = ''\n",
    "trainDepth = [1,2,3,4]\n",
    "testDepth = [1,2,3,4]\n",
    "dropout = 0.2\n",
    "lr = 0.00001 # learning rate\n",
    "mom = 0.7 # momentum\n",
    "wd = 0.001 # weight decay\n",
    "optimizer = \"adam\" # name of optimizer\n",
    "num_epochs = 2 # number of training epochs\n",
    "load_epoch = 0 # load pre-trained model starting from epoch number load_epoch\n",
    "model_prefix = \"notebookModel/model0/trained_model \" # path to save model checkpoints\n",
    "kv_store = \"device\" # KV_store \n",
    "\n",
    "# refer to section 1. for an explanation about different neural network classes\n",
    "tTypeStr = 'lstmTreeInpOutMSE' # neural network type. Other options: 'nnTreeInpOut', 'nnTree', 'lstmTree'\n",
    "tType = lstmTreeInpOutMSE  # neural network type. Other options: nnTreeInpOut, nnTree, lstmTree\n",
    "\n",
    "# path to data: below is a smaller dataset that runs faster.\n",
    "# file data/data4000_orig_inpuOut_with_neg.json is the data used in the paper\n",
    "path = \"data/data1000_depth4_inpuOut.json\" \n",
    "result_path = \"notebookModel/model0/res\" # path for saving results such as train/test accuracies and other metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Reading data\n",
    "<a id=\"sec:read\"></a>\n",
    "\n",
    "In this section we explain how to load the data generated using the dataset generation method. Function *readBlockJsonEquations* available in module neuralAlgonometry loads equations saved in a json format. You can input the train/test splitting ratio *splitRatio* and if desired a *devSplit* which holds out a portion of the train data for validation. This is set to 1 in this notebook meaning no validation data is used, but one can set it to, say, 0.9 to keep $10%$ of the data for validation. Validation data can be useful for assessing the models overfitting behavior during training.\n",
    "\n",
    "The returned data includes the train/test/devEquations that are lists of objects of class *EquationTree*. train/test/devVars contains the list of variables in each equation. train/test/devLabels is a list of labels corresponding to each equation. Labels are either <font color=blue>mx.nd.array([0], dtype='float32')</font> , or <font color=blue>mx.nd.array([1], dtype='float32')</font> for incorrect and correct equations, respectively.\n",
    "\n",
    "It should be noted that since the bucketing module needs to see all the neural network blocks when forming the first computation graph, The first equation is a synthetic equation including all the functions and terminals in the grammar.  Flag *containsNumeric* indicates weather the data contains function evaluations or if it only contains symbolic expressions. If the data contains fnuction evaluation data set this flag to True. In that case the first equation will be appended with a Number block.\n",
    "\n",
    "In case the random seed is not sat, the created data split may not be reproducible. Therefore, one can save the original split, if further analysis needs to be done on the data using the saved data. Function *writeJson* is a function available in the neuralAlgonometry module. It saves the trees in a the json format. The format of the equations are described in the *generateData.ipynb* notebook. This is commented out in the last three lines of the cell below. But can be uncommented if it is desirable to save the splits. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "random.seed(1)\n",
    "mx.random.seed(1)\n",
    "np.random.seed(1)\n",
    "[trainEquations, trainVars, _, trainLabels,\n",
    "devEquations, devVars, _, devLabels,\n",
    "testEquations, testVars, _, testLabels] \\\n",
    "          = readBlockJsonEquations(path, trainDepth=trainDepth, testDepth=testDepth,\n",
    "                                   excludeFunc=excludeFunc, splitRatio=splitRatio, devSplit=devSplit, containsNumeric=True)\n",
    "    \n",
    "# uncomment if need to store the original data split\n",
    "# writeJson(result_path+'trainData.json', trainEquations, ranges, trainVars, trainLabels, 6)\n",
    "# writeJson(result_path+'devData.json', devEquations, ranges, devVars, devLabels, 6)\n",
    "# writeJson(result_path+'testData.json', testEquations, ranges, testVars, testLabels, 6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to handle both symbolic and input-output expressions we use different loss functions and terminal types for each type of expression. The loss function for symbolic expressions is the logistic loss and the loss function for the input-output expressions is the MSE loss. In the following code block function genMSEeq, converts input-output expressions into a format compatible with the MSE loss function. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "eqToChange = copy.deepcopy(trainEquations[0])\n",
    "origNode = copy.deepcopy(eqToChange.findNode(106))\n",
    "child1 = EquationTree(func='Number', varname='1.1')\n",
    "# child1 = EquationTree(func='NumberDecoder', args=[child1])\n",
    "rootNode = EquationTree(func='Add', args=[child1, origNode])\n",
    "changedEq = putNodeInTree(eqToChange, rootNode, 106)\n",
    "changedEq.args[0] = EquationTree(func='NumberDecoder', args=[changedEq.args[0]], varname='')\n",
    "changedEq.args[1] = EquationTree(func='NumberDecoder', args=[changedEq.args[1]], varname='')\n",
    "changedEq.enumerize_queue()\n",
    "# changedEq.preOrder()\n",
    "trainEquations[0] = copy.deepcopy(changedEq)\n",
    "# devEquations[0]   = copy.deepcopy(changedEq)\n",
    "testEquations[0]  = copy.deepcopy(changedEq)\n",
    "\n",
    "oldTrainEquations = trainEquations\n",
    "trainEquations = []\n",
    "oldTrainVars = trainVars\n",
    "trainVars = []\n",
    "oldTrainLabels = trainLabels\n",
    "trainLabels = []\n",
    "domainNumbers = []\n",
    "for ind, tup in enumerate(zip(oldTrainEquations, oldTrainVars, oldTrainLabels)):\n",
    "    te, tv, tl = tup\n",
    "    if te.isNumeric() and ind!=0:\n",
    "        if tl.asnumpy()==1:\n",
    "            mseTe, mseTl, mseTv = genMSEeq(te, tl, tv, domainNumbers)\n",
    "            # print 'label outside:', mseTl[2].asnumpy()\n",
    "            trainEquations.extend(mseTe)\n",
    "            trainVars.extend(mseTv)\n",
    "            trainLabels.extend(mseTl)\n",
    "        else:\n",
    "            continue\n",
    "    else:\n",
    "        trainEquations.append(te)\n",
    "        trainVars.append(tv)\n",
    "        trainLabels.append(tl)\n",
    "\n",
    "assert len(domainNumbers)==len(set(domainNumbers)), 'domainNumbers contains repeated entries'\n",
    "assert len(trainLabels) == len(trainEquations), 'labels and equations have different lengths'\n",
    "assert len(trainLabels) == len(trainVars)\n",
    "\n",
    "oldTestEquations = testEquations\n",
    "testEquations = []\n",
    "oldTestVars = testVars\n",
    "testVars = []\n",
    "oldTestLabels = testLabels\n",
    "testLabels = []\n",
    "testDomainNumbers = []\n",
    "for ind, tup in enumerate(zip(oldTestEquations, oldTestVars, oldTestLabels)):\n",
    "    te, tv, tl = tup\n",
    "    if te.isNumeric() and ind!=0:\n",
    "        if tl.asnumpy()==1:\n",
    "            mseTe, mseTl, mseTv = genMSEeq(te, tl, tv, testDomainNumbers)\n",
    "            testEquations.extend(mseTe)\n",
    "            testVars.extend(mseTv)\n",
    "            testLabels.extend(mseTl)\n",
    "        else:\n",
    "            continue\n",
    "    else:\n",
    "        testEquations.append(te)\n",
    "        testVars.append(tv)\n",
    "        testLabels.append(tl)\n",
    "assert len(testDomainNumbers) == len(set(testDomainNumbers))\n",
    "assert len(testEquations) == len(testLabels)\n",
    "assert len(testEquations) == len(testVars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Construct neural network classes\n",
    "<a id=\"sec:nn\"></a>\n",
    "\n",
    "In this block we construct the neural network classes for each equation in the train and test set. If you have sat devSplit to something other than 1, then you should also construct the network for your validation set. This can be done by uncommenting the last part of the code in the cell below.\n",
    "\n",
    "Function *buildNNTree*, that is implemented in module neuralAlgonometry, traverses the input equation and constructs a treeLSTM (or another model if tType is different) that mirrors the equation's structure.  \n",
    "\n",
    "The figures below depict the neural network constrcted using this function for the equations in [this section](#sec:import)\n",
    "\n",
    "$\\sin^2(\\theta) + \\cos^2(\\theta) = 1$ | $\\sin(-2.5) = -0.6$ \n",
    "- | -\n",
    "<img src=\"figs/network_sym.png\", width=\"400\", height=\"400\"/>  | <img src=\"figs/network_num.png\", width=\"400\", height=\"400\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "trainSamples = []\n",
    "trainDataNames = []\n",
    "for i, eq in enumerate(trainEquations):\n",
    "    currNNTree = buildNNTree(treeType=tType , parsedEquation=eq, \n",
    "                                num_hidden=num_hidden, params=params, \n",
    "                                emb_dimension=emb_dimension, dropout=dropout)\n",
    "\n",
    "    [currDataNames, _] = currNNTree.getDataNames(dataNames=[], nodeNumbers=[])\n",
    "    trainDataNames.append(currDataNames)\n",
    "    trainSamples.append(currNNTree)\n",
    "\n",
    "testSamples = []\n",
    "testDataNames = []\n",
    "for i, eq in enumerate(testEquations):\n",
    "    currNNTree = buildNNTree(treeType=tType , parsedEquation=eq, \n",
    "                                num_hidden=num_hidden, params=params, \n",
    "                                emb_dimension=emb_dimension, dropout=dropout)\n",
    "\n",
    "    [currDataNames, _] = currNNTree.getDataNames(dataNames=[], nodeNumbers=[])\n",
    "    testDataNames.append(currDataNames)\n",
    "    testSamples.append(currNNTree)\n",
    "    \n",
    "# devSamples = []\n",
    "# devDataNames = []\n",
    "# for i, eq in enumerate(devEquations):\n",
    "#     currNNTree = buildNNTree(treeType=tType , parsedEquation=eq, \n",
    "#                            num_hidden=num_hidden, params=params, \n",
    "#                            emb_dimension=emb_dimension, dropout=dropout)\n",
    "\n",
    "#     [currDataNames, _] = currNNTree.getDataNames(dataNames=[], nodeNumbers=[])\n",
    "#     devDataNames.append(currDataNames)\n",
    "#     devSamples.append(currNNTree)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. Construct data iterators\n",
    "<a id=\"sec:iter\"></a>\n",
    "\n",
    "Class *BucketEqIteratorInpOut* which is a subclass of <font color=blue>mx.io.DataIter</font>. It constructs the data iterators for the train and test equations. If you have sat devSplit to something other than 1, you need to have a data iterator for your validation set. This can be done by uncommenting the code below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "numTrainSamples = len(trainEquations)\n",
    "trainBuckets = list(xrange(numTrainSamples))\n",
    "\n",
    "numTestSamples = len(testEquations)\n",
    "testBuckets = list(xrange(numTestSamples))\n",
    "\n",
    "train_eq, _ = encode_equations(trainEquations, vocab=functionDictionary, invalid_label=-1, \n",
    "                                     invalid_key='\\n', start_label=0)\n",
    "data_train  = BucketEqIteratorInpOut(enEquations=train_eq, eqTreeList=trainSamples, batch_size=batch_size, \n",
    "                             buckets=trainBuckets, labels=trainLabels, vocabSize=len(functionDictionary),\n",
    "                                    invalid_label=-1)\n",
    "\n",
    "test_eq, _ = encode_equations(testEquations, vocab=functionDictionary, invalid_label=-1, \n",
    "                             invalid_key='\\n', start_label=0)\n",
    "data_test  = BucketEqIteratorInpOut(enEquations=test_eq, eqTreeList=testSamples, batch_size=batch_size, \n",
    "                             buckets=testBuckets, labels=testLabels, vocabSize=len(functionDictionary),\n",
    "                                    invalid_label=-1, devFlag=1)\n",
    "\n",
    "\n",
    "# numDevSamples = len(devEquations)\n",
    "# devBuckets = list(xrange(numDevSamples))\n",
    "# dev_eq, _ = encode_equations(devEquations, vocab=functionDictionary, invalid_label=-1, \n",
    "#                              invalid_key='\\n', start_label=0)\n",
    "# data_dev  = BucketEqIteratorInpOut(enEquations=dev_eq, eqTreeList=devSamples, batch_size=batch_size, \n",
    "#                              buckets=devBuckets, labels=devLabels, vocabSize=len(functionDictionary),\n",
    "#                                     invalid_label=-1, devFlag=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. Symbol generator function for the bucketing module\n",
    "<a id=\"sec:sym_gen\"></a>\n",
    "\n",
    "Defining the sym_gen function for MxNet's [bucketing module](https://mxnet.incubator.apache.org/how_to/bucketing.html) and constructing the neural network model that forms the computation graph. This function returns the prediction symbol as well as the data names and label names. *data_names_corr* is a list that contains the data which contains the terminals' names. \n",
    "\n",
    "For the terminals that are represented with their one-hot encoding, we have a one-layer neural network block that is responsible for embedding the representation of that cell (cell $W_{sym}$ in [section 5](#sec:nn)). For floting point numbers in the range $[-3.14, 3.14]$ we have a two-layer neural network that is responsible for encoding its values (cell $W_{num}$ in [section 5](#sec:nn)). Floating point numbers are inputed as is to the neural network.\n",
    "\n",
    "The final model is an instance of MxNet's *BucketingModule*. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "cell = trainSamples[0]\n",
    "\n",
    "def sym_gen(bucketIndexObj):\n",
    "\n",
    "    label = mx.sym.Variable('softmax_label')\n",
    "\n",
    "    bucketIDX = bucketIndexObj.bucketIDX\n",
    "    devFlag = bucketIndexObj.devFlag\n",
    "\n",
    "    if devFlag == 0:\n",
    "        tree = trainSamples[bucketIDX]\n",
    "    else:\n",
    "        tree = testSamples[bucketIDX]\n",
    "\n",
    "\n",
    "    [dataNames, nodeNumbers] = tree.getDataNames(dataNames=[], nodeNumbers=[])\n",
    "    data_names_corr = [dataNames[i]+'_%d'%(nodeNumbers[i]) for i in range(len(dataNames))]\n",
    "    nameDict = {}\n",
    "    for i, dn in enumerate(dataNames):\n",
    "        if dn not in nameDict:\n",
    "            nameDict[dn+'_%d'%(nodeNumbers[i])] = mx.sym.Variable(name=dn+'_%d'%(nodeNumbers[i]))\n",
    "        else:\n",
    "            raise AssertionError(\"data name should not have been in the dictionary\")\n",
    "\n",
    "    if 'lstm' in tTypeStr:\n",
    "        outputs, _ = tree.unroll(nameDict)\n",
    "    else:\n",
    "        outputs = tree.unroll(nameDict)\n",
    "\n",
    "    pred = mx.sym.LogisticRegressionOutput(data=outputs, label=label, name='softmax')\n",
    "\n",
    "    return pred, (data_names_corr), ('softmax_label',)\n",
    "\n",
    "model = mx.mod.BucketingModule(\n",
    "    sym_gen             = sym_gen,\n",
    "    default_bucket_key  = bucketIndex(0, 0),\n",
    "    context             = contexts,\n",
    "    fixed_param_names  = [str(tTypeStr)+'_Equality_i2h_weight'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7. Training\n",
    "<a id=\"sec:train\"></a>\n",
    "\n",
    "In this section we perform the training using model.fit($\\dots$). Once ran, the training and test accuracies will be shown in the output log. Function *saveResults*, saves the precision, recall and accuracy metrics for the train and test data in the *result_path* whose value is sat in [Section 3](#sec:params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2017-11-21 20:43:46,063 Epoch[0] Train-accuracy=0.544708\n",
      "2017-11-21 20:43:46,064 Epoch[0] Train-precision=0.480236\n",
      "2017-11-21 20:43:46,065 Epoch[0] Train-recall=0.762740\n",
      "2017-11-21 20:43:46,066 Epoch[0] Time cost=50.873\n",
      "2017-11-21 20:43:46,075 Saved checkpoint to \"notebookModel/model0/trained_model -0001.params\"\n",
      "2017-11-21 20:43:56,582 Epoch[0] Validation-accuracy=0.554780\n",
      "2017-11-21 20:43:56,582 Epoch[0] Validation-precision=0.472826\n",
      "2017-11-21 20:43:56,583 Epoch[0] Validation-recall=0.739796\n",
      "2017-11-21 20:44:07,173 Epoch[1] Train-accuracy=0.575969\n",
      "2017-11-21 20:44:07,173 Epoch[1] Train-precision=0.503524\n",
      "2017-11-21 20:44:07,174 Epoch[1] Train-recall=0.728088\n",
      "2017-11-21 20:44:07,175 Epoch[1] Time cost=9.470\n",
      "2017-11-21 20:44:07,185 Saved checkpoint to \"notebookModel/model0/trained_model -0002.params\"\n",
      "2017-11-21 20:44:08,857 Epoch[1] Validation-accuracy=0.578507\n",
      "2017-11-21 20:44:08,858 Epoch[1] Validation-precision=0.491051\n",
      "2017-11-21 20:44:08,858 Epoch[1] Validation-recall=0.746599\n"
     ]
    }
   ],
   "source": [
    "import logging\n",
    "head = '%(asctime)-15s %(message)s'\n",
    "logging.basicConfig(level=logging.DEBUG, format=head)\n",
    "\n",
    "acc = Accuracy()\n",
    "prc = precision()\n",
    "rcl = recall()\n",
    "\n",
    "if load_epoch:\n",
    "    _, arg_params, aux_params = mx.rnn.load_rnn_checkpoint(\n",
    "        cell, model_prefix, load_epoch)\n",
    "else:\n",
    "    arg_params = None\n",
    "    aux_params = None\n",
    "\n",
    "opt_params = {\n",
    "'learning_rate': lr,\n",
    "'wd': wd\n",
    "}\n",
    "\n",
    "if optimizer not in ['adadelta', 'adagrad', 'adam', 'rmsprop']:\n",
    "    opt_params['momentum'] = mom\n",
    "\n",
    "model.fit(\n",
    "train_data          = data_train,\n",
    "eval_data           = data_test,\n",
    "kvstore             = kv_store,\n",
    "eval_metric         = [acc, prc, rcl],\n",
    "optimizer           = optimizer,\n",
    "optimizer_params    = opt_params,\n",
    "initializer         = mx.init.Mixed([str(tTypeStr)+'_Equality_i2h_weight', '.*'], \n",
    "                                [mx.init.One(), mx.init.Xavier(factor_type=\"in\", magnitude=2.34)]),\n",
    "arg_params          = arg_params,\n",
    "aux_params          = aux_params,\n",
    "begin_epoch         = load_epoch,\n",
    "num_epoch           = num_epochs,\n",
    "epoch_end_callback  = mx.rnn.do_rnn_checkpoint(cell, model_prefix, 1) \\\n",
    "                                               if model_prefix else None)\n",
    "\n",
    "accTrain = [acc.allVals[i] for i in range(0,len(acc.allVals),2)]\n",
    "accVal   = [acc.allVals[i] for i in range(1,len(acc.allVals),2)]\n",
    "prcTrain = [prc.allVals[i] for i in range(0,len(prc.allVals),2)]\n",
    "prcVal   = [prc.allVals[i] for i in range(1,len(prc.allVals),2)]\n",
    "rclTrain = [rcl.allVals[i] for i in range(0,len(rcl.allVals),2)]\n",
    "rclVal   = [rcl.allVals[i] for i in range(1,len(rcl.allVals),2)]\n",
    "\n",
    "trainMetrics = [accTrain, prcTrain, rclTrain]\n",
    "valMetrics   = [accVal,     prcVal,   rclVal]\n",
    "\n",
    "# args\n",
    "if result_path:\n",
    "    saveResults(result_path+'_train.json', {}, trainMetrics, valMetrics)\n",
    "    trainPredicts = model.predict(data_train).asnumpy()\n",
    "    np.save(result_path+'_train_predictions.npy', trainPredicts)\n",
    "    with open(result_path+'_train_labels.txt', 'wa') as labelFile:\n",
    "        for lbl in trainLabels:\n",
    "            labelFile.write('{0}\\n'.format(lbl.asscalar()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "### Appendix: Pre-setup\n",
    "\n",
    "**Note: ** MxNet's Parameter allow_extra_params should be sat to True as shown in [this commit](https://github.com/Mega-DatA-Lab/mxnet/commit/13505824699cfc39d8ea52537c56bd5aaf9639b6) for this code to work properly. This is used for handling dynamic graphs.\n",
    "\n",
    "**Note: ** Use the _update_params(...) function in [this commit](https://github.com/Mega-DatA-Lab/mxnet/commit/960af8aa713f00e4dd6240dcc4f03867e8ac9f23) in MxNet's [model.py](https://github.com/apache/incubator-mxnet/blob/master/python/mxnet/model.py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}