RNN_LSTM_from_scratch.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "5.1-EXE-Recurrent-Neural-Networks.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "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.4"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y-CptVs7iACc"
      },
      "source": [
        "# How to build RNNs and LSTMs from scratch with NumPy\n",
        "\n",
        "Originally developed by Nicklas Hansen, Peter E. Christensen and Alexander R. Johansen as educational material for the graduate deep learning course at the Technical University of Denmark (DTU). You can access the full course material [here](https://github.com/DeepLearningDTU/02456-deep-learning-with-PyTorch). Inspired by the great [Andrej Karpathy](https://karpathy.ai/).\n",
        "\n",
        "____\n",
        "\n",
        "\n",
        "In this lab, we will introduce different ways of learning from sequential data.\n",
        "\n",
        "As a recurring example, we will train neural networks to do language modelling, i.e. predict the next token in a sentence. In the context of natural language processing a token could be a character or a word, but mind you that the concepts introduced here apply to all kinds of sequential data, such as e.g. protein sequences, weather measurements, audio signals, or videos, just to name a few.\n",
        "\n",
        "To really get a grasp of what is going on inside a recurrent neural network (RNN), we will carry out a substantial part of this exercise in NumPy rather than PyTorch. We start off with a simple toy problem, build an RNN using NumPy, train and it, and see for ourselves that it really works. Once we're convinced, we proceed to build and train a Long Short-Term Memory (LSTM) cell, also in NumPy. This is *not* simply to cause you frustration, but rather to provide you with a deeper understanding of the recurrence in RNNs, which will become very beneficial to you in the following weeks. Once you understand the inner workings of an RNN, we will proceed to a PyTorch implementation that you may use for the remainder of the course and in your projects.\n",
        "\n",
        "To summarize, in this notebook we will show you:\n",
        "* How to represent sequences of categorical variables\n",
        "* How to build and train an RNN in NumPy\n",
        "* How to build and train an LSTM network in NumPy\n",
        "* How to build and train an LSTM network in PyTorch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XapO8SLwiACd"
      },
      "source": [
        "## Representing tokens or text\n",
        "\n",
        "In previous labs we mainly considered data $x \\in \\mathrm{R}^d$, where $d$ is the feature space dimension.\n",
        "With time sequences our data can be represented as $x \\in \\mathrm{R}^{t \\, \\times \\, d}$, where $t$ is the sequence length. \n",
        "This emphasises sequence dependence and that the samples along the sequence are not independent and identically distributed (i.i.d.).\n",
        "We will model functions as $\\mathrm{R}^{t \\, \\times \\, d} \\rightarrow \\mathrm{R}^c$, where $c$ is the amount of classes in the output.\n",
        "\n",
        "There are several ways to represent sequences. With text, the challenge is how to represent a word as a feature vector in $d$ dimensions, as we are required to represent text with decimal numbers in order to apply neural networks to it.\n",
        "\n",
        "In this exercise we will use a simple one-hot encoding but for categorical variables that can take on many values (e.g. words in the English language) this may be infeasible. For such scenarios, you can project the encodings into a smaller space by use of embeddings. If you want to learn more about tokens, encodings and embeddings than what is covered in this exercise, we highly recommend [this lecture](https://www.youtube.com/watch?v=kEMJRjEdNzM&list=PLoROMvodv4rOhcuXMZkNm7j3fVwBBY42z)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bdA4LPsFiACe"
      },
      "source": [
        "### One-hot encoding over vocabulary\n",
        "\n",
        "One way to represent a fixed amount of words is by making a one-hot encoded vector, which consists of 0s in all cells with the exception of a single 1 in a cell used uniquely to identify each word.\n",
        "\n",
        "| vocabulary    | one-hot encoded vector   |\n",
        "| ------------- |--------------------------|\n",
        "| Paris         | $= [1, 0, 0, \\ldots, 0]$ |\n",
        "| Rome          | $= [0, 1, 0, \\ldots, 0]$ |\n",
        "| Copenhagen    | $= [0, 0, 1, \\ldots, 0]$ |\n",
        "\n",
        "Representing a large vocabulary with one-hot encodings often becomes inefficient because of the size of each sparse vector.\n",
        "To overcome this challenge it is common practice to truncate the vocabulary to contain the $k$ most used words and represent the rest with a special symbol, $\\mathtt{UNK}$, to define unknown/unimportant words.\n",
        "This often causes entities such as names to be represented with $\\mathtt{UNK}$ because they are rare.\n",
        "\n",
        "Consider the following text\n",
        "> I love the corny jokes in Spielberg's new movie.\n",
        "\n",
        "where an example result would be similar to\n",
        "> I love the corny jokes in $\\mathtt{UNK}$'s new movie."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KNmyPw7zk2gY"
      },
      "source": [
        "## Generating a dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M9IEA4t2k2gb"
      },
      "source": [
        "For this exercise we will create a simple dataset that we can learn from. We generate sequences of the form:\n",
        "\n",
        "`a b EOS`,\n",
        "\n",
        "`a a b b EOS`,\n",
        "\n",
        "`a a a a a b b b b b EOS`\n",
        "\n",
        "where `EOS` is a special character denoting the end of a sequence. The task is to predict the next token $t_n$, i.e. `a`, `b`, `EOS` or the unknown token `UNK` given a sequence of tokens $\\{ t_{1}, t_{2}, \\dots , t_{n-1}\\}$, and we are to process sequences in a sequential manner. As such, the network will need to learn that e.g. 5 `b`s and an `EOS` token will follow 5 `a`s."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dcoN-kb7k2gc"
      },
      "source": [
        "import numpy as np\n",
        "\n",
        "# Set seed such that we always get the same dataset\n",
        "# (this is a good idea in general)\n",
        "np.random.seed(42)\n",
        "\n",
        "def generate_dataset(num_sequences=2**8):\n",
        "    \"\"\"\n",
        "    Generates a number of sequences as our dataset.\n",
        "    \n",
        "    Args:\n",
        "     `num_sequences`: the number of sequences to be generated.\n",
        "     \n",
        "    Returns a list of sequences.\n",
        "    \"\"\"\n",
        "    samples = []\n",
        "    \n",
        "    for _ in range(num_sequences): \n",
        "        num_tokens = np.random.randint(1, 12)\n",
        "        sample = ['a'] * num_tokens + ['b'] * num_tokens + ['EOS']\n",
        "        samples.append(sample)\n",
        "        \n",
        "    return samples\n",
        "\n",
        "\n",
        "sequences = generate_dataset()\n",
        "\n",
        "print('A single sample from the generated dataset:')\n",
        "print(sequences[0])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YMLd3Gzak2gp"
      },
      "source": [
        "## Representing tokens as indices"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S9LSqaJSk2gp"
      },
      "source": [
        "To build a one-hot encoding, we need to assign each possible word in our vocabulary an index. We do that by creating two dictionaries: one that allows us to go from a given word to its corresponding index in our vocabulary, and one for the reverse direction. Let's call them `word_to_idx` and `idx_to_word`. The keyword `vocab_size` specifies the maximum size of our vocabulary. If we try to access a word that does not exist in our vocabulary, it is automatically replaced by the `UNK` token or its corresponding index."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Smdo70UMk2gr"
      },
      "source": [
        "from collections import defaultdict\n",
        "\n",
        "def sequences_to_dicts(sequences):\n",
        "    \"\"\"\n",
        "    Creates word_to_idx and idx_to_word dictionaries for a list of sequences.\n",
        "    \"\"\"\n",
        "    # A bit of Python-magic to flatten a nested list\n",
        "    flatten = lambda l: [item for sublist in l for item in sublist]\n",
        "    \n",
        "    # Flatten the dataset\n",
        "    all_words = flatten(sequences)\n",
        "    \n",
        "    # Count number of word occurences\n",
        "    word_count = defaultdict(int)\n",
        "    for word in flatten(sequences):\n",
        "        word_count[word] += 1\n",
        "\n",
        "    # Sort by frequency\n",
        "    word_count = sorted(list(word_count.items()), key=lambda l: -l[1])\n",
        "\n",
        "    # Create a list of all unique words\n",
        "    unique_words = [item[0] for item in word_count]\n",
        "    \n",
        "    # Add UNK token to list of words\n",
        "    unique_words.append('UNK')\n",
        "\n",
        "    # Count number of sequences and number of unique words\n",
        "    num_sentences, vocab_size = len(sequences), len(unique_words)\n",
        "\n",
        "    # Create dictionaries so that we can go from word to index and back\n",
        "    # If a word is not in our vocabulary, we assign it to token 'UNK'\n",
        "    word_to_idx = defaultdict(lambda: vocab_size-1)\n",
        "    idx_to_word = defaultdict(lambda: 'UNK')\n",
        "\n",
        "    # Fill dictionaries\n",
        "    for idx, word in enumerate(unique_words):\n",
        "        # YOUR CODE HERE!\n",
        "        word_to_idx[word] = \n",
        "        idx_to_word[idx] = \n",
        "\n",
        "    return word_to_idx, idx_to_word, num_sentences, vocab_size\n",
        "\n",
        "\n",
        "word_to_idx, idx_to_word, num_sequences, vocab_size = sequences_to_dicts(sequences)\n",
        "\n",
        "print(f'We have {num_sequences} sentences and {len(word_to_idx)} unique tokens in our dataset (including UNK).\\n')\n",
        "print('The index of \\'b\\' is', word_to_idx['b'])\n",
        "print(f'The word corresponding to index 1 is \\'{idx_to_word[1]}\\'')\n",
        "\n",
        "assert idx_to_word[word_to_idx['b']] == 'b', \\\n",
        "    'Consistency error: something went wrong in the conversion.'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sNY1OOS_k2gy"
      },
      "source": [
        "## Exercise A:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xc-iHbBmk2gz"
      },
      "source": [
        "Complete the `sequences_to_dicts` function above. You will need to fill the `word_to_idx` and `idx_to_word` dictionaries so that we can go back and forth between the two representations."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cGSoDRgHk2g1"
      },
      "source": [
        "## Partitioning the dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UMTn1iLIk2g1"
      },
      "source": [
        "To build our dataset, we need to create inputs and targets for each sequences and partition sentences it into training, validation and test sets. 80%, 10% and 10% is a common distribution, but mind you that this largely depends on the size of the dataset. Since we are doing next-word predictions, our target sequence is simply the input sequence shifted by one word.\n",
        "\n",
        "We can use PyTorch's `Dataset` class to build a simple dataset where we can easily retrieve (inputs, targets) pairs for each of our sequences."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9dW7MrPnk2g3"
      },
      "source": [
        "from torch.utils import data\n",
        "\n",
        "class Dataset(data.Dataset):\n",
        "    def __init__(self, inputs, targets):\n",
        "        self.inputs = inputs\n",
        "        self.targets = targets\n",
        "\n",
        "    def __len__(self):\n",
        "        # Return the size of the dataset\n",
        "        return len(self.targets)\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        # Retrieve inputs and targets at the given index\n",
        "        X = self.inputs[index]\n",
        "        y = self.targets[index]\n",
        "\n",
        "        return X, y\n",
        "\n",
        "    \n",
        "def create_datasets(sequences, dataset_class, p_train=0.8, p_val=0.1, p_test=0.1):\n",
        "    # Define partition sizes\n",
        "    num_train = int(len(sequences)*p_train)\n",
        "    num_val = int(len(sequences)*p_val)\n",
        "    num_test = int(len(sequences)*p_test)\n",
        "\n",
        "    # Split sequences into partitions\n",
        "    sequences_train = sequences[:num_train]\n",
        "    sequences_val = sequences[num_train:num_train+num_val]\n",
        "    sequences_test = sequences[-num_test:]\n",
        "\n",
        "    def get_inputs_targets_from_sequences(sequences):\n",
        "        # Define empty lists\n",
        "        inputs, targets = [], []\n",
        "        \n",
        "        # Append inputs and targets s.t. both lists contain L-1 words of a sentence of length L\n",
        "        # but targets are shifted right by one so that we can predict the next word\n",
        "        for sequence in sequences:\n",
        "            inputs.append(sequence[:-1])\n",
        "            targets.append(sequence[1:])\n",
        "            \n",
        "        return inputs, targets\n",
        "\n",
        "    # Get inputs and targets for each partition\n",
        "    inputs_train, targets_train = get_inputs_targets_from_sequences(sequences_train)\n",
        "    inputs_val, targets_val = get_inputs_targets_from_sequences(sequences_val)\n",
        "    inputs_test, targets_test = get_inputs_targets_from_sequences(sequences_test)\n",
        "\n",
        "    # Create datasets\n",
        "    training_set = dataset_class(inputs_train, targets_train)\n",
        "    validation_set = dataset_class(inputs_val, targets_val)\n",
        "    test_set = dataset_class(inputs_test, targets_test)\n",
        "\n",
        "    return training_set, validation_set, test_set\n",
        "    \n",
        "\n",
        "training_set, validation_set, test_set = create_datasets(sequences, Dataset)\n",
        "\n",
        "print(f'We have {len(training_set)} samples in the training set.')\n",
        "print(f'We have {len(validation_set)} samples in the validation set.')\n",
        "print(f'We have {len(test_set)} samples in the test set.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4xMMSm7Mk2g9"
      },
      "source": [
        "When working with more complex data than what we use in this exercise, creating a PyTorch `DataLoader` on top of the dataset can be beneficial. A data loader is basically a fancy generator/iterator that we can use to abstract away all of the data handling and pre-processing + it's super useful for processing batches of data as well! Data loaders will come in handy later when you start to work on your projects, so be sure to check them out!\n",
        "\n",
        "For more information on how to use datasets and data loaders in PyTorch, [consult the official guide](https://pytorch.org/tutorials/beginner/data_loading_tutorial.html)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dzmryk72k2g-"
      },
      "source": [
        "## One-hot encodings"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "abRN9f8Xk2g_"
      },
      "source": [
        "We now create a simple function that returns the one-hot encoded representation of a given index of a word in our vocabulary. Notice that the shape of the one-hot encoding is equal to the entire vocabulary (which can be huge!). Additionally, we define a function to automatically one-hot encode a sentence."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IZruCIHJk2hB"
      },
      "source": [
        "def one_hot_encode(idx, vocab_size):\n",
        "    \"\"\"\n",
        "    One-hot encodes a single word given its index and the size of the vocabulary.\n",
        "    \n",
        "    Args:\n",
        "     `idx`: the index of the given word\n",
        "     `vocab_size`: the size of the vocabulary\n",
        "    \n",
        "    Returns a 1-D numpy array of length `vocab_size`.\n",
        "    \"\"\"\n",
        "    # Initialize the encoded array\n",
        "    one_hot = np.zeros(vocab_size)\n",
        "    \n",
        "    # Set the appropriate element to one\n",
        "    one_hot[idx] = 1.0\n",
        "\n",
        "    return one_hot\n",
        "\n",
        "\n",
        "def one_hot_encode_sequence(sequence, vocab_size):\n",
        "    \"\"\"\n",
        "    One-hot encodes a sequence of words given a fixed vocabulary size.\n",
        "    \n",
        "    Args:\n",
        "     `sentence`: a list of words to encode\n",
        "     `vocab_size`: the size of the vocabulary\n",
        "     \n",
        "    Returns a 3-D numpy array of shape (num words, vocab size, 1).\n",
        "    \"\"\"\n",
        "    # Encode each word in the sentence\n",
        "    encoding = np.array([one_hot_encode(word_to_idx[word], vocab_size) for word in sequence])\n",
        "\n",
        "    # Reshape encoding s.t. it has shape (num words, vocab size, 1)\n",
        "    encoding = encoding.reshape(encoding.shape[0], encoding.shape[1], 1)\n",
        "    \n",
        "    return encoding\n",
        "\n",
        "\n",
        "test_word = one_hot_encode(word_to_idx['a'], vocab_size)\n",
        "print(f'Our one-hot encoding of \\'a\\' has shape {test_word.shape}.')\n",
        "\n",
        "test_sentence = one_hot_encode_sequence(['a', 'b'], vocab_size)\n",
        "print(f'Our one-hot encoding of \\'a b\\' has shape {test_sentence.shape}.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "erI_MXvKk2hG"
      },
      "source": [
        "Great! Now that we have our one-hot encodings in place, we can move on to the RNNs!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MA6bxjGWjeSB"
      },
      "source": [
        "# Introduction to Recurrent Neural Networks (RNN)\n",
        "\n",
        "Reading material: [blog post](http://karpathy.github.io/2015/05/21/rnn-effectiveness/) and (optionally) [this lecture](https://www.youtube.com/watch?v=iWea12EAu6U&list=PLoROMvodv4rOhcuXMZkNm7j3fVwBBY42z).\n",
        "\n",
        "___\n",
        "\n",
        "A recurrent neural network (RNN) is a type of neural network that has been succesful in modelling sequential data, e.g. language, speech, protein sequences, etc.\n",
        "\n",
        "A RNN performs its computations in a cyclic manner, where the same computation is applied to every sample of a given sequence.\n",
        "The idea is that the network should be able to use the previous computations as some form of memory and apply this to future computations.\n",
        "An image may best explain how this is to be understood,\n",
        "\n",
        "![rnn-unroll image](https://github.com/DeepLearningDTU/02456-deep-learning-with-PyTorch/blob/master/static_files/rnn-unfold.png?raw=1)\n",
        "\n",
        "\n",
        "where it the network contains the following elements:\n",
        "\n",
        "- $x$ is the input sequence of samples, \n",
        "- $U$ is a weight matrix applied to the given input sample,\n",
        "- $V$ is a weight matrix used for the recurrent computation in order to pass memory along the sequence,\n",
        "- $W$ is a weight matrix used to compute the output of the every timestep (given that every timestep requires an output),\n",
        "- $h$ is the hidden state (the network's memory) for a given time step, and\n",
        "- $o$ is the resulting output.\n",
        "\n",
        "When the network is unrolled as shown, it is easier to refer to a timestep, $t$.\n",
        "We have the following computations through the network:\n",
        "\n",
        "- $h_t = f(U\\,{x_t} + V\\,{h_{t-1}})$, where $f$ is a non-linear activation function, e.g. $\\mathrm{tanh}$.\n",
        "- $o_t = W\\,{h_t}$\n",
        "\n",
        "When we are doing language modelling using a cross-entropy loss, we additionally apply the softmax function to the output $o_{t}$:\n",
        "\n",
        "- $\\hat{y}_t = \\mathrm{softmax}(o_{t})$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GuvwbvsGz9KE"
      },
      "source": [
        "## Implementing an RNN"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "20WtL0hak2hJ"
      },
      "source": [
        "We will implement the forward pass, backward pass, optimization and training loop for an RNN in numpy so that you can get familiar with the recurrent nature of RNNs. Later, we will go back to PyTorch and appreciate how convenient the implementation becomes!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oBYpnqINk2hK"
      },
      "source": [
        "Let's first define the necessary model parameters. Recall that an $n \\times m$ weight matrix maps $\\mathbb{R}^{m} \\rightarrow \\mathbb{R}^{n}$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m6Q8gTJ0k2hK"
      },
      "source": [
        "hidden_size = 50 # Number of dimensions in the hidden state\n",
        "vocab_size  = len(word_to_idx) # Size of the vocabulary used\n",
        "\n",
        "def init_orthogonal(param):\n",
        "    \"\"\"\n",
        "    Initializes weight parameters orthogonally.\n",
        "    This is a common initiailization for recurrent neural networks.\n",
        "    \n",
        "    Refer to this paper for an explanation of this initialization:\n",
        "    https://arxiv.org/abs/1312.6120\n",
        "    \"\"\"\n",
        "    if param.ndim < 2:\n",
        "        raise ValueError(\"Only parameters with 2 or more dimensions are supported.\")\n",
        "\n",
        "    rows, cols = param.shape\n",
        "    \n",
        "    new_param = np.random.randn(rows, cols)\n",
        "    \n",
        "    if rows < cols:\n",
        "        new_param = new_param.T\n",
        "    \n",
        "    # Compute QR factorization\n",
        "    q, r = np.linalg.qr(new_param)\n",
        "    \n",
        "    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf\n",
        "    d = np.diag(r, 0)\n",
        "    ph = np.sign(d)\n",
        "    q *= ph\n",
        "\n",
        "    if rows < cols:\n",
        "        q = q.T\n",
        "    \n",
        "    new_param = q\n",
        "    \n",
        "    return new_param\n",
        "\n",
        "\n",
        "def init_rnn(hidden_size, vocab_size):\n",
        "    \"\"\"\n",
        "    Initializes our recurrent neural network.\n",
        "    \n",
        "    Args:\n",
        "     `hidden_size`: the dimensions of the hidden state\n",
        "     `vocab_size`: the dimensions of our vocabulary\n",
        "    \"\"\"\n",
        "    # Weight matrix (input to hidden state)\n",
        "    # YOUR CODE HERE!\n",
        "    U = np.zeros((, ))\n",
        "\n",
        "    # Weight matrix (recurrent computation)\n",
        "    # YOUR CODE HERE!\n",
        "    V = np.zeros((, ))\n",
        "\n",
        "    # Weight matrix (hidden state to output)\n",
        "    # YOUR CODE HERE!\n",
        "    W = np.zeros((, ))\n",
        "\n",
        "    # Bias (hidden state)\n",
        "    # YOUR CODE HERE!\n",
        "    b_hidden = np.zeros((, ))\n",
        "\n",
        "    # Bias (output)\n",
        "    # YOUR CODE HERE!\n",
        "    b_out = np.zeros((, ))\n",
        "    \n",
        "    # Initialize weights\n",
        "    U = init_orthogonal(U)\n",
        "    V = init_orthogonal(V)\n",
        "    W = init_orthogonal(W)\n",
        "    \n",
        "    # Return parameters as a tuple\n",
        "    return U, V, W, b_hidden, b_out\n",
        "\n",
        "\n",
        "params = init_rnn(hidden_size=hidden_size, vocab_size=vocab_size)\n",
        "print('U:', params[0].shape)\n",
        "print('V:', params[1].shape)\n",
        "print('W:', params[2].shape)\n",
        "print('b_hidden:', params[3].shape)\n",
        "print('b_out:', params[4].shape)\n",
        "\n",
        "for param in params:\n",
        "    assert param.ndim == 2, \\\n",
        "        'all parameters should be 2-dimensional '\\\n",
        "        '(hint: a dimension can simply have size 1)'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yolo5dKrk2hR"
      },
      "source": [
        "## Exercise B:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SIyQO1S3k2hS"
      },
      "source": [
        "Define the appropriate shape of the weights and biases in the `init_rnn` function.\n",
        "\n",
        "You only have to worry about the weight and bias dimensions; you can leave the scaling as is. Refer to the equations and figure above if you're in doubt. Note that `np.zeros()` expects a tuple as its first argument (shape)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KFH6jRXik2hV"
      },
      "source": [
        "### Function definitions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nRtByep8k2hW"
      },
      "source": [
        "Below we have defined the `sigmoid`, `tanh` and `softmax` functions for you. You might need them in a bit!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rxKigMDdk2hY"
      },
      "source": [
        "#### Sigmoid activation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qe2sVT0Hk2hY"
      },
      "source": [
        "$\\sigma(x) = \\frac{1}{1 + e^{-x}}$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sKw9FNa2k2hZ"
      },
      "source": [
        "def sigmoid(x, derivative=False):\n",
        "    \"\"\"\n",
        "    Computes the element-wise sigmoid activation function for an array x.\n",
        "\n",
        "    Args:\n",
        "     `x`: the array where the function is applied\n",
        "     `derivative`: if set to True will return the derivative instead of the forward pass\n",
        "    \"\"\"\n",
        "    x_safe = x + 1e-12\n",
        "    f = 1 / (1 + np.exp(-x_safe))\n",
        "    \n",
        "    if derivative: # Return the derivative of the function evaluated at x\n",
        "        return f * (1 - f)\n",
        "    else: # Return the forward pass of the function at x\n",
        "        return f"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QgGTIbNpk2hd"
      },
      "source": [
        "#### Hyperbolic Tangent activation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s9-qTul2k2he"
      },
      "source": [
        "$\\tanh{x} = \\frac{e^{x} - e^{-x}}{e^{x} + e^{-x}}$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NJPhICbKk2hf"
      },
      "source": [
        "def tanh(x, derivative=False):\n",
        "    \"\"\"\n",
        "    Computes the element-wise tanh activation function for an array x.\n",
        "\n",
        "    Args:\n",
        "     `x`: the array where the function is applied\n",
        "     `derivative`: if set to True will return the derivative instead of the forward pass\n",
        "    \"\"\"\n",
        "    x_safe = x + 1e-12\n",
        "    f = (np.exp(x_safe)-np.exp(-x_safe))/(np.exp(x_safe)+np.exp(-x_safe))\n",
        "    \n",
        "    if derivative: # Return the derivative of the function evaluated at x\n",
        "        return 1-f**2\n",
        "    else: # Return the forward pass of the function at x\n",
        "        return f"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OR4wltaGk2hk"
      },
      "source": [
        "#### Softmax"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uDVaozqLk2hl"
      },
      "source": [
        "$\\mathrm{softmax}(x) = \\frac{e^{x_{i}}}{\\sum_{j=1}^{K} e^{x_{j}}} \\text { for } i=1, \\ldots, K \\text { and } \\mathbf{x}=\\left(x_{1}, \\ldots, x_{K}\\right) \\in \\mathbb{R}^{K}$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c-3VX41qk2hn"
      },
      "source": [
        "def softmax(x, derivative=False):\n",
        "    \"\"\"\n",
        "    Computes the softmax for an array x.\n",
        "    \n",
        "    Args:\n",
        "     `x`: the array where the function is applied\n",
        "     `derivative`: if set to True will return the derivative instead of the forward pass\n",
        "    \"\"\"\n",
        "    x_safe = x + 1e-12\n",
        "    f = np.exp(x_safe) / np.sum(np.exp(x_safe))\n",
        "    \n",
        "    if derivative: # Return the derivative of the function evaluated at x\n",
        "        pass # We will not need this one\n",
        "    else: # Return the forward pass of the function at x\n",
        "        return f"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WFUtyOtesrfe"
      },
      "source": [
        "### Implement the forward pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xfRLRqrWk2hs"
      },
      "source": [
        "Now that we have all the definitions in place, we can start to implement a forward pass."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RNZ3IVjs4jFB"
      },
      "source": [
        "def forward_pass(inputs, hidden_state, params):\n",
        "    \"\"\"\n",
        "    Computes the forward pass of a vanilla RNN.\n",
        "    \n",
        "    Args:\n",
        "     `inputs`: sequence of inputs to be processed\n",
        "     `hidden_state`: an already initialized hidden state\n",
        "     `params`: the parameters of the RNN\n",
        "    \"\"\"\n",
        "    # First we unpack our parameters\n",
        "    U, V, W, b_hidden, b_out = params\n",
        "    \n",
        "    # Create a list to store outputs and hidden states\n",
        "    outputs, hidden_states = [], []\n",
        "    \n",
        "    # For each element in input sequence\n",
        "    for t in range(len(inputs)):\n",
        "\n",
        "        # Compute new hidden state\n",
        "        # YOUR CODE HERE!\n",
        "        hidden_state = \n",
        "\n",
        "        # Compute output\n",
        "        # YOUR CODE HERE!\n",
        "        out = \n",
        "        \n",
        "        # Save results and continue\n",
        "        outputs.append(out)\n",
        "        hidden_states.append(hidden_state.copy())\n",
        "    \n",
        "    return outputs, hidden_states\n",
        "\n",
        "\n",
        "# Get first sequence in training set\n",
        "test_input_sequence, test_target_sequence = training_set[0]\n",
        "\n",
        "# One-hot encode input and target sequence\n",
        "test_input = one_hot_encode_sequence(test_input_sequence, vocab_size)\n",
        "test_target = one_hot_encode_sequence(test_target_sequence, vocab_size)\n",
        "\n",
        "# Initialize hidden state as zeros\n",
        "hidden_state = np.zeros((hidden_size, 1))\n",
        "\n",
        "# Now let's try out our new function\n",
        "outputs, hidden_states = forward_pass(test_input, hidden_state, params)\n",
        "\n",
        "print('Input sequence:')\n",
        "print(test_input_sequence)\n",
        "\n",
        "print('\\nTarget sequence:')\n",
        "print(test_target_sequence)\n",
        "\n",
        "print('\\nPredicted sequence:')\n",
        "print([idx_to_word[np.argmax(output)] for output in outputs])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WvE8zu8Ck2hx"
      },
      "source": [
        "## Exercise C:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aDk85eRTk2hx"
      },
      "source": [
        "Implement the forward pass in the code above. Use $\\tanh{x}$ as the non-linear activation function $f$. You can use `np.dot()` to compute dot products. Refer to the equations and the figure if you're in doubt."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vug6QCohsx_S"
      },
      "source": [
        "### Implement the backward pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M4unVioMk2h0"
      },
      "source": [
        "It's now time to implement the backward pass. This can be a bit tricky so it may be helpful to take another look at the RNN equations, figure and your forward pass implementation. Note that (depending on your implementation) you may sometimes need to transpose using `np.transpose()` or simply `.T`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "roXK0HW6s2Z_"
      },
      "source": [
        "def clip_gradient_norm(grads, max_norm=0.25):\n",
        "    \"\"\"\n",
        "    Clips gradients to have a maximum norm of `max_norm`.\n",
        "    This is to prevent the exploding gradients problem.\n",
        "    \"\"\" \n",
        "    # Set the maximum of the norm to be of type float\n",
        "    max_norm = float(max_norm)\n",
        "    total_norm = 0\n",
        "    \n",
        "    # Calculate the L2 norm squared for each gradient and add them to the total norm\n",
        "    for grad in grads:\n",
        "        grad_norm = np.sum(np.power(grad, 2))\n",
        "        total_norm += grad_norm\n",
        "    \n",
        "    total_norm = np.sqrt(total_norm)\n",
        "    \n",
        "    # Calculate clipping coeficient\n",
        "    clip_coef = max_norm / (total_norm + 1e-6)\n",
        "    \n",
        "    # If the total norm is larger than the maximum allowable norm, then clip the gradient\n",
        "    if clip_coef < 1:\n",
        "        for grad in grads:\n",
        "            grad *= clip_coef\n",
        "    \n",
        "    return grads\n",
        "\n",
        "\n",
        "def backward_pass(inputs, outputs, hidden_states, targets, params):\n",
        "    \"\"\"\n",
        "    Computes the backward pass of a vanilla RNN.\n",
        "    \n",
        "    Args:\n",
        "     `inputs`: sequence of inputs to be processed\n",
        "     `outputs`: sequence of outputs from the forward pass\n",
        "     `hidden_states`: sequence of hidden_states from the forward pass\n",
        "     `targets`: sequence of targets\n",
        "     `params`: the parameters of the RNN\n",
        "    \"\"\"\n",
        "    # First we unpack our parameters\n",
        "    U, V, W, b_hidden, b_out = params\n",
        "    \n",
        "    # Initialize gradients as zero\n",
        "    d_U, d_V, d_W = np.zeros_like(U), np.zeros_like(V), np.zeros_like(W)\n",
        "    d_b_hidden, d_b_out = np.zeros_like(b_hidden), np.zeros_like(b_out)\n",
        "    \n",
        "    # Keep track of hidden state derivative and loss\n",
        "    d_h_next = np.zeros_like(hidden_states[0])\n",
        "    loss = 0\n",
        "    \n",
        "    # For each element in output sequence\n",
        "    # NB: We iterate backwards s.t. t = N, N-1, ... 1, 0\n",
        "    for t in reversed(range(len(outputs))):\n",
        "\n",
        "        # Compute cross-entropy loss (as a scalar)\n",
        "        # When taking logarithms, it's a good idea to add a small constant (e.g. 1e-9)\n",
        "        # YOUR CODE HERE!\n",
        "        loss += \n",
        "        \n",
        "        # Backpropagate into output (derivative of cross-entropy)\n",
        "        # If you're confused about this step, see this link for an explanation:\n",
        "        # http://cs231n.github.io/neural-networks-case-study/#grad\n",
        "        d_o = outputs[t].copy()\n",
        "        d_o[np.argmax(targets[t])] -= 1\n",
        "        \n",
        "        # Backpropagate into W\n",
        "        # YOUR CODE HERE!\n",
        "        d_W += \n",
        "        d_b_out += d_o\n",
        "        \n",
        "        # Backpropagate into h\n",
        "        d_h = np.dot(W.T, d_o) + d_h_next\n",
        "        \n",
        "        # Backpropagate through non-linearity\n",
        "        # (we assume tanh is used here)\n",
        "        d_f = (1 - hidden_states[t]**2) * d_h\n",
        "        d_b_hidden += d_f\n",
        "        \n",
        "        # Backpropagate into U\n",
        "        # YOUR CODE HERE!\n",
        "        d_U += \n",
        "        \n",
        "        # Backpropagate into V\n",
        "        # YOUR CODE HERE!\n",
        "        d_V += \n",
        "        d_h_next = np.dot(V.T, d_f)\n",
        "    \n",
        "    # Pack gradients\n",
        "    grads = d_U, d_V, d_W, d_b_hidden, d_b_out    \n",
        "    \n",
        "    # Clip gradients\n",
        "    grads = clip_gradient_norm(grads)\n",
        "    \n",
        "    return loss, grads\n",
        "\n",
        "\n",
        "loss, grads = backward_pass(test_input, outputs, hidden_states, test_target, params)\n",
        "\n",
        "print('We get a loss of:')\n",
        "print(loss)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4OguQafQk2h5"
      },
      "source": [
        "## Exercise D:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ezSRiVJzk2h5"
      },
      "source": [
        "Implement the missing code in the backward pass above using a cross-entropy loss and $\\tanh{x}$ as non-linear activation function $f$. Refer to previous weeks for the definition of cross-entropy, see week 1 exercise i) for an explanation of its derivative, or follow [this link](http://cs231n.github.io/neural-networks-case-study/#grad) for a hint.\n",
        "\n",
        "**Note:** to compute a cross-entropy loss that is independent of the number of target classes, it is common to divide the loss by the number of targets (i.e. take the mean rather than a sum). It is optional whether you do this or not.\n",
        "\n",
        "**Hint:**\n",
        "To complete the above backward pass implementation, we need to compute the partial derivatives\n",
        "$\n",
        "\\frac{\\partial E}{\\partial W},~\\frac{\\partial E}{\\partial U},~\\frac{\\partial E}{\\partial V}\n",
        "$. Recall though, that RNNs are recurrent and the weights $W,~U,~V$ are shared across time, i.e. we do not have separate weights for each time step. Therefore, to compute e.g. the partial derivative $\\frac{\\partial E}{\\partial W}$, we need to 1) sum up across time, and 2) apply the chain rule:\n",
        "\n",
        "$$\\frac{\\partial E}{\\partial W} = \\sum_{t} \\frac{\\partial E}{\\partial W_{t}} =  \\sum_{t} \\frac{\\partial E}{\\partial o_{t}} \\frac{\\partial o_{t}}{\\partial W_{t}}\\,,$$\n",
        "where $o_{t}$ is the resulting output (logits) at time step $t$, as previously covered in week 1 and 2.\n",
        "\n",
        "Now, recall the definition of the RNN forward pass. We have the following computations through the network:\n",
        "\n",
        "- $h_t = f(U\\,{x_t} + V\\,{h_{t-1}})$, where $f$ usually is an activation function, e.g. $\\mathrm{tanh}$.\n",
        "- $o_t = W\\,{h_t}$\n",
        "- $\\hat{y}_t = \\mathrm{softmax}(o_{t})$\n",
        "\n",
        "where\n",
        "- $U$ is a weight matrix applied to the given input sample,\n",
        "- $V$ is a weight matrix used for the recurrent computation in order to pass memory along the sequence,\n",
        "- $W$ is a weight matrix used to compute the output of the every timestep (given that every timestep requires an output), and\n",
        "- $h$ is the hidden state (the network's memory) for a given time step.\n",
        "\n",
        "From week 1 (exercise i) we have that\n",
        "$$\\frac{\\partial E}{\\partial o_{t}} = \\hat{y}_{t} - y_{t}\\,,$$\n",
        "where $\\hat{y}_{t}$ is a softmax distribution over model outputs $o_{t}$ at time $t$, and $y_{t}$ is the target label at time $t$. In the above code, this corresponds to\n",
        "```\n",
        "d_o = outputs[t].copy()\n",
        "d_o[np.argmax(targets[t])] -= 1\n",
        "```\n",
        "where `outputs[t]` is the normalized output $\\hat{y}_{t}$, and `targets[t]` is a one-hot encoded target.\n",
        "\n",
        "Now, to compute $\\frac{\\partial E}{\\partial o_{t}} \\frac{\\partial o_{t}}{\\partial W_{t}}$ (denoted `d_W` in the code) we need to determine $\\frac{\\partial o_{t}}{\\partial W_{t}}$ using the function definitions further up.\n",
        "\n",
        "Likewise, we need to determine $\\frac{\\partial E}{\\partial U_{t}},~\\frac{\\partial E}{\\partial V_{t}}$ using the chain rule, where the intermediate steps $\\frac{\\partial E}{\\partial h_{t}}\\,\\frac{\\partial h_{t}}{\\partial f_{t}}$ are pre-computed for you in the code as `d_f`. To compute $\\frac{\\partial E}{\\partial U_{t}},~\\frac{\\partial E}{\\partial V_{t}}$ you will therefore need to determine $\\frac{\\partial f_{t}}{\\partial U_{t}}$ and $\\frac{\\partial f_{t}}{\\partial V_{t}}$, respectively, using the function definitions.\n",
        "\n",
        "For further explanation, [follow this link](https://www.deeplearningbook.org/contents/rnn.html) (section 10.2.2).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pxi7sbLus82R"
      },
      "source": [
        "### Optimization"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NkLXBmU0k2h7"
      },
      "source": [
        "Now that we can do forward passes and compute gradients with backpropagation, we're ready to train our network. For that we will need an optimizer. A common and easy to implement optimization method is stochastic gradient descent (SGD), which has the update rule: $\\theta_{n+1} = \\theta_{n} - \\eta \\frac{\\partial E}{\\partial \\theta_{n}}$, where $\\eta$ is the learning rate and $E$ is our cost function. This is essentially what's going on behind the scenes when you run `optimizer.step()` in PyTorch using the SGD optimizer. If you want to learn more about optimization in a deep learning context, [this is a great starting point](https://arxiv.org/abs/1609.04747)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N4yTZYkBtAVr"
      },
      "source": [
        "def update_parameters(params, grads, lr=1e-3):\n",
        "    # Take a step\n",
        "    for param, grad in zip(params, grads):\n",
        "        param -= lr * grad\n",
        "    \n",
        "    return params"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K7ArhViPtEck"
      },
      "source": [
        "### Training loop"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dOZlMx9bk2iB"
      },
      "source": [
        "We need to define a full training loop with a forward pass, backward pass, optimization step and validation. Training will take approximately 5 minutes, so you might want to read on while the notebook is running."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YMcg90qYs6kW",
        "scrolled": false
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline\n",
        "\n",
        "# Hyper-parameters\n",
        "num_epochs = 1000\n",
        "\n",
        "# Initialize a new network\n",
        "params = init_rnn(hidden_size=hidden_size, vocab_size=vocab_size)\n",
        "\n",
        "# Initialize hidden state as zeros\n",
        "hidden_state = np.zeros((hidden_size, 1))\n",
        "\n",
        "# Track loss\n",
        "training_loss, validation_loss = [], []\n",
        "\n",
        "# For each epoch\n",
        "for i in range(num_epochs):\n",
        "    \n",
        "    # Track loss\n",
        "    epoch_training_loss = 0\n",
        "    epoch_validation_loss = 0\n",
        "    \n",
        "     # For each sentence in validation set\n",
        "    for inputs, targets in validation_set:\n",
        "        \n",
        "        # One-hot encode input and target sequence\n",
        "        inputs_one_hot = one_hot_encode_sequence(inputs, vocab_size)\n",
        "        targets_one_hot = one_hot_encode_sequence(targets, vocab_size)\n",
        "        \n",
        "        # Re-initialize hidden state\n",
        "        hidden_state = np.zeros_like(hidden_state)\n",
        "\n",
        "        # Forward pass\n",
        "        # YOUR CODE HERE!\n",
        "        outputs, hidden_states = \n",
        "\n",
        "        # Backward pass\n",
        "        # YOUR CODE HERE!\n",
        "        loss, _ = \n",
        "        \n",
        "        # Update loss\n",
        "        epoch_validation_loss += loss\n",
        "    \n",
        "    # For each sentence in training set\n",
        "    for inputs, targets in training_set:\n",
        "        \n",
        "        # One-hot encode input and target sequence\n",
        "        inputs_one_hot = one_hot_encode_sequence(inputs, vocab_size)\n",
        "        targets_one_hot = one_hot_encode_sequence(targets, vocab_size)\n",
        "        \n",
        "        # Re-initialize hidden state\n",
        "        hidden_state = np.zeros_like(hidden_state)\n",
        "\n",
        "        # Forward pass\n",
        "        # YOUR CODE HERE!\n",
        "        outputs, hidden_states = \n",
        "\n",
        "        # Backward pass\n",
        "        # YOUR CODE HERE!\n",
        "        loss, grads = \n",
        "        \n",
        "        if np.isnan(loss):\n",
        "            raise ValueError('Gradients have vanished/exploded!')\n",
        "        \n",
        "        # Update parameters\n",
        "        # YOUR CODE HERE!\n",
        "        params = \n",
        "        \n",
        "        # Update loss\n",
        "        epoch_training_loss += loss\n",
        "        \n",
        "    # Save loss for plot\n",
        "    training_loss.append(epoch_training_loss/len(training_set))\n",
        "    validation_loss.append(epoch_validation_loss/len(validation_set))\n",
        "\n",
        "    # Print loss every 100 epochs\n",
        "    if i % 100 == 0:\n",
        "        print(f'Epoch {i}, training loss: {training_loss[-1]}, validation loss: {validation_loss[-1]}')\n",
        "\n",
        "\n",
        "# Get first sentence in test set\n",
        "inputs, targets = test_set[1]\n",
        "\n",
        "# One-hot encode input and target sequence\n",
        "inputs_one_hot = one_hot_encode_sequence(inputs, vocab_size)\n",
        "targets_one_hot = one_hot_encode_sequence(targets, vocab_size)\n",
        "\n",
        "# Initialize hidden state as zeros\n",
        "hidden_state = np.zeros((hidden_size, 1))\n",
        "\n",
        "# Forward pass\n",
        "outputs, hidden_states = forward_pass(inputs_one_hot, hidden_state, params)\n",
        "output_sentence = [idx_to_word[np.argmax(output)] for output in outputs]\n",
        "print('Input sentence:')\n",
        "print(inputs)\n",
        "\n",
        "print('\\nTarget sequence:')\n",
        "print(targets)\n",
        "\n",
        "print('\\nPredicted sequence:')\n",
        "print([idx_to_word[np.argmax(output)] for output in outputs])\n",
        "\n",
        "# Plot training and validation loss\n",
        "epoch = np.arange(len(training_loss))\n",
        "plt.figure()\n",
        "plt.plot(epoch, training_loss, 'r', label='Training loss',)\n",
        "plt.plot(epoch, validation_loss, 'b', label='Validation loss')\n",
        "plt.legend()\n",
        "plt.xlabel('Epoch'), plt.ylabel('NLL')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Uo_SNJYlk2iG"
      },
      "source": [
        "## Exercise E:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q14DnOhZk2iG"
      },
      "source": [
        "Complete the training loop above and run the training. You can leave the hyper-parameters and network size unchanged.\n",
        "\n",
        "A correct implementation should yield a loss of around **1** (using mean CE) or around **4** (using sum CE) after 1000 epochs. Does it work? If not, try to identify the issue -- perhaps something in the backward pass is not right?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nn7QpUZXk2iH"
      },
      "source": [
        "## Extrapolation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BJIJInBhk2iI"
      },
      "source": [
        "Now that we have trained an RNN, it's time to put it to test. We will provide the network with a starting sentence and let it `freestyle` from there!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EOMxSE8Hk2iI"
      },
      "source": [
        "def freestyle(params, sentence='', num_generate=10):\n",
        "    \"\"\"\n",
        "    Takes in a sentence as a string and outputs a sequence\n",
        "    based on the predictions of the RNN.\n",
        "    \n",
        "    Args:\n",
        "     `params`: the parameters of the network\n",
        "     `sentence`: string with whitespace-separated tokens\n",
        "     `num_generate`: the number of tokens to generate\n",
        "    \"\"\"\n",
        "    sentence = sentence.split(' ')\n",
        "    \n",
        "    sentence_one_hot = one_hot_encode_sequence(sentence, vocab_size)\n",
        "    \n",
        "    # Initialize hidden state as zeros\n",
        "    hidden_state = np.zeros((hidden_size, 1))\n",
        "\n",
        "    # Generate hidden state for sentence\n",
        "    outputs, hidden_states = forward_pass(sentence_one_hot, hidden_state, params)\n",
        "    \n",
        "    # Output sentence\n",
        "    output_sentence = sentence\n",
        "    \n",
        "    # Append first prediction\n",
        "    word = idx_to_word[np.argmax(outputs[-1])]    \n",
        "    output_sentence.append(word)\n",
        "    \n",
        "    # Forward pass\n",
        "    for i in range(num_generate):\n",
        "\n",
        "        # Get the latest prediction and latest hidden state\n",
        "        output = outputs[-1]\n",
        "        hidden_state = hidden_states[-1]\n",
        "    \n",
        "        # Reshape our output to match the input shape of our forward pass\n",
        "        output = output.reshape(1, output.shape[0], output.shape[1])\n",
        "    \n",
        "        # Forward pass\n",
        "        outputs, hidden_states = forward_pass(output, hidden_state, params)\n",
        "        \n",
        "        # Compute the index of the most likely word and look up the corresponding word\n",
        "        word = idx_to_word[np.argmax(outputs)]\n",
        "        \n",
        "        output_sentence.append(word)\n",
        "        \n",
        "        if word == 'EOS':\n",
        "            break\n",
        "        \n",
        "    return output_sentence\n",
        "\n",
        "\n",
        "# Perform freestyle (extrapolation)\n",
        "test_examples = ['a a b', 'a a a a b', 'a a a a a a b', 'a', 'r n n']\n",
        "for i, test_example in enumerate(test_examples):\n",
        "    print(f'Example {i}:', test_example)\n",
        "    print('Predicted sequence:', freestyle(params, sentence=test_example), end='\\n\\n')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W8jeDBtFk2iM"
      },
      "source": [
        "## Exercise F:\n",
        "\n",
        "How well does your RNN extrapolate -- does it work as expected? Are there any imperfections? If yes, why could that be?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oKdqe04-k2iN"
      },
      "source": [
        "## Exercise G (optional):"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C7gp_oS8k2iN"
      },
      "source": [
        "Alter the forward pass, backward pass and training loop to handle batches of samples. You will see great improvements!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sGeKM44bvISA"
      },
      "source": [
        "# Introduction to the Long Short-Term Memory (LSTM) Cell\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X44hQ653vNCj"
      },
      "source": [
        "Reading material: [Christopher Olah's walk-through](http://colah.github.io/posts/2015-08-Understanding-LSTMs/).\n",
        "\n",
        "___\n",
        "\n",
        "\n",
        "A vanilla RNN suffers from [the vanishing gradients problem](http://neuralnetworksanddeeplearning.com/chap5.html#the_vanishing_gradient_problem) which gives challenges in saving memory over longer sequences. To combat these issues the gated hidden units were created. The two most prominent gated hidden units are the Long Short-Term Memory (LSTM) cell and the Gated Recurrent Unit (GRU), both of which have shown increased performance in saving and reusing memory in later timesteps. In this exercise, we will focus on LSTM but you would easily be able to go ahead and implement the GRU as well based on the principles that you learn here.\n",
        "\n",
        "Below is a figure of the LSTM cell:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Rgc-g3zwV9f"
      },
      "source": [
        "![lstm](https://i.imgur.com/3VkmUCe.png)\n",
        "Source: https://arxiv.org/abs/1412.7828"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ytasZ5cqw4W1"
      },
      "source": [
        "\n",
        "The LSTM cell contains three gates, input, forget, output gates and a memory cell.\n",
        "The output of the LSTM unit is computed with the following functions, where $\\sigma = \\mathrm{softmax}$.\n",
        "We have input gate $i$, forget gate $f$, and output gate $o$ defines as\n",
        "\n",
        "- $i = \\sigma ( W^i [h_{t-1}, x_t])$\n",
        "\n",
        "- $f = \\sigma ( W^f [h_{t-1},x_t])$\n",
        "\n",
        "- $o = \\sigma ( W^o [h_{t-1},x_t])$\n",
        "\n",
        "where $W^i, W^f, W^o$ are weight matrices applied to a concatenated $h_{t-1}$ (hidden state vector) and $x_t$ (input vector)  for each respective gate.\n",
        "\n",
        "$h_{t-1}$, from the previous time step along with the current input $x_t$ are used to compute the a candidate $g$\n",
        "\n",
        "- $g = \\mathrm{tanh}( W^g [h_{t-1}, x_t])$\n",
        "\n",
        "The value of the cell's memory, $c_t$, is updated as\n",
        "\n",
        "- $c_t = c_{t-1} \\circ f + g \\circ i$\n",
        "\n",
        "where $c_{t-1}$ is the previous memory, and $\\circ$ refers to element-wise multiplication.\n",
        "\n",
        "The output, $h_t$, is computed as\n",
        "\n",
        "- $h_t = \\mathrm{tanh}(c_t) \\circ o$\n",
        "\n",
        "and it is used for both the timestep's output and the next timestep, whereas $c_t$ is exclusively sent to the next timestep.\n",
        "This makes $c_t$ a memory feature, and is not used directly to compute the output of the timestep."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m8_4RWp3k2iQ"
      },
      "source": [
        "## Initialiation of an LSTM network"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MVn8z4VRk2iR"
      },
      "source": [
        "Similarly to before we will implement the forward pass, backward pass, optimization and training loop, now for an LSTM in numpy. Later, we will go back to PyTorch and appreciate how convenient the implementation becomes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w4HETMthk2iS"
      },
      "source": [
        "# Size of concatenated hidden + input vector\n",
        "z_size = hidden_size + vocab_size \n",
        "\n",
        "def init_lstm(hidden_size, vocab_size, z_size):\n",
        "    \"\"\"\n",
        "    Initializes our LSTM network.\n",
        "    \n",
        "    Args:\n",
        "     `hidden_size`: the dimensions of the hidden state\n",
        "     `vocab_size`: the dimensions of our vocabulary\n",
        "     `z_size`: the dimensions of the concatenated input \n",
        "    \"\"\"\n",
        "    # Weight matrix (forget gate)\n",
        "    # YOUR CODE HERE!\n",
        "    W_f = np.zeros((, ))\n",
        "    \n",
        "    # Bias for forget gate\n",
        "    b_f = np.zeros((hidden_size, 1))\n",
        "\n",
        "    # Weight matrix (input gate)\n",
        "    # YOUR CODE HERE!\n",
        "    W_i = np.zeros((, ))\n",
        "    \n",
        "    # Bias for input gate\n",
        "    b_i = np.zeros((hidden_size, 1))\n",
        "\n",
        "    # Weight matrix (candidate)\n",
        "    # YOUR CODE HERE!\n",
        "    W_g = np.zeros((, ))\n",
        "    \n",
        "    # Bias for candidate\n",
        "    b_g = np.zeros((hidden_size, 1))\n",
        "\n",
        "    # Weight matrix of the output gate\n",
        "    # YOUR CODE HERE!\n",
        "    W_o = np.zeros((, ))\n",
        "    \n",
        "    # Bias for output gate\n",
        "    b_o = np.zeros((hidden_size, 1))\n",
        "\n",
        "    # Weight matrix relating the hidden-state to the output\n",
        "    # YOUR CODE HERE!\n",
        "    W_v = np.zeros((, ))\n",
        "    \n",
        "    # Bias for logits\n",
        "    b_v = np.zeros((vocab_size, 1))\n",
        "    \n",
        "    # Initialize weights according to https://arxiv.org/abs/1312.6120\n",
        "    W_f = init_orthogonal(W_f)\n",
        "    W_i = init_orthogonal(W_i)\n",
        "    W_g = init_orthogonal(W_g)\n",
        "    W_o = init_orthogonal(W_o)\n",
        "    W_v = init_orthogonal(W_v)\n",
        "\n",
        "    return W_f, W_i, W_g, W_o, W_v, b_f, b_i, b_g, b_o, b_v\n",
        "\n",
        "\n",
        "params = init_lstm(hidden_size=hidden_size, vocab_size=vocab_size, z_size=z_size)\n",
        "print('W_f:', params[0].shape)\n",
        "print('W_i:', params[1].shape)\n",
        "print('W_g:', params[2].shape)\n",
        "print('W_o:', params[3].shape)\n",
        "print('W_v:', params[4].shape)\n",
        "print('b_i:', params[5].shape)\n",
        "print('b_g:', params[6].shape)\n",
        "print('b_o:', params[7].shape)\n",
        "print('b_v:', params[8].shape)\n",
        "\n",
        "for param in params:\n",
        "    assert param.ndim == 2, \\\n",
        "        'all parameters should be 2-dimensional '\\\n",
        "        '(hint: a dimension can simply have size 1)'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "siljVza0k2iV"
      },
      "source": [
        "## Exercise H:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uU6Bn1vBk2iW"
      },
      "source": [
        "Implement the initializations in the code above. Refer to the equations and the figure if you're in doubt."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gKu-bfhzk2iY"
      },
      "source": [
        "## Forward pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZBa5XNvEk2iZ"
      },
      "source": [
        "As before we need to bring the equations into pieces of code. This is done in the following function:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wGeiytn9k2ia"
      },
      "source": [
        "def forward(inputs, h_prev, C_prev, p):\n",
        "    \"\"\"\n",
        "    Arguments:\n",
        "    x -- your input data at timestep \"t\", numpy array of shape (n_x, m).\n",
        "    h_prev -- Hidden state at timestep \"t-1\", numpy array of shape (n_a, m)\n",
        "    C_prev -- Memory state at timestep \"t-1\", numpy array of shape (n_a, m)\n",
        "    p -- python list containing:\n",
        "                        W_f -- Weight matrix of the forget gate, numpy array of shape (n_a, n_a + n_x)\n",
        "                        b_f -- Bias of the forget gate, numpy array of shape (n_a, 1)\n",
        "                        W_i -- Weight matrix of the update gate, numpy array of shape (n_a, n_a + n_x)\n",
        "                        b_i -- Bias of the update gate, numpy array of shape (n_a, 1)\n",
        "                        W_g -- Weight matrix of the first \"tanh\", numpy array of shape (n_a, n_a + n_x)\n",
        "                        b_g --  Bias of the first \"tanh\", numpy array of shape (n_a, 1)\n",
        "                        W_o -- Weight matrix of the output gate, numpy array of shape (n_a, n_a + n_x)\n",
        "                        b_o --  Bias of the output gate, numpy array of shape (n_a, 1)\n",
        "                        W_v -- Weight matrix relating the hidden-state to the output, numpy array of shape (n_v, n_a)\n",
        "                        b_v -- Bias relating the hidden-state to the output, numpy array of shape (n_v, 1)\n",
        "    Returns:\n",
        "    z_s, f_s, i_s, g_s, C_s, o_s, h_s, v_s -- lists of size m containing the computations in each forward pass\n",
        "    outputs -- prediction at timestep \"t\", numpy array of shape (n_v, m)\n",
        "    \"\"\"\n",
        "    assert h_prev.shape == (hidden_size, 1)\n",
        "    assert C_prev.shape == (hidden_size, 1)\n",
        "\n",
        "    # First we unpack our parameters\n",
        "    W_f, W_i, W_g, W_o, W_v, b_f, b_i, b_g, b_o, b_v = p\n",
        "    \n",
        "    # Save a list of computations for each of the components in the LSTM\n",
        "    x_s, z_s, f_s, i_s,  = [], [] ,[], []\n",
        "    g_s, C_s, o_s, h_s = [], [] ,[], []\n",
        "    v_s, output_s =  [], [] \n",
        "    \n",
        "    # Append the initial cell and hidden state to their respective lists\n",
        "    h_s.append(h_prev)\n",
        "    C_s.append(C_prev)\n",
        "    \n",
        "    for x in inputs:\n",
        "        \n",
        "        # Concatenate input and hidden state\n",
        "        z = np.row_stack((h_prev, x))\n",
        "        z_s.append(z)\n",
        "        \n",
        "        # Calculate forget gate\n",
        "        # YOUR CODE HERE!\n",
        "        f = \n",
        "        f_s.append(f)\n",
        "        \n",
        "        # Calculate input gate\n",
        "        # YOUR CODE HERE!\n",
        "        i = \n",
        "        i_s.append(i)\n",
        "        \n",
        "        # Calculate candidate\n",
        "        g = tanh(np.dot(W_g, z) + b_g)\n",
        "        g_s.append(g)\n",
        "        \n",
        "        # Calculate memory state\n",
        "        # YOUR CODE HERE!\n",
        "        C_prev = \n",
        "        C_s.append(C_prev)\n",
        "        \n",
        "        # Calculate output gate\n",
        "        # YOUR CODE HERE!\n",
        "        o = \n",
        "        o_s.append(o)\n",
        "        \n",
        "        # Calculate hidden state\n",
        "        h_prev = o * tanh(C_prev)\n",
        "        h_s.append(h_prev)\n",
        "\n",
        "        # Calculate logits\n",
        "        v = np.dot(W_v, h_prev) + b_v\n",
        "        v_s.append(v)\n",
        "        \n",
        "        # Calculate softmax\n",
        "        output = softmax(v)\n",
        "        output_s.append(output)\n",
        "\n",
        "    return z_s, f_s, i_s, g_s, C_s, o_s, h_s, v_s, output_s\n",
        "\n",
        "\n",
        "# Get first sentence in test set\n",
        "inputs, targets = test_set[1]\n",
        "\n",
        "# One-hot encode input and target sequence\n",
        "inputs_one_hot = one_hot_encode_sequence(inputs, vocab_size)\n",
        "targets_one_hot = one_hot_encode_sequence(targets, vocab_size)\n",
        "\n",
        "# Initialize hidden state as zeros\n",
        "h = np.zeros((hidden_size, 1))\n",
        "c = np.zeros((hidden_size, 1))\n",
        "\n",
        "# Forward pass\n",
        "z_s, f_s, i_s, g_s, C_s, o_s, h_s, v_s, outputs = forward(inputs_one_hot, h, c, params)\n",
        "\n",
        "output_sentence = [idx_to_word[np.argmax(output)] for output in outputs]\n",
        "print('Input sentence:')\n",
        "print(inputs)\n",
        "\n",
        "print('\\nTarget sequence:')\n",
        "print(targets)\n",
        "\n",
        "print('\\nPredicted sequence:')\n",
        "print([idx_to_word[np.argmax(output)] for output in outputs])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o4u5sjlKk2id"
      },
      "source": [
        "## Exercise I:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AhNbDEuMk2id"
      },
      "source": [
        "Complete the implementation of the LSTM forward pass above. Refer to the equations and figures further up if you're in doubt."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G-6tLfsuk2ie"
      },
      "source": [
        "## Backward pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H5-ALOUMk2if"
      },
      "source": [
        "Similar to the RNN in numpy we also need to specify a backward pass. Fortunately, we have already done the work for you here :-)\n",
        "\n",
        "Feel free to dive into the code to get a better intuition of what is going on -- otherwise you can jump straight to the training loop."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "snUFSDBDk2ig"
      },
      "source": [
        "def backward(z, f, i, g, C, o, h, v, outputs, targets, p = params):\n",
        "    \"\"\"\n",
        "    Arguments:\n",
        "    z -- your concatenated input data  as a list of size m.\n",
        "    f -- your forget gate computations as a list of size m.\n",
        "    i -- your input gate computations as a list of size m.\n",
        "    g -- your candidate computations as a list of size m.\n",
        "    C -- your Cell states as a list of size m+1.\n",
        "    o -- your output gate computations as a list of size m.\n",
        "    h -- your Hidden state computations as a list of size m+1.\n",
        "    v -- your logit computations as a list of size m.\n",
        "    outputs -- your outputs as a list of size m.\n",
        "    targets -- your targets as a list of size m.\n",
        "    p -- python list containing:\n",
        "                        W_f -- Weight matrix of the forget gate, numpy array of shape (n_a, n_a + n_x)\n",
        "                        b_f -- Bias of the forget gate, numpy array of shape (n_a, 1)\n",
        "                        W_i -- Weight matrix of the update gate, numpy array of shape (n_a, n_a + n_x)\n",
        "                        b_i -- Bias of the update gate, numpy array of shape (n_a, 1)\n",
        "                        W_g -- Weight matrix of the first \"tanh\", numpy array of shape (n_a, n_a + n_x)\n",
        "                        b_g --  Bias of the first \"tanh\", numpy array of shape (n_a, 1)\n",
        "                        W_o -- Weight matrix of the output gate, numpy array of shape (n_a, n_a + n_x)\n",
        "                        b_o --  Bias of the output gate, numpy array of shape (n_a, 1)\n",
        "                        W_v -- Weight matrix relating the hidden-state to the output, numpy array of shape (n_v, n_a)\n",
        "                        b_v -- Bias relating the hidden-state to the output, numpy array of shape (n_v, 1)\n",
        "    Returns:\n",
        "    loss -- crossentropy loss for all elements in output\n",
        "    grads -- lists of gradients of every element in p\n",
        "    \"\"\"\n",
        "\n",
        "    # Unpack parameters\n",
        "    W_f, W_i, W_g, W_o, W_v, b_f, b_i, b_g, b_o, b_v = p\n",
        "\n",
        "    # Initialize gradients as zero\n",
        "    W_f_d = np.zeros_like(W_f)\n",
        "    b_f_d = np.zeros_like(b_f)\n",
        "\n",
        "    W_i_d = np.zeros_like(W_i)\n",
        "    b_i_d = np.zeros_like(b_i)\n",
        "\n",
        "    W_g_d = np.zeros_like(W_g)\n",
        "    b_g_d = np.zeros_like(b_g)\n",
        "\n",
        "    W_o_d = np.zeros_like(W_o)\n",
        "    b_o_d = np.zeros_like(b_o)\n",
        "\n",
        "    W_v_d = np.zeros_like(W_v)\n",
        "    b_v_d = np.zeros_like(b_v)\n",
        "    \n",
        "    # Set the next cell and hidden state equal to zero\n",
        "    dh_next = np.zeros_like(h[0])\n",
        "    dC_next = np.zeros_like(C[0])\n",
        "        \n",
        "    # Track loss\n",
        "    loss = 0\n",
        "    \n",
        "    for t in reversed(range(len(outputs))):\n",
        "        \n",
        "        # Compute the cross entropy\n",
        "        loss += -np.mean(np.log(outputs[t]) * targets[t])\n",
        "        # Get the previous hidden cell state\n",
        "        C_prev= C[t-1]\n",
        "        \n",
        "        # Compute the derivative of the relation of the hidden-state to the output gate\n",
        "        dv = np.copy(outputs[t])\n",
        "        dv[np.argmax(targets[t])] -= 1\n",
        "\n",
        "        # Update the gradient of the relation of the hidden-state to the output gate\n",
        "        W_v_d += np.dot(dv, h[t].T)\n",
        "        b_v_d += dv\n",
        "\n",
        "        # Compute the derivative of the hidden state and output gate\n",
        "        dh = np.dot(W_v.T, dv)        \n",
        "        dh += dh_next\n",
        "        do = dh * tanh(C[t])\n",
        "        do = sigmoid(o[t], derivative=True)*do\n",
        "        \n",
        "        # Update the gradients with respect to the output gate\n",
        "        W_o_d += np.dot(do, z[t].T)\n",
        "        b_o_d += do\n",
        "\n",
        "        # Compute the derivative of the cell state and candidate g\n",
        "        dC = np.copy(dC_next)\n",
        "        dC += dh * o[t] * tanh(tanh(C[t]), derivative=True)\n",
        "        dg = dC * i[t]\n",
        "        dg = tanh(g[t], derivative=True) * dg\n",
        "        \n",
        "        # Update the gradients with respect to the candidate\n",
        "        W_g_d += np.dot(dg, z[t].T)\n",
        "        b_g_d += dg\n",
        "\n",
        "        # Compute the derivative of the input gate and update its gradients\n",
        "        di = dC * g[t]\n",
        "        di = sigmoid(i[t], True) * di\n",
        "        W_i_d += np.dot(di, z[t].T)\n",
        "        b_i_d += di\n",
        "\n",
        "        # Compute the derivative of the forget gate and update its gradients\n",
        "        df = dC * C_prev\n",
        "        df = sigmoid(f[t]) * df\n",
        "        W_f_d += np.dot(df, z[t].T)\n",
        "        b_f_d += df\n",
        "\n",
        "        # Compute the derivative of the input and update the gradients of the previous hidden and cell state\n",
        "        dz = (np.dot(W_f.T, df)\n",
        "             + np.dot(W_i.T, di)\n",
        "             + np.dot(W_g.T, dg)\n",
        "             + np.dot(W_o.T, do))\n",
        "        dh_prev = dz[:hidden_size, :]\n",
        "        dC_prev = f[t] * dC\n",
        "        \n",
        "    grads= W_f_d, W_i_d, W_g_d, W_o_d, W_v_d, b_f_d, b_i_d, b_g_d, b_o_d, b_v_d\n",
        "    \n",
        "    # Clip gradients\n",
        "    grads = clip_gradient_norm(grads)\n",
        "    \n",
        "    return loss, grads\n",
        "\n",
        "\n",
        "# Perform a backward pass\n",
        "loss, grads = backward(z_s, f_s, i_s, g_s, C_s, o_s, h_s, v_s, outputs, targets_one_hot, params)\n",
        "\n",
        "print('We get a loss of:')\n",
        "print(loss)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z4r4mgFsk2ik"
      },
      "source": [
        "## Training loop"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7ok73S7ck2ik"
      },
      "source": [
        "It's time to check whether the LSTM really works! Let's train it for some number of epochs and observe the result."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "23N8SVeik2ik"
      },
      "source": [
        "# Hyper-parameters\n",
        "num_epochs = 200\n",
        "\n",
        "# Initialize a new network\n",
        "z_size = hidden_size + vocab_size # Size of concatenated hidden + input vector\n",
        "params = init_lstm(hidden_size=hidden_size, vocab_size=vocab_size, z_size=z_size)\n",
        "\n",
        "# Initialize hidden state as zeros\n",
        "hidden_state = np.zeros((hidden_size, 1))\n",
        "\n",
        "# Track loss\n",
        "training_loss, validation_loss = [], []\n",
        "\n",
        "# For each epoch\n",
        "for i in range(num_epochs):\n",
        "    \n",
        "    # Track loss\n",
        "    epoch_training_loss = 0\n",
        "    epoch_validation_loss = 0\n",
        "    \n",
        "    # For each sentence in validation set\n",
        "    for inputs, targets in validation_set:\n",
        "        \n",
        "        # One-hot encode input and target sequence\n",
        "        inputs_one_hot = one_hot_encode_sequence(inputs, vocab_size)\n",
        "        targets_one_hot = one_hot_encode_sequence(targets, vocab_size)\n",
        "\n",
        "        # Initialize hidden state and cell state as zeros\n",
        "        h = np.zeros((hidden_size, 1))\n",
        "        c = np.zeros((hidden_size, 1))\n",
        "\n",
        "        # Forward pass\n",
        "        z_s, f_s, i_s, g_s, C_s, o_s, h_s, v_s, outputs = forward(inputs_one_hot, h, c, params)\n",
        "        \n",
        "        # Backward pass\n",
        "        loss, _ = backward(z_s, f_s, i_s, g_s, C_s, o_s, h_s, v_s, outputs, targets_one_hot, params)\n",
        "        \n",
        "        # Update loss\n",
        "        epoch_validation_loss += loss\n",
        "    \n",
        "    # For each sentence in training set\n",
        "    for inputs, targets in training_set:\n",
        "        \n",
        "        # One-hot encode input and target sequence\n",
        "        inputs_one_hot = one_hot_encode_sequence(inputs, vocab_size)\n",
        "        targets_one_hot = one_hot_encode_sequence(targets, vocab_size)\n",
        "\n",
        "        # Initialize hidden state and cell state as zeros\n",
        "        h = np.zeros((hidden_size, 1))\n",
        "        c = np.zeros((hidden_size, 1))\n",
        "\n",
        "        # Forward pass\n",
        "        z_s, f_s, i_s, g_s, C_s, o_s, h_s, v_s, outputs = forward(inputs_one_hot, h, c, params)\n",
        "        \n",
        "        # Backward pass\n",
        "        loss, grads = backward(z_s, f_s, i_s, g_s, C_s, o_s, h_s, v_s, outputs, targets_one_hot, params)\n",
        "        \n",
        "        # Update parameters\n",
        "        params = update_parameters(params, grads, lr=1e-1)\n",
        "        \n",
        "        # Update loss\n",
        "        epoch_training_loss += loss\n",
        "                \n",
        "    # Save loss for plot\n",
        "    training_loss.append(epoch_training_loss/len(training_set))\n",
        "    validation_loss.append(epoch_validation_loss/len(validation_set))\n",
        "\n",
        "    # Print loss every 10 epochs\n",
        "    if i % 10 == 0:\n",
        "        print(f'Epoch {i}, training loss: {training_loss[-1]}, validation loss: {validation_loss[-1]}')\n",
        "\n",
        "    \n",
        "# Get first sentence in test set\n",
        "inputs, targets = test_set[1]\n",
        "\n",
        "# One-hot encode input and target sequence\n",
        "inputs_one_hot = one_hot_encode_sequence(inputs, vocab_size)\n",
        "targets_one_hot = one_hot_encode_sequence(targets, vocab_size)\n",
        "\n",
        "# Initialize hidden state as zeros\n",
        "h = np.zeros((hidden_size, 1))\n",
        "c = np.zeros((hidden_size, 1))\n",
        "\n",
        "# Forward pass\n",
        "z_s, f_s, i_s, g_s, C_s, o_s, h_s, v_s, outputs = forward(inputs_one_hot, h, c, params)\n",
        "\n",
        "# Print example\n",
        "print('Input sentence:')\n",
        "print(inputs)\n",
        "\n",
        "print('\\nTarget sequence:')\n",
        "print(targets)\n",
        "\n",
        "print('\\nPredicted sequence:')\n",
        "print([idx_to_word[np.argmax(output)] for output in outputs])\n",
        "\n",
        "# Plot training and validation loss\n",
        "epoch = np.arange(len(training_loss))\n",
        "plt.figure()\n",
        "plt.plot(epoch, training_loss, 'r', label='Training loss',)\n",
        "plt.plot(epoch, validation_loss, 'b', label='Validation loss')\n",
        "plt.legend()\n",
        "plt.xlabel('Epoch'), plt.ylabel('NLL')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_3nTj0CRk2in"
      },
      "source": [
        "## Exercise J:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GFlWfDOTk2io"
      },
      "source": [
        "Run the training loop above. Compare your LSTM learning curve (NLL and number of epochs) to the vanilla RNN from earlier. Do you observe any improvements? Motivate your answer."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gi51eWgKxyOk"
      },
      "source": [
        "## PyTorch implementation of the LSTM\n",
        "\n",
        "Now that we know how the LSTM cell works, let's see how easy it is to use in PyTorch!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O6HDdJLuk2ip"
      },
      "source": [
        "Definition of our LSTM network. We define a LSTM layer using the [nn.LSTM](https://pytorch.org/docs/stable/nn.html#lstm) class. The LSTM layer takes as argument the size of the input and the size of the hidden state like in our numpy implementation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8UGrvknfk2ip"
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "class MyRecurrentNet(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(MyRecurrentNet, self).__init__()\n",
        "        \n",
        "        # Recurrent layer\n",
        "        # YOUR CODE HERE!\n",
        "        self.lstm = nn.LSTM(input_size=vocab_size,\n",
        "                         hidden_size=50,\n",
        "                         num_layers=1,\n",
        "                         bidirectional=False)\n",
        "        \n",
        "        # Output layer\n",
        "        self.l_out = nn.Linear(in_features=50,\n",
        "                            out_features=vocab_size,\n",
        "                            bias=False)\n",
        "        \n",
        "    def forward(self, x):\n",
        "        # RNN returns output and last hidden state\n",
        "        x, (h, c) = self.lstm(x)\n",
        "        \n",
        "        # Flatten output for feed-forward layer\n",
        "        x = x.view(-1, self.lstm.hidden_size)\n",
        "        \n",
        "        # Output layer\n",
        "        x = self.l_out(x)\n",
        "        \n",
        "        return x\n",
        "\n",
        "net = MyRecurrentNet()\n",
        "print(net)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J6r3bPwYk2is"
      },
      "source": [
        "## Exercise K:\n",
        "\n",
        "Define an LSTM for our recurrent neural network `MyRecurrentNet` above. A single LSTM layer is sufficient. What should the input size and hidden size be? Hint: use the PyTorch documentation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4WpALf2-x7Ty"
      },
      "source": [
        "### Training loop"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IJ8LWatlk2is"
      },
      "source": [
        "It's time for us to train our network. In the section below, you will get to put your deep learning skills to use and create your own training loop. You may want to consult previous exercises if you cannot recall how to define the training loop."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2URKsyFDx8xG"
      },
      "source": [
        "# Hyper-parameters\n",
        "num_epochs = 200\n",
        "\n",
        "# Initialize a new network\n",
        "net = MyRecurrentNet()\n",
        "\n",
        "# Define a loss function and optimizer for this problem\n",
        "# YOUR CODE HERE!\n",
        "criterion = \n",
        "optimizer = \n",
        "\n",
        "# Track loss\n",
        "training_loss, validation_loss = [], []\n",
        "\n",
        "# For each epoch\n",
        "for i in range(num_epochs):\n",
        "    \n",
        "    # Track loss\n",
        "    epoch_training_loss = 0\n",
        "    epoch_validation_loss = 0\n",
        "    \n",
        "    net.eval()\n",
        "        \n",
        "    # For each sentence in validation set\n",
        "    for inputs, targets in validation_set:\n",
        "        \n",
        "        # One-hot encode input and target sequence\n",
        "        inputs_one_hot = one_hot_encode_sequence(inputs, vocab_size)\n",
        "        targets_idx = [word_to_idx[word] for word in targets]\n",
        "        \n",
        "        # Convert input to tensor\n",
        "        inputs_one_hot = torch.Tensor(inputs_one_hot)\n",
        "        inputs_one_hot = inputs_one_hot.permute(0, 2, 1)\n",
        "        \n",
        "        # Convert target to tensor\n",
        "        targets_idx = torch.LongTensor(targets_idx)\n",
        "        \n",
        "        # Forward pass\n",
        "        # YOUR CODE HERE!\n",
        "        outputs = \n",
        "        \n",
        "        # Compute loss\n",
        "        # YOUR CODE HERE!\n",
        "        loss = \n",
        "        \n",
        "        # Update loss\n",
        "        epoch_validation_loss += loss.detach().numpy()\n",
        "    \n",
        "    net.train()\n",
        "    \n",
        "    # For each sentence in training set\n",
        "    for inputs, targets in training_set:\n",
        "        \n",
        "        # One-hot encode input and target sequence\n",
        "        inputs_one_hot = one_hot_encode_sequence(inputs, vocab_size)\n",
        "        targets_idx = [word_to_idx[word] for word in targets]\n",
        "        \n",
        "        # Convert input to tensor\n",
        "        inputs_one_hot = torch.Tensor(inputs_one_hot)\n",
        "        inputs_one_hot = inputs_one_hot.permute(0, 2, 1)\n",
        "        \n",
        "        # Convert target to tensor\n",
        "        targets_idx = torch.LongTensor(targets_idx)\n",
        "        \n",
        "        # Forward pass\n",
        "        # YOUR CODE HERE!\n",
        "        outputs = \n",
        "        \n",
        "        # Compute loss\n",
        "        # YOUR CODE HERE!\n",
        "        loss = \n",
        "        \n",
        "        # Backward pass\n",
        "        # YOUR CODE HERE!\n",
        "        # zero grad, backward, step...\n",
        "        \n",
        "        # Update loss\n",
        "        epoch_training_loss += loss.detach().numpy()\n",
        "        \n",
        "    # Save loss for plot\n",
        "    training_loss.append(epoch_training_loss/len(training_set))\n",
        "    validation_loss.append(epoch_validation_loss/len(validation_set))\n",
        "\n",
        "    # Print loss every 10 epochs\n",
        "    if i % 10 == 0:\n",
        "        print(f'Epoch {i}, training loss: {training_loss[-1]}, validation loss: {validation_loss[-1]}')\n",
        "\n",
        "        \n",
        "# Get first sentence in test set\n",
        "inputs, targets = test_set[1]\n",
        "\n",
        "# One-hot encode input and target sequence\n",
        "inputs_one_hot = one_hot_encode_sequence(inputs, vocab_size)\n",
        "targets_idx = [word_to_idx[word] for word in targets]\n",
        "\n",
        "# Convert input to tensor\n",
        "inputs_one_hot = torch.Tensor(inputs_one_hot)\n",
        "inputs_one_hot = inputs_one_hot.permute(0, 2, 1)\n",
        "\n",
        "# Convert target to tensor\n",
        "targets_idx = torch.LongTensor(targets_idx)\n",
        "\n",
        "# Forward pass\n",
        "outputs = net.forward(inputs_one_hot).data.numpy()\n",
        "\n",
        "print('\\nInput sequence:')\n",
        "print(inputs)\n",
        "\n",
        "print('\\nTarget sequence:')\n",
        "print(targets)\n",
        "\n",
        "print('\\nPredicted sequence:')\n",
        "print([idx_to_word[np.argmax(output)] for output in outputs])\n",
        "\n",
        "# Plot training and validation loss\n",
        "epoch = np.arange(len(training_loss))\n",
        "plt.figure()\n",
        "plt.plot(epoch, training_loss, 'r', label='Training loss',)\n",
        "plt.plot(epoch, validation_loss, 'b', label='Validation loss')\n",
        "plt.legend()\n",
        "plt.xlabel('Epoch'), plt.ylabel('NLL')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ydr7Czg_k2iw"
      },
      "source": [
        "## Exercise L:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M74xjOq0k2ix"
      },
      "source": [
        "Finish the training loop above and run it. Compare your loss to the NumPy implementation. Are they similar? How do they differ?\n",
        "\n",
        "Try to play around with the choice of hyper-parameters, optimizer, and hidden dimensions. How much can you improve the negative log-likelihood by these simple changes?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M93ORx95k2ix"
      },
      "source": [
        "## Exercise M (optional):\n",
        "\n",
        "Aside from the LSTM cell, various other RNN cells exist. The gated recurrent unit (GRU) is a variation of the LSTM cell that uses less gating mechanisms. Try to look it up in the [PyTorch documentation](https://pytorch.org/docs/stable/nn.html#gru) and switch out the LSTM cell in the code above. What do you notice in terms of performance and convergence speed?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "juN400Ekk2iz"
      },
      "source": [
        "## Exercise N (optional):\n",
        "\n",
        "Go back and generate a more complex patterned dataset to learn from. Do you see any significant differences between the vanilla RNN and LSTM when you increase the difficulty of the task?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v68YEkEBk2iz"
      },
      "source": [
        "# It works, now what?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NjpqSrSuk2i0"
      },
      "source": [
        "In this notebook you have learned how to use embeddings, recurrent neural networks, and the LSTM cell in particular.\n",
        "\n",
        "As we have already seen, RNNs are excellent for sequential data such as language. But what do we do if we're modelling data with strong dependency in both directions? Like in many things deep learning, we can build powerful models by stacking layers on top of each other; *bi-directional* RNNs consist of two LSTM cells, one for each direction. A sequence is first fed into the forward LSTM cell and the reversed sequence is then used as input to the backward LSTM cell together with the last hidden state from the forward LSTM cell. Follow [this link](https://pdfs.semanticscholar.org/4b80/89bc9b49f84de43acc2eb8900035f7d492b2.pdf) for the original paper from 1997(!).\n",
        "\n",
        "For even deeper representations, multiple layers of both uni-directional and bi-directional RNNs can be stacked ontop of each other, just like feed-forward and convolutional layers. For more information on this, check out the [LSTM PyTorch documentation](https://pytorch.org/docs/stable/nn.html#lstm). Next week we will also explore ways to combine RNNs with other types of layers for even more expressive function approximators."
      ]
    }
  ]
}