NRI_student.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "NRI_student.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "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.0"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/timlacroix/nri_practical_session/blob/master/NRI_student.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QS6qWID378AP",
        "colab_type": "text"
      },
      "source": [
        "# Neural Relational Inference\n",
        "\n",
        "In this session, we implement the ideas described in https://arxiv.org/pdf/1802.04687.pdf .\n",
        "Most of the code in the solution has been adapted from https://github.com/ethanfetaya/NRI .\n",
        "\n",
        "First add this drive folder to your own google drive account :\n",
        "https://drive.google.com/open?id=10Awx22Z8vah5MxBrCSgQGuQaS2HdG2ae\n",
        "\n",
        "Then follow these setup instructions. The `ls` should show one data folder and a utils.py file."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J4L1qRoN8CYv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "## SETUP\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CWqIcOtf8oSx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%cd /content/drive/My\\ Drive/Summer_School\n",
        "!ls"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kRg_2aDa78AQ",
        "colab_type": "text"
      },
      "source": [
        "## Data, baselines and evaluations"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "grvv5Gl078AR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from utils import load_data\n",
        "\n",
        "loaders, location_range, velocity_range = load_data(batch_size=1, suffix='_springs5')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BfUs5pTv78AT",
        "colab_type": "text"
      },
      "source": [
        "### Plotting the input data\n",
        "\n",
        "Find a good way to display the input data. Display both the particles and the relation matrix "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X06q8gxZ78AT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from matplotlib import pyplot as plt\n",
        "import numpy as np\n",
        "\n",
        "plt.style.use('seaborn-notebook')\n",
        "plt.figure(dpi=150)\n",
        "\n",
        "\n",
        "num_atoms = 5\n",
        "off_diag_idx = np.ravel_multi_index(\n",
        "    np.where(np.ones((num_atoms, num_atoms)) - np.eye(num_atoms)),\n",
        "    [num_atoms, num_atoms]\n",
        ")\n",
        "for atoms, edges in loaders['train']:\n",
        "    # edges contains the off-diagonal elements of the interaction matrix\n",
        "    interactions = np.reshape(np.zeros((num_atoms, num_atoms)), [-1, 25])\n",
        "    interactions[0][off_diag_idx] = edges\n",
        "    interactions = np.reshape(interactions, [5,5])\n",
        "    # now interactions_{i,j} contains whether or not i and j interact\n",
        "\n",
        "    # TODO : plot the atoms and their interactions\n",
        "    \n",
        "    break\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RuVI0xht78AW",
        "colab_type": "text"
      },
      "source": [
        "## LSTM Baseline\n",
        "\n",
        "It's always a good idea to start with \"simple\" algorithms first. Here we'll start by trying to predict the joint trajectories with an LSTM. To do this, we first map the coordinates to a feature space with an MLP. Then we apply the LSTM in this feature space, before decoding the output of the LSTM to the original coordinates. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4XuyUZTm78AX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from __future__ import division\n",
        "from __future__ import print_function\n",
        "\n",
        "import time\n",
        "import pickle\n",
        "import os\n",
        "import datetime\n",
        "\n",
        "import torch.optim as optim\n",
        "from torch.optim import lr_scheduler\n",
        "\n",
        "suffix = '_springs5'\n",
        "n_atoms =  5\n",
        "\n",
        "epochs = 1\n",
        "\n",
        "n_hidden = 256\n",
        "n_layers = 2\n",
        "\n",
        "batch_size = 128\n",
        "learning_rate = 5e-4\n",
        "dropout = 0\n",
        "temp = 0.5\n",
        "\n",
        "timesteps = 49\n",
        "prediction_steps = 10\n",
        "valid_freq = 1\n",
        "\n",
        "var = 5e-5\n",
        "\n",
        "loaders, location_range, velocity_range = load_data(batch_size=batch_size, suffix=suffix)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eDhROnRa78AY",
        "colab_type": "text"
      },
      "source": [
        "### Model\n",
        "\n",
        "Complete the following skeleton to implement the LSTM baseline described in the appendix of the paper.\n",
        "    <ul>\n",
        "        <li> step: <ul>\n",
        "            <li>$f_t = \\textrm{pos_encoder}(x_t)$</li>\n",
        "            <li>$\\delta, h_{t+1} = \\textrm{LSTM}(f_t, h_t)$</li>\n",
        "            <li>return  $x_{t+1} = x_t + \\textrm{pos_decoder}(\\delta)$ and $h_{t+1}$</li>\n",
        "            </ul>\n",
        "        </li>\n",
        "        <li> forward: <ul>\n",
        "            <li>run step for $b$ <em>burn-in</em> steps with true data as input.</li>\n",
        "            <li>then use the previous prediction as input. </li>\n",
        "            <li>return full sequence of outputs</li>\n",
        "            </ul></li>\n",
        "    </ul>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RDBBMrGt78AZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "from torch import nn\n",
        "from torch.functional import F\n",
        "\n",
        "class RecurrentBaseline(nn.Module):\n",
        "    \"\"\"LSTM model for joint trajectory prediction.\"\"\"\n",
        "\n",
        "    def __init__(self, n_dims, n_hid, n_atoms, n_layers, do_prob=0.):\n",
        "        super(RecurrentBaseline, self).__init__()\n",
        "        \n",
        "        # Encoder from positions to n_hid dimensional space\n",
        "        # The architecture is linear / relu / dropout(do_prob) / linear / relu\n",
        "        self.pos_encoder = ...\n",
        "        \n",
        "        # RNN : n_atoms * n_hid -> n_atoms * n_hid. LSTM with n_layers.\n",
        "        self.rnn = nn.LSTM(n_atoms * n_hid, n_atoms * n_hid, n_layers)  # TODO\n",
        "\n",
        "        # Decode predicted *joint* configuration to physical *joint* location\n",
        "        # The architecture is linear / relu / linear\n",
        "        self.pos_decoder = ...\n",
        "\n",
        "    def step(self, ins, hidden=None):\n",
        "        # Input shape: [num_sims, n_atoms, n_dims]\n",
        "        \n",
        "        # Apply first MLP to encode the coordinates\n",
        "\n",
        "        # View to join the last two dimensions. Add a dummy time dimension at the beginning of x\n",
        "\n",
        "        # Apply LSTM given hidden and encoded input\n",
        "\n",
        "        # remove extraneous time dimension\n",
        "\n",
        "        # Apply second MLP to decode the output of the LSTM and compute delta\n",
        "\n",
        "        # View to separate the last two dimensions again\n",
        "\n",
        "        # Add delta to inputs\n",
        "\n",
        "\n",
        "        # Return both output and hidden\n",
        "        return x, hidden\n",
        "\n",
        "    def forward(self, inputs, burn_in_steps=1):\n",
        "        # Input shape: [num_sims, num_things, num_timesteps, n_in]\n",
        "\n",
        "        outputs = []\n",
        "        hidden = None\n",
        "\n",
        "        for step in range(0, inputs.size(2) - 1):\n",
        "            # If step <= burn_in_steps, the input is the true input\n",
        "            # Otherwise it's the output of the previous step\n",
        "            if step <= burn_in_steps:\n",
        "                ins = ...\n",
        "            else:\n",
        "                ins = ...\n",
        "\n",
        "            output, hidden = self.step(ins, hidden)\n",
        "            outputs.append(output)\n",
        "\n",
        "        outputs = torch.stack(outputs, dim=2)\n",
        "\n",
        "        return outputs"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VesTtGqR78Aa",
        "colab_type": "text"
      },
      "source": [
        "### Eval\n",
        "\n",
        "We use the negative log likelihood as a measure of performance, as in the paper and provided in the github repo. The evaluation code is given.\n",
        "\n",
        "Fill out the training loop to use the RecurrentBaseline implemented above to minimize this negative log likelihood. Also compute the negative log-likelihood of the predictions after burn-in. \n",
        "\n",
        "Notice that the negative log-likelihood is lower after the burn-in than before. Why is this surprising ? Why is this happening ?\n",
        "\n",
        "Comment on the speed of this baseline. How involved would it be to make it faster ?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C4NYuDOn78Ac",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from tqdm import tqdm_notebook, tnrange\n",
        "\n",
        "model = RecurrentBaseline(4, 256, 5, 2, 0.2).cuda()\n",
        "optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
        "\n",
        "def nll_gaussian(preds, target, variance):\n",
        "    neg_log_p = ((preds - target) ** 2 / (2 * variance))\n",
        "    return neg_log_p.sum() / (target.size(0) * target.size(1))\n",
        "\n",
        "# One epoch of training\n",
        "def train(epoch):\n",
        "    t = time.time()\n",
        "    loss = []\n",
        "    pred_loss = []\n",
        "    model.train()\n",
        "    with tqdm_notebook(loaders['train'], desc=f'training') as t:\n",
        "        for data, relations in t:\n",
        "            data, relations = data.cuda(), relations.cuda()\n",
        "            # compute the predicted trajectory with burn_in = timesteps - prediction_steps\n",
        "            output = ...\n",
        "\n",
        "            # output_t is data_{t+1}. Select a time-shifted slice of target to make loss computations easier\n",
        "            target = ...\n",
        "            \n",
        "            # Compute the training loss and nll on steps after burn in\n",
        "            l = ...\n",
        "            predicted_loss = ...\n",
        "\n",
        "            # take a gradient step\n",
        "\n",
        "\n",
        "            loss.append(l.item())\n",
        "            pred_loss.append(predicted_loss.item())\n",
        "            \n",
        "            t.set_postfix(loss=l.item(), pred=predicted_loss.item())\n",
        "\n",
        "    # return average loss and pred_loss\n",
        "    return ... \n",
        "\n",
        "# Train model\n",
        "for epoch in tnrange(1):\n",
        "    nll, pred = train(epoch)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6UKJVvOz78Ae",
        "colab_type": "text"
      },
      "source": [
        "## Neural Relational Inference model\n",
        "\n",
        "There are two parts to the model : an encoder that estimates the relation matrix and a decoder that produces a sequence given an estimation of the relation matrix.\n",
        "\n",
        "### The encoder\n",
        "The equations for the encoder in the paper are :\n",
        "$${\\bf h}^1_j = f_{emb}({\\bf x}_j)$$\n",
        "$$v\\rightarrow e:\\quad {\\bf h}^1_{(i,j)} = f_e^1([{\\bf h}^1_i, {\\bf h}^1_j])$$\n",
        "$$e\\rightarrow v:\\quad{\\bf h}^2_j = f_v^1\\big(\\sum_{i \\neq j}{\\bf h}^1_{(i,j)}\\big)$$\n",
        "$$v\\rightarrow e:\\quad{\\bf h}^2_{(i,j)} = f_e^2([{\\bf h}_i^2, {\\bf h}_j^2])$$\n",
        "\n",
        "Finally, we do a logistic regression on ${\\bf h}^2_{(i,j)}$ to obtain the probabilities of edge / non-edge.\n",
        "\n",
        "We will represent all functions as multi-layer perceptrons."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z-shQxHm78Af",
        "colab_type": "text"
      },
      "source": [
        "The implementation challenge in the encoder is to efficiently concatenate the $h_i$, $h_j$. We do this using ```index_select(input, dim, indices)```.\n",
        "\n",
        "Given an input of dimension $atoms \\times d$, create two index tensors such that for\n",
        "```python\n",
        "    x = torch.index_select(input, 0, id1)\n",
        "    y = torch.index_select(input, 0, id2)\n",
        "```\n",
        "We have $x_{i*atoms + j} = input_i$ and $x_{i*atoms + j} = input_j$. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v1tCfOH578Af",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "n_atoms = 5\n",
        "d = 2\n",
        "features = torch.FloatTensor([[i] * d for i in range(n_atoms)])\n",
        "\n",
        "id1 = ...\n",
        "id2 = ..."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6G__uh3h78Ai",
        "colab_type": "text"
      },
      "source": [
        "We can now easily write the concatenation in the $v\\rightarrow e$ step:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LniowJXr78Aj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def v_to_e(x, id1, id2):\n",
        "    return torch.cat([\n",
        "        torch.index_select(x, 0, id1),\n",
        "        torch.index_select(x, 0, id2),\n",
        "    ], 1)\n",
        "\n",
        "print(v_to_e(features, id1, id2))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U7Gtj1bh78Al",
        "colab_type": "text"
      },
      "source": [
        "Read and understand this implementation of the aggregation in the $e \\rightarrow v$ step.\n",
        "\n",
        "What differs in this implementation from the equation above ?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HUz4cXj178Al",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "aggregator = torch.FloatTensor([\n",
        "    [1./n_atoms if row * n_atoms <= col < (row + 1) * n_atoms else 0 for col in range(n_atoms * n_atoms)]\n",
        "    for row in range(n_atoms)\n",
        "])\n",
        "\n",
        "def e_to_v(x, matrix):\n",
        "    return matrix @ x"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bZbmQdM078An",
        "colab_type": "text"
      },
      "source": [
        "In order to remove self-loops, we will use another index select. Given a tensor resulting from the v\\_to\\_e function above, write a function using index select that returns all edges except self edges. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mMbdWVFz78An",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def remove_self_edges(features):\n",
        "    ...\n",
        "\n",
        "print(remove_self_edges(v_to_e(features, id1, id2)))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RVQ6W1CYVYH9",
        "colab_type": "text"
      },
      "source": [
        "#### MLP Encoder\n",
        "\n",
        "Let's hide this all inside a neat MLPEncoder class. An MLP class is given that will be used for all function approximation within this encoder.\n",
        "\n",
        "`ids_and_agg` returns (id1, id2, agg, id3) where id1 and id2 are used for tiling the features in the v->e step. agg is the same aggregator as before except if no_self_edges is true, in which case it doesn't aggregate the self-edges. Finally, id3 is the index used for removing self-edges from the v->e step.\n",
        "\n",
        "Complete the MLPEncoder class below, to match the equations of the encoder above.\n",
        "\n",
        "**Note** : Compared to the equations above, we add a skip connection from the output of equation 2 to the input of equation 4 (ie, $f_e^2$ acts on the concatenation of $[{\\bf h}_i^2, {\\bf h}_j^2]$ and ${\\bf h^1_{(i,j)}}$)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "glTzstK578Ao",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch.nn.functional as F\n",
        "\n",
        "class MLP(nn.Module):\n",
        "    \"\"\"Two-layer fully-connected ELU net with batch norm.\"\"\"\n",
        "\n",
        "    def __init__(self, n_in, n_hid, n_out, do_prob=0.):\n",
        "        super(MLP, self).__init__()\n",
        "        self.fc1 = nn.Linear(n_in, n_hid)\n",
        "        self.fc2 = nn.Linear(n_hid, n_out)\n",
        "        self.bn = nn.BatchNorm1d(n_out)\n",
        "        self.dropout_prob = do_prob\n",
        "\n",
        "    def batch_norm(self, inputs):\n",
        "        x = inputs.view(inputs.size(0) * inputs.size(1), -1)\n",
        "        x = self.bn(x)\n",
        "        return x.view(inputs.size(0), inputs.size(1), -1)\n",
        "\n",
        "    def forward(self, inputs):\n",
        "        # Input shape: [num_sims, num_things, num_features]\n",
        "        x = F.elu(self.fc1(inputs))\n",
        "        x = F.dropout(x, self.dropout_prob, training=self.training)\n",
        "        x = F.elu(self.fc2(x))\n",
        "        return self.batch_norm(x)\n",
        "\n",
        "\n",
        "\n",
        "def ids_and_agg(n_atoms, no_self_edges=False):\n",
        "    n_for_agg = (n_atoms - 1) if no_self_edges else n_atoms\n",
        "    return (\n",
        "        torch.cuda.LongTensor(sum([[i] * n_atoms for i in range(n_atoms)], [])),\n",
        "        torch.cuda.LongTensor(sum([list(range(n_atoms)) for i in range(n_atoms)], [])),\n",
        "        torch.cuda.FloatTensor([\n",
        "            [1. / n_for_agg if row * n_for_agg <= col < (row + 1) * n_for_agg else 0\n",
        "             for col in range(n_for_agg * n_atoms)]\n",
        "            for row in range(n_atoms)\n",
        "        ]),\n",
        "        torch.cuda.LongTensor([\n",
        "            i for i in range(n_atoms * n_atoms)\n",
        "            if i not in set([j*n_atoms + j for j in range(n_atoms)])\n",
        "        ])\n",
        "    )\n",
        "\n",
        "class MLPEncoder(nn.Module):\n",
        "    def __init__(self, n_atoms, n_in, n_hid, n_out, do_prob=0.):\n",
        "        \"\"\"\n",
        "        Given an input of shape [batch_size, num_atoms, num_timesteps, num_dims],\n",
        "        output a tensor of shape [num_atoms * (num_atoms - 1), n_out] with class logits for each atom-atom interaction edge.\n",
        "        \n",
        "        :param n_atoms number of atoms in the simulation\n",
        "        :param n_in total number of features for one atom, ie, n_dim * n_timesteps\n",
        "        :param n_hid size of the hidden layer\n",
        "        :param n_out number of classes to output for the encoder (ie, edge types)\n",
        "        \"\"\"\n",
        "        super(MLPEncoder, self).__init__()\n",
        "        # mlp1 is f_emb\n",
        "        self.mlp1 = ...\n",
        "        # mlp2 is f_e^1 \n",
        "        self.mlp2 = ...\n",
        "        # mlp3 is f_v^1\n",
        "        self.mlp3 = ...\n",
        "        # mlp4 is f_e^2\n",
        "        self.mlp4 = ...\n",
        "        # fc_out to output the logits of each interaction class\n",
        "        self.fc_out = ...\n",
        "        \n",
        "        self.id1, self.id2, self.aggregator, self.id3 = ids_and_agg(n_atoms)\n",
        "        \n",
        "    def tile(self, x):\n",
        "        # v -> e\n",
        "        return ...\n",
        "\n",
        "    def aggregate(self, x):\n",
        "        # e -> v\n",
        "        return ...\n",
        "\n",
        "    def forward(self, inputs):\n",
        "        # Input shape: [num_sims, num_atoms, num_timesteps, num_dims]\n",
        "        x = inputs.view(inputs.size(0), inputs.size(1), -1)\n",
        "        # New shape: [num_sims, num_atoms, num_timesteps*num_dims]\n",
        "        \n",
        "        x = ...                  # eq 1\n",
        "        x_skip = ...             # eq 2\n",
        "        x = ...                  # eq 3\n",
        "        x = ...                  # eq 4\n",
        "        \n",
        "        logits = ...\n",
        "        return ...       # remove self-edges"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UPKTVV3l78Aq",
        "colab_type": "text"
      },
      "source": [
        "Before adding the decoder, let's verify that this MLP Encoder can at least overfit the training data.\n",
        "\n",
        "The function edge_accuracy computes the argmax of the logits along the last dimension of `preds` (either 0 or 1 for two classes), and returns the proportion of correct classifications compared to `target`.\n",
        "\n",
        "Complete the following code to train your encoder to match the relations in the train set. After two epochs, you should reach train accuracies above 0.95"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-UVugHfq78Aq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from tqdm import tqdm_notebook, tnrange\n",
        "from torch.nn import BCEWithLogitsLoss\n",
        "n_atoms = 5\n",
        "model = MLPEncoder(n_atoms, int(4 * timesteps), 256, 2, 0.2).cuda()\n",
        "optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
        "\n",
        "def edge_accuracy(preds, target):\n",
        "    \"\"\"\n",
        "    :param preds: edge logits\n",
        "    :param target: ground truth\n",
        "    :return: precision of the prediction\n",
        "    \"\"\"\n",
        "    _, preds = preds.max(-1)\n",
        "    correct = preds.float().data.eq(\n",
        "        target.float().data.view_as(preds)).cpu().sum()\n",
        "    return np.float(correct) / (target.size(0) * target.size(1))\n",
        "\n",
        "# One epoch of training\n",
        "def train(epoch):\n",
        "    t = time.time()\n",
        "    loss_train = []\n",
        "    acc_train = []\n",
        "    # pick the right loss\n",
        "    loss = ...\n",
        "    model.train()\n",
        "    with tqdm_notebook(loaders['train'], 'training') as t:\n",
        "        for data, relations in t:\n",
        "            data, relations = data.cuda(), relations.cuda()\n",
        "            # compute the loss and take a gradient step\n",
        "            \n",
        "            \n",
        "            # compute the accuracy\n",
        "            edge_acc = ...\n",
        "            \n",
        "            loss_train.append(l.item())\n",
        "            acc_train.append(edge_acc)\n",
        "            \n",
        "            t.set_postfix(loss=l.item(), acc=edge_acc)\n",
        "\n",
        "    return np.mean(loss_train), np.mean(acc_train)\n",
        "\n",
        "# Train model\n",
        "t_total = time.time()\n",
        "best_epoch = 0\n",
        "for epoch in tnrange(2):\n",
        "    train_loss, train_acc = train(epoch)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mNO-aL2E78As",
        "colab_type": "text"
      },
      "source": [
        "### The decoder\n",
        "\n",
        "The equations of the recurrent decoder are :\n",
        "$$\n",
        "\\begin{aligned} v \\rightarrow e : \\tilde{\\mathbf{h}}_{(i, j)}^{t} &=\\sum_{k} z_{i j, k} \\tilde{f}_{e}^{k}\\left(\\left[\\tilde{\\mathbf{h}}_{i}^{t}, \\tilde{\\mathbf{h}}_{j}^{t}\\right]\\right) \\\\ e \\rightarrow v : \\operatorname{MSG}_{j}^{t} &=\\sum_{i \\neq j} \\tilde{\\mathbf{h}}_{(i, j)}^{t} \\\\ \\tilde{\\mathbf{h}}_{j}^{t+1} &=\\operatorname{GRU}\\left(\\left[\\operatorname{MSG}_{j}^{t}, \\mathbf{x}_{j}^{t}\\right], \\tilde{\\mathbf{h}}_{j}^{t}\\right) \\\\ \\boldsymbol{\\mu}_{j}^{t+1} &=\\mathbf{x}_{j}^{t}+f_{\\text { out }}\\left(\\tilde{\\mathbf{h}}_{j}^{t+1}\\right) \\\\ p\\left(\\mathbf{x}^{t+1} | \\mathbf{x}^{t}, \\mathbf{z}\\right) &=\\mathcal{N}\\left(\\boldsymbol{\\mu}^{t+1}, \\sigma^{2} \\mathbf{I}\\right) \\end{aligned}\n",
        "$$\n",
        "\n",
        "We will use only one edge type (ie 2 classes, _on_ or _off_) for simplicity.\n",
        "\n",
        "The last equation describes the likelihood of one point, by optimizing the nll_gaussian, we ensure that this probability is high on the train set.\n",
        "\n",
        "Complete the following skeleton code for the `RNNDecoder` module."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3lHevAqY78At",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class RNNDecoder(nn.Module):\n",
        "    \"\"\"Recurrent decoder module.\"\"\"\n",
        "\n",
        "    def __init__(self, n_dims, n_hid, do_prob=0.):\n",
        "        super(RNNDecoder, self).__init__()\n",
        "        \n",
        "        # Linear, Tanh, Dropout, Linear, Tanh\n",
        "        self.edge_mlp = ...\n",
        "        \n",
        "        # GruCell\n",
        "        self.gru = ...\n",
        "        \n",
        "        # Linear, ReLU, Linear, ReLU, Linear\n",
        "        self.decoder = ...\n",
        "        self.n_dims = n_dims\n",
        "        self.n_hid = n_hid\n",
        "        \n",
        "        self.id1, self.id2, self.aggregator, self.id3 = ids_and_agg(n_atoms, True)\n",
        "        \n",
        "    def step(self, inputs, edges, hidden):\n",
        "        \"\"\"\n",
        "        Compute one step of the decoder\n",
        "        :param inputs: a tensor of shape [batch_size x n_atoms x dims]\n",
        "        :param edges: a tensor of shape  [batch_size x n_edges x edge_type]\n",
        "        :param hidden: a tensor of shape [batch_size x n_atoms x hidden_size]\n",
        "        \"\"\"\n",
        "        # concatenate the features (equation 1)\n",
        "        hidden_state = ...\n",
        "        # remove the self edges\n",
        "        hidden_without_self = ...\n",
        "        # multiply by the z, the probability that the edge is active\n",
        "        transformed = ...\n",
        "        # aggregate\n",
        "        hidden_state = ...\n",
        "        \n",
        "        # compute the next_hidden state with the gru\n",
        "        next_hidden = ...\n",
        "        \n",
        "        # compute the output\n",
        "        output = ...\n",
        "\n",
        "        return output, next_hidden\n",
        "\n",
        "    def forward(self, data, edges, burn_in_steps=1):\n",
        "\n",
        "        inputs = data.transpose(1, 2).contiguous()\n",
        "\n",
        "        time_steps = inputs.size(1)\n",
        "\n",
        "        hidden = torch.zeros(inputs.size(0), inputs.size(2), self.n_hid).cuda()\n",
        "        pred_all = []\n",
        "\n",
        "        for step in range(inputs.size(1) - 1):\n",
        "            # similar as for the LSTM baseline\n",
        "            if step <= burn_in_steps:\n",
        "                ins = ...\n",
        "            else:\n",
        "                ins = ...\n",
        "\n",
        "            pred, hidden = ...\n",
        "            pred_all.append(pred)\n",
        "\n",
        "        preds = torch.stack(pred_all, dim=1)\n",
        "\n",
        "        return preds.transpose(1, 2)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eOKPsITI78Au",
        "colab_type": "text"
      },
      "source": [
        "### Minimizing the negative ELBO\n",
        "\n",
        "Our goal is to minimize the following loss :\n",
        "$$\n",
        "\\mathcal{L}=-\\mathbb{E}_{q_{\\phi}(\\mathbf{z} | \\mathbf{x})}\\left[\\log p_{\\theta}(\\mathbf{x} | \\mathbf{z})\\right]+\\mathrm{KL}\\left[q_{\\phi}(\\mathbf{z} | \\mathbf{x}) \\| p_{\\theta}(\\mathbf{z})\\right]\n",
        "$$\n",
        "\n",
        "Following the paper, we first encode the trajectory to obtain $q_\\phi(z|x)$. Instead of sampling from the logits, we use the gumbel_softmax to obtain backpropable edge probabilities.\n",
        "\n",
        "With these edge probabilities $z$, we use our rnn decoder and the nll_gaussian to compute the likelihood of $x$ given $z$.\n",
        "\n",
        "Complete the following code, which implements the method described in the Neural Relational Inference paper."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CKelTu6-78Av",
        "colab_type": "text"
      },
      "source": [
        "### Putting it all together\n",
        "\n",
        "What do you notice about the edge accuracy during training ? Why is this surprising ? "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E-qiO0eU78Av",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from torch.nn.functional import gumbel_softmax, softmax\n",
        "\n",
        "def kl_categorical_uniform(preds, num_atoms):\n",
        "    \"\"\"\n",
        "    :param preds tensor of shape [n_edge x 2] containing the probabilities of no edge and edge in columns 0 and 1.\n",
        "    \"\"\"\n",
        "    kl_div = preds * torch.log(preds + 1e-16)\n",
        "    return kl_div.sum() / (num_atoms * preds.size(0))\n",
        "\n",
        "def train(data_loader, optimizer, encoder, decoder):\n",
        "    loss_train = []\n",
        "\n",
        "    encoder.train()\n",
        "    decoder.train()\n",
        "    with tqdm_notebook(data_loader, desc=f'training') as t:\n",
        "        for data, relations in t:\n",
        "            data, relations = data.cuda(), relations.cuda()     \n",
        "\n",
        "            # Encode\n",
        "            logits = ...\n",
        "\n",
        "            # Compute edges with soft gumbel_softmax\n",
        "            edges = gumbel_softmax(\n",
        "                logits.view(-1, 2), tau=temp, hard=False\n",
        "            ).view(logits.shape)\n",
        "\n",
        "            # Decode using the edge weights\n",
        "            output = ...\n",
        "\n",
        "            nll = ...\n",
        "            kl = ...\n",
        "            l = nll + kl\n",
        "            \n",
        "            # Compute edge accuracy\n",
        "            edge_acc = ...\n",
        "\n",
        "            # Gradient step\n",
        "\n",
        "            loss_train.append(l.item())\n",
        "            t.set_postfix(loss=l.item(), nll = nll.item(), kl = kl.item(), acc=edge_acc)\n",
        "        \n",
        "    return np.mean(loss_train)\n",
        "\n",
        "dropout = 0\n",
        "n_dims = 4\n",
        "hidden = 256\n",
        "\n",
        "encoder = MLPEncoder(num_atoms, int(n_dims * timesteps), hidden, 2, dropout).cuda()\n",
        "decoder = RNNDecoder(n_dims, hidden, dropout).cuda()\n",
        "\n",
        "optimizer = optim.Adam(\n",
        "    list(encoder.parameters()) + list(decoder.parameters()),\n",
        "    lr=learning_rate\n",
        ")\n",
        "\n",
        "for e in range(10):\n",
        "    loss = train(loaders['train'], optimizer, encoder, decoder)\n",
        "    print(f\"{loss}\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ljPj_ZMCeqo_",
        "colab_type": "text"
      },
      "source": [
        "### Bonus\n",
        "On the test set, plot the predicted trajectories for a few inputs. Differentiate the burn-in phase and pure prediction phase. Also plot the true trajectories and compare."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UggrMKQCe2wI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}