02_keras_LM_train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# keras training\n",
    "\n",
    "this file loads our processed data and trains a recurrent language model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from keras import backend as K\n",
    "from keras.models import Model, load_model\n",
    "from keras.layers import Input, Embedding, Dropout, LSTM, Lambda, Dense, Activation \n",
    "from keras.callbacks import ModelCheckpoint\n",
    "import h5py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# restrict GPU usage here, if using multi-gpu\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read in data\n",
    "cardtext = [list(x) for x in list(np.load('data/card_texts.npy'))]\n",
    "c2i = np.load('data/c2i.npy').item()\n",
    "i2c = np.load('data/i2c.npy').item()\n",
    "xcards = np.load('data/xcards.npy')\n",
    "ycards = np.load('data/ycards.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add axis for sparse_categorical_Crossentropy\n",
    "ycards = ycards[:, :, np.newaxis]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## define the hyperparameters\n",
    "\n",
    "**dropout rate** : prevents over-fitting; by 'ignoring' every few characters, the language model must learn to generalize. typical values are 0.25 to 0.50  \n",
    "**embedding size** : the size of the character embeddings, which are learned through training  \n",
    "**hidden_size** : the size of the LSTM gates and cells; i.e. the size of its 'memory'  \n",
    "**vocab_size** : the model will predict one of *n* characters where *n* is the vocabulary size  \n",
    "**batch size** : we will use *minibatch gradient descent*; this is the number of examples we will train on each batch  \n",
    "**number of epochs** : one *epoch* is one pass through all the training data  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set parameters\n",
    "DROP_RATE = 0.50              # dropout: between 0.25 and 0.5 is common\n",
    "HIDDEN_SIZE = 500             # lstm feature vector size\n",
    "MAX_Y_LEN = ycards.shape[1]   # maximum card length\n",
    "VOCAB_SIZE = len(c2i.keys())  # number of characters\n",
    "BATCH_SIZE = 32               # cards per batch\n",
    "NUM_EPOCHS = 10               # number of epochs to train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## define the model\n",
    "\n",
    "the model we will use is a *recurrent language model*. essentially, our network will predict the next character, given the previous chracters it has seen/generated. we could try to help the network realize that it is at the beginning of a card by initializing the states to a fixed value such as zeroes (using `initial_state`), but we will leave the state initialization random here, and rely on the initial start-of-sentence token to signal to the network that we are starting a card. hopefully, the random initial state might help randomize the generated cards.\n",
    "\n",
    "we have already divided the cards into lists of characters, *indexed* the strings into integer arrays, and *padded* the arrays to a fixed length in the previous files. we also created input and output sequences that are offset by one, such that the first element of the *output* corresponds with the *second* element of the input etc. this is because we will train the model with *teacher forcing* : at each step, we will input the *true* character, and induce the network to output the next element. on decode, of course, since we are randomly generating the card sequences, we must input the *actual* previous output.\n",
    "\n",
    "due to this, our training and decode networks are slightly different. our training network takes full sequence inputs, and outputs full sequence outputs (the inputs, offset by one). this is because we already know the full sequences we are training on: the actual cards. on decode, we want new cards, so we will generate each character at a time, and feed *that* predicted character (sampled randomly from the softmax distribution, for randomness) back into the LSTM to generate the next character. because the LSTM relies on a 'memory' of what it has already generated, we must also input the previous *states*, which we can do with `return_sequences=True`.\n",
    "\n",
    "this network is adapted from the decoder in [the keras blog seeq2seq article](https://blog.keras.io/a-ten-minute-introduction-to-sequence-to-sequence-learning-in-keras.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the decoder, using `encoder_states` as initial state.\n",
    "decoder_input  = Input(shape=(MAX_Y_LEN, ), name='lm_input')\n",
    "decoder_embed  = Embedding(VOCAB_SIZE, HIDDEN_SIZE, \n",
    "                           mask_zero=True, trainable=True, name='lm_emb')\n",
    "decoder_lstm1  = LSTM(HIDDEN_SIZE, \n",
    "                      return_sequences=True, \n",
    "                      return_state=True, \n",
    "                      name='lm_lstm1')\n",
    "decoder_lstm2  = LSTM(HIDDEN_SIZE, \n",
    "                      return_sequences=True, \n",
    "                      return_state=True, \n",
    "                      name='lm_lstm2')\n",
    "\n",
    "decoder_dense_1  = Dense(HIDDEN_SIZE, activation='relu', name='lm_dns_1')\n",
    "# decoder_dense_2  = Dense(VOCAB_SIZE, activation='softmax', name='lm_dns_final')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## weight tying\n",
    "\n",
    "this concept is from Ofir Press & Lior Wolf 2017, [\"Using the Output Embedding to Improve Language Models\"](https://arxiv.org/pdf/1608.05859.pdf)\n",
    "\n",
    "a conceptual outline by Ofir Press can be seen [on his blog](http://ofir.io/Neural-Language-Modeling-From-Scratch/)\n",
    "\n",
    "they define a Language Model generally as taking a word _c_, represented by a one-hot vector, embedding it into a dense vector with a weight matrix **U**, doing some computation on it (passing it through two LSTM layers, in this model, as well as in the blog example) to get a dense vector _h_, and then converting this back to a word prediction using matrix **V** followed by a softmax activation i.e. a `Dense` layer in `keras`. however, they note that _c_ and _h_ both share the property of being 'word vectors', and that the matrices **U** and **V** are of the same dimension (size of vocabulary x embedding size), and are conducting inverse operations (mapping words to dense vectors, and mapping dense vectors to words). so they propose \"weight tying\", which sets the weights **U** = **V** (though for the math, one is _transposed_). \n",
    "\n",
    "this can be demonstrated in `numpy` with the following: \n",
    "\n",
    "1. our parameters can be as follows:\n",
    "\n",
    "```\n",
    "word vocabulary = 10\n",
    "embedding dims  =  4\n",
    "```\n",
    "\n",
    "2. we can make an artificial word embedding `e` of size `(word vocabulary x embedding dims) == (10, 4)`\n",
    "   here each column = a 'word embedding' which here is filled with the n-th value:   \n",
    "\n",
    "```\n",
    ">>> e\n",
    "array([[ 1,  1,  1,  1],\n",
    "       [ 2,  2,  2,  2],\n",
    "       [ 3,  3,  3,  3],\n",
    "       [ 4,  4,  4,  4],\n",
    "       [ 5,  5,  5,  5],\n",
    "       [ 6,  6,  6,  6],\n",
    "       [ 7,  7,  7,  7],\n",
    "       [ 8,  8,  8,  8],\n",
    "       [ 9,  9,  9,  9],\n",
    "       [10, 10, 10, 10]])\n",
    "```\n",
    "\n",
    "3. a word is represented as a one-hot vector of length `word vocabulary`.  \n",
    "   here `w` is the vector for the third word in the index\n",
    "\n",
    "```\n",
    ">>> w\n",
    "array([0, 0, 1, 0, 0, 0, 0, 0, 0, 0])\n",
    "```\n",
    "\n",
    "4. to retrieve the embedding, we use the dot product,  \n",
    "   which multiplies each non-target row by 0 and the target embedding row by 1:\n",
    "\n",
    "```\n",
    ">>> np.dot(w, e)\n",
    "array([3, 3, 3, 3])\n",
    "```\n",
    "\n",
    "5. now we assume we do some calculation here and get an LSTM output vector `h`  of `embedding size`\n",
    "\n",
    "```\n",
    ">>> h\n",
    "array([2, 2, 2, 2])\n",
    "```\n",
    "\n",
    "6. then we can expand back to a 10-word vocabulary using `e.T` i.e. setting **U** = **V**   \n",
    "   (this only applies when our input and outut vocabularies are the same)  \n",
    "   (again, this means that the LSTM hidden size must `==` the embedding size)  \n",
    "\n",
    "```\n",
    ">>> np.dot(h, e.T)\n",
    "array([ 8, 16, 24, 32, 40, 48, 56, 64, 72, 80])\n",
    "```\n",
    "\n",
    "7. this output will then be turned into a probability distribution over the ten words using softmax.  \n",
    "\n",
    "we can use a `Lambda` layer that will take the recurrent outputs and `dot` them with the transposed embedding weights:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# weight-tying Lambda\n",
    "def weight_tying(layer_input):\n",
    "    result = K.dot(layer_input, K.transpose(decoder_embed.weights[0]))\n",
    "    return result\n",
    "\n",
    "decoder_dense_2 = Lambda(weight_tying, name='weight_tying')\n",
    "decoder_dense_3 = Activation('softmax')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the actual model\n",
    "x = decoder_embed(decoder_input)\n",
    "x = Dropout(DROP_RATE)(x)\n",
    "x, h1, c1 = decoder_lstm1(x)\n",
    "x = Dropout(DROP_RATE)(x)\n",
    "x, h2, c2 = decoder_lstm2(x)\n",
    "x = Dropout(DROP_RATE)(x)\n",
    "x = decoder_dense_1(x)\n",
    "x = Dropout(DROP_RATE)(x)\n",
    "x = decoder_dense_2(x)\n",
    "x = decoder_dense_3(x)\n",
    "\n",
    "model = Model(decoder_input, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "lm_input (InputLayer)        (None, 256)               0         \n",
      "_________________________________________________________________\n",
      "lm_emb (Embedding)           (None, 256, 500)          51000     \n",
      "_________________________________________________________________\n",
      "dropout_1 (Dropout)          (None, 256, 500)          0         \n",
      "_________________________________________________________________\n",
      "lm_lstm1 (LSTM)              [(None, 256, 500), (None, 2002000   \n",
      "_________________________________________________________________\n",
      "dropout_2 (Dropout)          (None, 256, 500)          0         \n",
      "_________________________________________________________________\n",
      "lm_lstm2 (LSTM)              [(None, 256, 500), (None, 2002000   \n",
      "_________________________________________________________________\n",
      "dropout_3 (Dropout)          (None, 256, 500)          0         \n",
      "_________________________________________________________________\n",
      "lm_dns_1 (Dense)             (None, 256, 500)          250500    \n",
      "_________________________________________________________________\n",
      "dropout_4 (Dropout)          (None, 256, 500)          0         \n",
      "_________________________________________________________________\n",
      "weight_tying (Lambda)        (None, 256, 102)          0         \n",
      "_________________________________________________________________\n",
      "activation_1 (Activation)    (None, 256, 102)          0         \n",
      "=================================================================\n",
      "Total params: 4,305,500\n",
      "Trainable params: 4,305,500\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compile\n",
    "model.compile(loss='sparse_categorical_crossentropy',\n",
    "              optimizer='adam', metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## model training\n",
    "\n",
    "we define a `ModelCheckpoint` that will save models as we train, in case the model training takes a long time. we then `fit` the model to train. we use `verbose=2` to view per-epoch stats; `verbose=1`, while it provides per-batch stats, can freeze Jupyter Lab, and `TQDMNotebook` doesn't work with Jupyter Lab yet (AFAIK).\n",
    "\n",
    "we then save weights at the end of training (and re-load them to test). we have a (commented by default) cell for loading weights before training, to allow us to continue training a partially-trained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "cpoint = ModelCheckpoint('model/weights.{epoch:04d}-{loss:.4f}.h5', \n",
    "                         monitor='loss',\n",
    "                        save_best_only=True,\n",
    "                        save_weights_only=True,\n",
    "                        period=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.load_weights('model/weights_final.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n",
      "30303/30303 [==============================] - 879s 29ms/step - loss: 1.5141 - acc: 0.5999\n",
      "Epoch 2/10\n",
      "30303/30303 [==============================] - 877s 29ms/step - loss: 0.8181 - acc: 0.7649\n",
      "Epoch 3/10\n",
      "30303/30303 [==============================] - 877s 29ms/step - loss: 0.5943 - acc: 0.8276\n",
      "Epoch 4/10\n",
      "30303/30303 [==============================] - 878s 29ms/step - loss: 0.5009 - acc: 0.8536\n",
      "Epoch 5/10\n",
      "30303/30303 [==============================] - 876s 29ms/step - loss: 0.4486 - acc: 0.8681\n",
      "Epoch 6/10\n",
      "30303/30303 [==============================] - 878s 29ms/step - loss: 0.4148 - acc: 0.8775\n",
      "Epoch 7/10\n",
      "30303/30303 [==============================] - 878s 29ms/step - loss: 0.3907 - acc: 0.8841\n",
      "Epoch 8/10\n",
      "30303/30303 [==============================] - 878s 29ms/step - loss: 0.3746 - acc: 0.8885\n",
      "Epoch 9/10\n",
      "30303/30303 [==============================] - 877s 29ms/step - loss: 0.3623 - acc: 0.8918\n",
      "Epoch 10/10\n",
      "30303/30303 [==============================] - 878s 29ms/step - loss: 0.3495 - acc: 0.8954\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7fe7dfcbfe48>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.fit(xcards, ycards, \n",
    "          batch_size=BATCH_SIZE, \n",
    "          epochs=NUM_EPOCHS, \n",
    "          callbacks=[cpoint], \n",
    "          verbose=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_weights('model/weights_tiedfinal.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_weights('model/weights_tiedfinal.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save architecture with json\n",
    "with open('model/weights_tiedfinal.json', 'w') as f:\n",
    "    f.write(model.to_json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "atlas",
   "language": "python",
   "name": "atlas"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}