homeworks/homework10.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Homework 10 - Berkeley STAT 157\n",
    "\n",
    "**Your name: XX, SID YY, teammates A,B,C (Please add your name, SID and teammates to ease\n",
    "Ryan and Rachel to grade.)**\n",
    "\n",
    "- Please submit your homework through gradescope\n",
    "- Handout 4/25/2019, due 5/2/2019 by 4pm.\n",
    "- We mentioned in the lecture that we can use bidirectional recurrent layers in the encoder to obtain better feature presentations. In this homework, you are asked to change the 2nd code block to turn the uni-directional LSTM layers into bi-directional. \n",
    "- This is an individual homework, please submit your work individually. \n",
    "\n",
    "The sequence to sequence (seq2seq) model is based on the encoder-decoder architecture to generate a sequence output for a sequence input. Both the encoder and the decoder use recurrent neural networks to handle sequence inputs. The hidden state of the encoder is used directly to initialize the decoder hidden state to pass information from the encoder to the decoder.\n",
    "\n",
    "![The sequence to sequence model architecture.](http://d2l.ai/_images/seq2seq.svg)\n",
    "\n",
    "The layers in the encoder and the decoder are illustrated in the following figure. \n",
    "\n",
    "![](http://d2l.ai/_images/seq2seq-details.svg)\n",
    "\n",
    "In this section we will implement the seq2seq model to train on the machine translation dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "1"
    }
   },
   "outputs": [],
   "source": [
    "# you may need to update the d2l package such as\n",
    "# pip install d2l==0.9.1\n",
    "\n",
    "import time\n",
    "from mxnet import nd, init, gluon, autograd\n",
    "from mxnet.gluon import nn, rnn, loss as gloss\n",
    "import d2l"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Encoder\n",
    "\n",
    "In the encoder, we use the word embedding layer to obtain a feature index from the word index of the input language and then input it into a multi-level LSTM recurrent unit. The input for the encoder is a batch of sequences, which is 2-D tensor with shape (batch size, sequence length). It outputs both the LSTM outputs, e.g the hidden state, for each time step and the hidden state and memory cell of the last time step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "3"
    }
   },
   "outputs": [],
   "source": [
    "class Seq2SeqEncoder(d2l.Encoder):\n",
    "    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,\n",
    "                 dropout=0, **kwargs):\n",
    "        super(Seq2SeqEncoder, self).__init__(**kwargs)\n",
    "        self.embedding = nn.Embedding(vocab_size, embed_size)\n",
    "        self.rnn = rnn.LSTM(num_hiddens, num_layers, dropout=dropout)\n",
    "\n",
    "    def forward(self, X, *args):\n",
    "        X = self.embedding(X) # X shape: (batch_size, seq_len, embed_size)\n",
    "        X = X.swapaxes(0, 1)  # RNN needs first axes to be time\n",
    "        state = self.rnn.begin_state(batch_size=X.shape[1], ctx=X.context)\n",
    "        out, state = self.rnn(X, state)\n",
    "        # The shape of out is (seq_len, batch_size, num_hiddens).\n",
    "        # state contains the hidden state and the memory cell\n",
    "        # of the last time step, the shape is (num_layers, batch_size, num_hiddens)\n",
    "        return out, state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we will create a mini-batch sequence input with a batch size of 4 and 7 time steps. We assume the number of hidden layers of the LSTM unit is 2 and the number of hidden units is 16. The output shape returned by the encoder after performing forward calculation on the input is (number of time steps, batch size, number of hidden units). The shape of the multi-layer hidden state of the gated recurrent unit in the final time step is (number of hidden layers, batch size, number of hidden units). For the gated recurrent unit, the `state` list contains only one element, which is the hidden state. If long short-term memory is used, the `state` list will also contain another element, which is the memory cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "4"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((7, 4, 16), 2, (2, 4, 16), (2, 4, 16))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoder = Seq2SeqEncoder(vocab_size=10, embed_size=8,\n",
    "                         num_hiddens=16, num_layers=2)\n",
    "encoder.initialize()\n",
    "X = nd.zeros((4, 7))\n",
    "output, state = encoder(X)\n",
    "output.shape, len(state), state[0].shape, state[1].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Decoder\n",
    "\n",
    "We directly use the hidden state of the encoder in the final time step as the initial hidden state of the decoder. This requires that the encoder and decoder RNNs have the same numbers of layers and hidden units.\n",
    "\n",
    "The forward calculation of the decoder is similar to the encoder's. The only difference is we add a dense layer with the hidden size to be the vocabulary size to output the predicted confidence score for each word."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "5"
    }
   },
   "outputs": [],
   "source": [
    "class Seq2SeqDecoder(d2l.Decoder):\n",
    "    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,\n",
    "                 dropout=0, **kwargs):\n",
    "        super(Seq2SeqDecoder, self).__init__(**kwargs)\n",
    "        self.embedding = nn.Embedding(vocab_size, embed_size)\n",
    "        self.rnn = rnn.LSTM(num_hiddens, num_layers, dropout=dropout)\n",
    "        self.dense = nn.Dense(vocab_size, flatten=False)\n",
    "\n",
    "    def init_state(self, enc_outputs, *args):\n",
    "        return enc_outputs[1]\n",
    "\n",
    "    def forward(self, X, state):\n",
    "        X = self.embedding(X).swapaxes(0, 1)\n",
    "        out, state = self.rnn(X, state)\n",
    "        # Make the batch to be the first dimension to simplify loss computation.\n",
    "        out = self.dense(out).swapaxes(0, 1)\n",
    "        return out, state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We create an decoder with the same hyper-parameters as the encoder. As can be seen, the output shape is changed to (batch size, the sequence length, vocabulary size)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "6"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((4, 7, 10), 2, (2, 4, 16), (2, 4, 16))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "decoder = Seq2SeqDecoder(vocab_size=10, embed_size=8,\n",
    "                         num_hiddens=16, num_layers=2)\n",
    "decoder.initialize()\n",
    "state = decoder.init_state(encoder(X))\n",
    "out, state = decoder(X, state)\n",
    "out.shape, len(state), state[0].shape, state[1].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Loss Function\n",
    "\n",
    "For each time step, the decoder outputs a vocabulary size confident score vector to predict words. Similar to language modeling, we can apply softmax to obtain the probabilities and then use cross entropy loss the calculate the loss. But note that we padded the target sentences to make them have the same length. We wound not like to compute the loss on the padding symbols.\n",
    "\n",
    "To implement the loss function that filters out some entries, we will use an operator called `SequenceMask`. It can specify to mask the first dimension (`axis=0`) or the second one (`axis=1`). If the second one is chosen, given a valid length vector `len` and 2-dim input `X`, this operator sets `X[i, len[i]:] = 0` for all $i$'s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "7"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       "[[1. 0. 0.]\n",
       " [4. 5. 0.]]\n",
       "<NDArray 2x3 @cpu(0)>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X = nd.array([[1,2,3], [4,5,6]])\n",
    "nd.SequenceMask(X, nd.array([1,2]), True, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Apply to $n$-dim tensor $X$, it sets `X[i, len[i]:, :, ..., :] = 0`. In addition, we can specify the filling value beyond 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "8"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       "[[[ 1.  1.  1.  1.]\n",
       "  [-1. -1. -1. -1.]\n",
       "  [-1. -1. -1. -1.]]\n",
       "\n",
       " [[ 1.  1.  1.  1.]\n",
       "  [ 1.  1.  1.  1.]\n",
       "  [-1. -1. -1. -1.]]]\n",
       "<NDArray 2x3x4 @cpu(0)>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X = nd.ones((2, 3, 4))\n",
    "nd.SequenceMask(X, nd.array([1,2]), True, value=-1, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can implement the masked version of the softmax cross-entropy loss. Note that each Gluon loss function allows to specify per-example weights, in default they are 1s. Then we can just use a zero weight for each example we would like to remove. So our customized loss function accepts an additional `valid_length` argument to ignore some failing elements in each sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "9"
    }
   },
   "outputs": [],
   "source": [
    "class MaskedSoftmaxCELoss(gloss.SoftmaxCELoss):\n",
    "    # pred shape: (batch_size, seq_len, vocab_size)\n",
    "    # label shape: (batch_size, seq_len)\n",
    "    # valid_length shape: (batch_size, )\n",
    "    def forward(self, pred, label, valid_length):\n",
    "        # the sample weights shape should be (batch_size, seq_len, 1)\n",
    "        weights = nd.ones_like(label).expand_dims(axis=-1)\n",
    "        weights = nd.SequenceMask(weights, valid_length, True, axis=1)\n",
    "        return super(MaskedSoftmaxCELoss, self).forward(pred, label, weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a sanity check, we create identical three sequences, keep 4 elements for the first sequence, 2 elements for the second sequence, and none for the last one. Then the first example loss should be 2 times larger than the second one, and the last loss should be 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "10"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       "[2.3025851 1.1512926 0.       ]\n",
       "<NDArray 3 @cpu(0)>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss = MaskedSoftmaxCELoss()\n",
    "loss(nd.ones((3, 4, 10)), nd.ones((3, 4)), nd.array([4, 2, 0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "During training, if the target sequence has length $n$, we feed the first $n-1$ tokens into the decoder as inputs, and the last $n-1$ tokens are used as ground truth label."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "13"
    }
   },
   "outputs": [],
   "source": [
    "def train_ch7(model, data_iter, lr, num_epochs, ctx):  # Saved in d2l\n",
    "    model.initialize(init.Xavier(), force_reinit=True, ctx=ctx)\n",
    "    trainer = gluon.Trainer(model.collect_params(),\n",
    "                            'adam', {'learning_rate': lr})\n",
    "    loss = MaskedSoftmaxCELoss()\n",
    "    tic = time.time()\n",
    "    for epoch in range(1, num_epochs+1):\n",
    "        l_sum, num_tokens_sum = 0.0, 0.0\n",
    "        for batch in data_iter:\n",
    "            X, X_vlen, Y, Y_vlen = [x.as_in_context(ctx) for x in batch]\n",
    "            Y_input, Y_label, Y_vlen = Y[:,:-1], Y[:,1:], Y_vlen-1\n",
    "            with autograd.record():\n",
    "                Y_hat, _ = model(X, Y_input, X_vlen, Y_vlen)\n",
    "                l = loss(Y_hat, Y_label, Y_vlen)\n",
    "            l.backward()\n",
    "            d2l.grad_clipping_gluon(model, 5, ctx)\n",
    "            num_tokens = Y_vlen.sum().asscalar()\n",
    "            trainer.step(num_tokens)\n",
    "            l_sum += l.sum().asscalar()\n",
    "            num_tokens_sum += num_tokens\n",
    "        if epoch % 50 == 0:\n",
    "            print(\"epoch %d, loss %.3f, time %.1f sec\" % (\n",
    "                epoch, l_sum/num_tokens_sum, time.time()-tic))\n",
    "            tic = time.time()\n",
    "            "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we create a model instance and set hyper-parameters. Then, we can train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "14"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 50, loss 0.126, time 10.2 sec\n",
      "epoch 100, loss 0.071, time 10.1 sec\n",
      "epoch 150, loss 0.044, time 10.1 sec\n",
      "epoch 200, loss 0.033, time 10.1 sec\n",
      "epoch 250, loss 0.028, time 10.2 sec\n",
      "epoch 300, loss 0.026, time 9.9 sec\n"
     ]
    }
   ],
   "source": [
    "embed_size, num_hiddens, num_layers, dropout = 32, 32, 2, 0.0\n",
    "batch_size, num_examples, max_len = 64, 1e3, 10\n",
    "lr, num_epochs, ctx = 0.005, 300, d2l.try_gpu()\n",
    "\n",
    "src_vocab, tgt_vocab, train_iter = d2l.load_data_nmt(\n",
    "    batch_size, max_len, num_examples)\n",
    "encoder = Seq2SeqEncoder(\n",
    "    len(src_vocab), embed_size, num_hiddens, num_layers, dropout)\n",
    "decoder = Seq2SeqDecoder(\n",
    "    len(tgt_vocab), embed_size, num_hiddens, num_layers, dropout)\n",
    "model = d2l.EncoderDecoder(encoder, decoder)\n",
    "train_ch7(model, train_iter, lr, num_epochs, ctx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predicting\n",
    "\n",
    "We introduced three methods to generate the output of the decoder at each time step in the [\"Beam Search\"](beam-search.md) section. Here we implement the simplest method, greedy search.\n",
    "\n",
    "\n",
    "![Sequence to sequence model predicting with greedy search](http://d2l.ai/_images/seq2seq_predict.svg)\n",
    "\n",
    "During predicting, we feed the same \"&lt;bos&gt;\" token to the decoder as training at time step 0. But the input token for a later time step is the predicted token from the previous time step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "15"
    }
   },
   "outputs": [],
   "source": [
    "def translate_ch7(model, src_sentence, src_vocab, tgt_vocab, max_len, ctx):\n",
    "    src_tokens = src_vocab[src_sentence.lower().split(' ')]\n",
    "    src_len = len(src_tokens)\n",
    "    if src_len < max_len:\n",
    "        src_tokens += [src_vocab.pad] * (max_len - src_len)\n",
    "    enc_X = nd.array(src_tokens, ctx=ctx)\n",
    "    enc_valid_length = nd.array([src_len], ctx=ctx)\n",
    "    # use expand_dim to add the batch_size dimension.\n",
    "    enc_outputs = model.encoder(enc_X.expand_dims(axis=0), enc_valid_length)\n",
    "    dec_state = model.decoder.init_state(enc_outputs, enc_valid_length)\n",
    "    dec_X = nd.array([tgt_vocab.bos], ctx=ctx).expand_dims(axis=0)\n",
    "    predict_tokens = []\n",
    "    for _ in range(max_len):\n",
    "        Y, dec_state = model.decoder(dec_X, dec_state)\n",
    "        # The token with highest score is used as the next time step input. \n",
    "        dec_X = Y.argmax(axis=2)\n",
    "        py = dec_X.squeeze(axis=0).astype('int32').asscalar()\n",
    "        if py == tgt_vocab.eos:\n",
    "            break\n",
    "        predict_tokens.append(py)\n",
    "    return ' '.join(tgt_vocab.to_tokens(predict_tokens))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try several examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "16"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Go . => va !\n",
      "Wow ! => <unk> !\n",
      "I'm OK . => je vais bien .\n",
      "I won ! => je l'ai emporté !\n"
     ]
    }
   ],
   "source": [
    "for sentence in ['Go .', 'Wow !', \"I'm OK .\", 'I won !']:\n",
    "    print(sentence + ' => ' + translate_ch7(\n",
    "        model, sentence, src_vocab, tgt_vocab, max_len, ctx))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.7.1"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}