bonus_property_prediction_tutorial.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "bonus_property_prediction_tutorial.ipynb",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "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/ajfisch/deeplearning_bootcamp_2020/blob/master/bonus_property_prediction_tutorial.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GA8LLAO9dchw",
        "colab_type": "text"
      },
      "source": [
        "# MLPs, RNNs, and CNNs for Property Prediction\n",
        "\n",
        "In this exercise, you'll implement a multi-layer perceptron (MPL), a recurrent neural network (RNN), and a convolutional neural network (CNN) to predict log p from SMILES strings.\n",
        "\n",
        "Let's get started!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kcCj4gBWdFeI",
        "colab_type": "text"
      },
      "source": [
        "# Preliminaries\n",
        "\n",
        "The next few sections will set up the necessary components of the exercise, including:\n",
        "\n",
        "\n",
        "1.   Installing PyTorch\n",
        "2.   Importing dependencies\n",
        "3.   Downloading and processing data\n",
        "4.   Defining training and evaluation procedures\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pg2T9UsXctbP",
        "colab_type": "text"
      },
      "source": [
        "## Download PyTorch"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eT_jLzv8do9z",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# http://pytorch.org/\n",
        "from os import path\n",
        "from wheel.pep425tags import get_abbr_impl, get_impl_ver, get_abi_tag\n",
        "\n",
        "platform = '{}{}-{}'.format(get_abbr_impl(), get_impl_ver(), get_abi_tag())\n",
        "accelerator = 'cu100' if path.exists('/opt/bin/nvidia-smi') else 'cpu'\n",
        "\n",
        "print(f'Platform = {platform}, Accelerator = {accelerator}')\n",
        "\n",
        "!pip install -q http://download.pytorch.org/whl/{accelerator}/torch-1.1.0-{platform}-linux_x86_64.whl\n",
        "!pip install -q torchvision\n",
        "\n",
        "import torch\n",
        "print(f'Torch version = {torch.__version__}')\n",
        "print(f'Cuda available = {torch.cuda.is_available()}')\n",
        "print(f'Cuda version = {torch.version.cuda}')\n",
        "print(f'Cuda device = {torch.cuda.get_device_name(0)}')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h7uKFRTBczh9",
        "colab_type": "text"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mcLJrU16dch0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import argparse\n",
        "from collections import Counter\n",
        "import csv\n",
        "import pickle\n",
        "import re\n",
        "\n",
        "import math\n",
        "import numpy as np\n",
        "from sklearn.metrics import mean_squared_error\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torch.utils.data import DataLoader, Dataset"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v7mgPaxPc1vx",
        "colab_type": "text"
      },
      "source": [
        "## Download and Process Data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DMJfz2f-JC9E",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!apt-get install wget\n",
        "!wget https://raw.githubusercontent.com/yala/introML_chem/master/lab1/data/chem/delaney_train.csv\n",
        "!wget https://raw.githubusercontent.com/yala/introML_chem/master/lab1/data/chem/delaney_val.csv\n",
        "!wget https://raw.githubusercontent.com/yala/introML_chem/master/lab1/data/chem/delaney_test.csv\n",
        "\n",
        "def get_data(split):\n",
        "    data_path = 'delaney_{}.csv'.format(split)\n",
        "    with open(data_path) as f:\n",
        "        data = csv.reader(f)\n",
        "    \n",
        "        # Skip header\n",
        "        next(data)\n",
        "        \n",
        "        # Get smiles and targets\n",
        "        smiles, Y = [], []\n",
        "        for row in data:\n",
        "            smiles.append(row[0])\n",
        "            Y.append(float(row[1]))\n",
        "    \n",
        "    return smiles, Y\n",
        "\n",
        "trainSmiles, trainY = get_data('train')\n",
        "devSmiles, devY = get_data('val')\n",
        "testSmiles, testY = get_data('test')\n",
        "\n",
        "allSmiles = trainSmiles + devSmiles + testSmiles\n",
        "\n",
        "print(f'Num Train = {len(trainSmiles):,}')\n",
        "print(f'Num Dev   = {len(devSmiles):,}')\n",
        "print(f'Num Test  = {len(testSmiles):,}')\n",
        "print()\n",
        "print(f'Example data point: smiles = {trainSmiles[0]}, logp = {trainY[0]}')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a_nrd1cJc_JM",
        "colab_type": "text"
      },
      "source": [
        "## Dataset Class"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4smEQXOddciA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class PropertyPredictionDataset(Dataset):\n",
        "    def __init__(self, X, Y):\n",
        "      self.X, self.Y = X, Y\n",
        "      assert len(X) == len(Y)\n",
        "\n",
        "    def __len__(self):\n",
        "       return len(self.X)\n",
        "\n",
        "    def __getitem__(self, i):\n",
        "      return np.array(self.X[i]), self.Y[i]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M9hvStGzdcjC",
        "colab_type": "text"
      },
      "source": [
        "## Model and Training Settings\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HgiOSHSrGFCU",
        "colab_type": "text"
      },
      "source": [
        "After building your MLP, RNN, and CNN below, return to this section and experiment with different values to see how they affect training and model performance."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n3Ht1mhhdcjE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "batch_size = 64\n",
        "epochs = 10\n",
        "lr = 1e-3\n",
        "weight_decay = 1e-4\n",
        "max_len = 100\n",
        "embedding_size = 300\n",
        "hidden_size = 300\n",
        "output_size = 1  # do not modify\n",
        "dropout = 0.6\n",
        "use_cuda = True"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eNLJoM1XPEz7",
        "colab_type": "text"
      },
      "source": [
        "## Utility Functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6Of3pYUgPE7R",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def param_count(model):\n",
        "    return sum(param.numel() for param in model.parameters() if param.requires_grad)\n",
        "  \n",
        "def rmse(targets, preds):\n",
        "    return math.sqrt(mean_squared_error(targets, preds))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gz35WFTOPDFo",
        "colab_type": "text"
      },
      "source": [
        "## Training Procedure"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wxIj0eWsdcjJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train_epoch(model, train_loader, optimizer, epoch):\n",
        "    model.train()  # Set the nn.Module to train mode. \n",
        "    total_loss = 0\n",
        "    total_rmse = 0\n",
        "    num_samples = len(train_loader.dataset)\n",
        "    num_batches = 0\n",
        "    for batch_idx, (data, target) in enumerate(train_loader):  # 1) get batch\n",
        "        # Adjust dimensions of target and cast to float\n",
        "        target = target.unsqueeze(1).float()\n",
        "      \n",
        "        # Move to cuda\n",
        "        if next(model.parameters()).is_cuda:\n",
        "            data, target = data.cuda(), target.cuda()\n",
        "      \n",
        "        # Reset gradient data to 0\n",
        "        optimizer.zero_grad()\n",
        "        \n",
        "        # Get prediction for batch\n",
        "        output = model(data)\n",
        "        \n",
        "        # 2) Compute loss\n",
        "        loss = F.mse_loss(output, target)\n",
        "        \n",
        "        # 3) Do backprop\n",
        "        loss.backward()\n",
        "        \n",
        "        # 4) Update model\n",
        "        optimizer.step()\n",
        "        \n",
        "        # Do book-keeping to track rmse and avg loss\n",
        "        total_loss += loss.detach()  # Don't keep computation graph \n",
        "        total_rmse += rmse(target.cpu().data.numpy(), output.cpu().data.numpy())\n",
        "        num_batches += 1\n",
        "\n",
        "    print(f'Train Epoch: {epoch} '\n",
        "          f'Loss: {total_loss / num_samples:.4f}, '\n",
        "          f'RMSE: {total_rmse / num_batches:.4f}')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xGZnlI58dcjN",
        "colab_type": "text"
      },
      "source": [
        "## Evaluation Procedure"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GPzqSuY3dcjO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def eval_epoch(model, test_loader, name):\n",
        "    model.eval()\n",
        "    test_loss = 0\n",
        "    test_rmse = 0\n",
        "    num_batches = 0\n",
        "    for data, target in test_loader:\n",
        "        target = target.unsqueeze(1).float()\n",
        "      \n",
        "        # Move to cuda\n",
        "        if next(model.parameters()).is_cuda:\n",
        "            data, target = data.cuda(), target.cuda()\n",
        "        \n",
        "        output = model(data)\n",
        "        \n",
        "        test_loss += F.mse_loss(output, target).item()  # sum up batch loss\n",
        "        test_rmse += rmse(target.cpu().data.numpy(), output.cpu().data.numpy())\n",
        "        num_batches += 1\n",
        "\n",
        "    test_loss /= len(test_loader.dataset)\n",
        "    test_rmse /= num_batches\n",
        "    print(f'\\n{name} set: '\n",
        "          f'Average loss: {test_loss:.4f}, '\n",
        "          f'RMSE: {test_rmse:.4f}')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6juQj-48uY5-",
        "colab_type": "text"
      },
      "source": [
        "# Character Embeddings for SMILES\n",
        "\n",
        "In the tutorial, we used word embeddings to encode each word in a sentence. Since a molecule's SMILES string is just a single \"word\" (i.e. a single sequence of characters), we'll instead use character embeddings, one for each character in the SMILES string. The embeddings will operate in essentially the same way as word embeddings."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jUt7IN02RJvc",
        "colab_type": "text"
      },
      "source": [
        "## Define Vocab and Character-to-Index Mapping"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BT93tF8iL1hr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Define vocab\n",
        "vocab = {char for smiles in allSmiles for char in smiles}\n",
        "\n",
        "print(f'Vocab = {vocab}')\n",
        "\n",
        "# Create word to index mapping\n",
        "padding_idx = 0\n",
        "char_to_index = {char: index + 1 for index, char in enumerate(vocab)}\n",
        "vocab_size = len(char_to_index) + 1\n",
        "\n",
        "print(f'Vocab size = {vocab_size:,}')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oDCNl7nKRHjt",
        "colab_type": "text"
      },
      "source": [
        "## Map Characters to Indices"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TEgwodfayJfn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "trainX = [[char_to_index[char] for char in smiles] for smiles in trainSmiles]\n",
        "devX =   [[char_to_index[char] for char in smiles] for smiles in devSmiles]\n",
        "testX =  [[char_to_index[char] for char in smiles] for smiles in testSmiles]\n",
        "\n",
        "print(f'Indices of first train SMILES = {trainX[0]}')\n",
        "print(f'Last five indices = {trainX[0][-5:]}')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i_OPFDOzQ8eX",
        "colab_type": "text"
      },
      "source": [
        "## Add Padding"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vMbCHXPavmyI",
        "colab_type": "text"
      },
      "source": [
        "Note: Since some SMILES are long, we've hard coded a maximum sentence length `max_len` in the Model and Training Settings section above."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S1K4Hg5DyimQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "trainX = [seq[:max_len] + [padding_idx] * (max_len - len(seq)) for seq in trainX]\n",
        "devX =   [seq[:max_len] + [padding_idx] * (max_len - len(seq)) for seq in devX]\n",
        "testX =  [seq[:max_len] + [padding_idx] * (max_len - len(seq)) for seq in testX]\n",
        "\n",
        "print(f'Indices of first train SMILES = {trainX[0]}')\n",
        "print(f'Last five indices = {trainX[0][-5:]}')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hZLu0JzjQ2LF",
        "colab_type": "text"
      },
      "source": [
        "## Build Dataset/DataLoader"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2qkAAPU60onl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Build Dataset\n",
        "train = PropertyPredictionDataset(trainX, trainY)\n",
        "dev = PropertyPredictionDataset(devX, devY)\n",
        "test = PropertyPredictionDataset(testX, testY)\n",
        "\n",
        "# Build DataLoader\n",
        "train_loader = DataLoader(train, batch_size=batch_size, shuffle=True)\n",
        "dev_loader = DataLoader(dev, batch_size=batch_size, shuffle=True)\n",
        "test_loader = DataLoader(test, batch_size=batch_size, shuffle=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZfXkqNFyqBsI",
        "colab_type": "text"
      },
      "source": [
        "# Multi-Layer Perceptron (MLP)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZePAuBdzEX_r",
        "colab_type": "text"
      },
      "source": [
        "Your first task is to build a multi-layer perceptron (MLP) to predict log p using a sum-of-embeddings approach. Replace all `raise NotImplementedError` lines below with your implementation. When you're ready, build the MLP and then train and test it."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JNQEIUKGqLM4",
        "colab_type": "text"
      },
      "source": [
        "## Define MLP"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rOD5dOZEqBAD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class MLP(nn.Module):\n",
        "    def __init__(self, vocab_size, padding_idx, embedding_size, hidden_size, output_size, dropout):\n",
        "        super(MLP, self).__init__()\n",
        "        \n",
        "        # Embedding layer\n",
        "        self.embed = nn.Embedding(vocab_size, embedding_size, padding_idx=padding_idx)\n",
        "        \n",
        "        # Fully connected layers\n",
        "        self.fc1 = nn.Linear(embedding_size, hidden_size)\n",
        "        self.fc2 = nn.Linear(hidden_size, hidden_size)\n",
        "        self.fc3 = nn.Linear(hidden_size, output_size)\n",
        "        \n",
        "        # Dropout (regularization)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):  # batch_size x seq_length\n",
        "        # Embed\n",
        "        embedded = self.embed(x)  # batch_size x seq_length x embedding_size\n",
        "        \n",
        "        # Sum embeddings\n",
        "        embedded = embedded.sum(dim=1)  # batch_size x embedding_size\n",
        "        \n",
        "        # MLP\n",
        "        hidden = self.dropout(F.relu(self.fc1(embedded)))  # batch_size x hidden_size\n",
        "        hidden = self.dropout(F.relu(self.fc2(hidden)))  # batch_size x hidden_size\n",
        "        logit = self.fc3(hidden)  # batch_size x output_size\n",
        "        \n",
        "        return logit"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IDTl_YiUqSMK",
        "colab_type": "text"
      },
      "source": [
        "## Build MLP"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "osUGOQsvqUzB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = MLP(vocab_size, padding_idx, embedding_size, hidden_size, output_size, dropout)\n",
        "\n",
        "print(model)\n",
        "print(f'Number of parameters = {param_count(model):,}')\n",
        "\n",
        "# Move to cuda\n",
        "if use_cuda and torch.cuda.is_available():\n",
        "    model = model.cuda()\n",
        "\n",
        "optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay) "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zCdqhdnVqV5k",
        "colab_type": "text"
      },
      "source": [
        "## Train MLP"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oXt9Oxz-qWzn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "for epoch in range(1, epochs + 1):\n",
        "    train_epoch(model, train_loader, optimizer, epoch)\n",
        "    eval_epoch(model,  dev_loader, \"Dev\")\n",
        "    print(\"---\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xE7LcxEPqW69",
        "colab_type": "text"
      },
      "source": [
        "## Test MLP"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gRTfwoZHqX5k",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "eval_epoch(model,  test_loader, \"Test\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZnsfujvesrvU",
        "colab_type": "text"
      },
      "source": [
        "# Recurrent Neural Network (RNN)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4UgmANcAMwjd",
        "colab_type": "text"
      },
      "source": [
        "Your next task is to build a recurrent neural network (RNN) to predict log p using a sequence of SMILES characteer embeddings. Replace all `raise NotImplementedError` lines below with your implementation. When you're ready, build the RNN and then train and test it."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zY_yh_IPA3hK",
        "colab_type": "text"
      },
      "source": [
        "## Define RNN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WdLyh493swKI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class RNN(nn.Module):\n",
        "    def __init__(self, vocab_size, padding_idx, embedding_size, hidden_size, output_size, dropout):\n",
        "        super(RNN, self).__init__()\n",
        "        \n",
        "        # Embedding layer\n",
        "        self.embed = nn.Embedding(vocab_size, embedding_size, padding_idx=padding_idx)\n",
        "        \n",
        "        # LSTM (RNN)\n",
        "        self.rnn = nn.LSTM(\n",
        "            input_size=embedding_size,\n",
        "            hidden_size=hidden_size,\n",
        "            batch_first=True\n",
        "        )\n",
        "        \n",
        "        # Fully connected layer\n",
        "        self.output = nn.Linear(hidden_size, output_size)\n",
        "        \n",
        "        # Dropout (regularization)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        \n",
        "    def forward(self, x):  # batch_size x seq_length\n",
        "        # Embed\n",
        "        embedded = self.embed(x)  # batch_size x seq_length x embedding_size\n",
        "      \n",
        "        # Run RNN\n",
        "        o, _ = self.rnn(embedded)  # batch_size x seq_length x hidden_size\n",
        "        \n",
        "        # Dropout\n",
        "        o = self.dropout(o)  # batch_size x seq_length x hidden_size\n",
        "        \n",
        "        # Max pooling across sequence\n",
        "        o, _ = torch.max(o, dim=1)    # batch_size x hidden_size\n",
        "        \n",
        "        # Output layer\n",
        "        logit = self.output(o)  # batch_size x output_size\n",
        "        \n",
        "        return logit"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QOxa8rZCA446",
        "colab_type": "text"
      },
      "source": [
        "## Build RNN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oTw0eOMW0nR1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = RNN(vocab_size, padding_idx, embedding_size, hidden_size, output_size, dropout)\n",
        "\n",
        "print(model)\n",
        "print(f'Number of parameters = {param_count(model):,}')\n",
        "\n",
        "# Move to cuda\n",
        "if use_cuda and torch.cuda.is_available():\n",
        "    model = model.cuda()\n",
        "    \n",
        "optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay) "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gaDSVZMXA7zO",
        "colab_type": "text"
      },
      "source": [
        "## Train RNN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DUR9k6471fFI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "for epoch in range(1, epochs + 1):\n",
        "    train_epoch(model, train_loader, optimizer, epoch)\n",
        "    eval_epoch(model,  dev_loader, \"Dev\")\n",
        "    print(\"---\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vAOPc283A9t2",
        "colab_type": "text"
      },
      "source": [
        "## Test RNN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tvS5w0e91fbH",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "eval_epoch(model,  test_loader, \"Test\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9sOI-H2e7sUi",
        "colab_type": "text"
      },
      "source": [
        "# Convolutional Neural Network (CNN)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gyfwzY5nM5T1",
        "colab_type": "text"
      },
      "source": [
        "Your next task is to build a recurrent neural network (CNN) to predict log p using a sequence of SMILES characteer embeddings. Replace all `raise NotImplementedError` lines below with your implementation. When you're ready, build the CNN and then train and test it."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ubU9eP8bBBFb",
        "colab_type": "text"
      },
      "source": [
        "## Define CNN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "klUA9Ugi7u5z",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class CNN(nn.Module):\n",
        "    def __init__(self, vocab_size, padding_idx, embedding_size, hidden_size, output_size, dropout):\n",
        "        super(CNN, self).__init__()\n",
        "        \n",
        "        # Embedding layer\n",
        "        self.embed = nn.Embedding(vocab_size, embedding_size, padding_idx=padding_idx)\n",
        "        \n",
        "        # Convolutional layers\n",
        "        self.conv1 = nn.Conv1d(in_channels=embedding_size, out_channels=hidden_size, kernel_size=3, padding=0)\n",
        "        self.conv2 = nn.Conv1d(in_channels=embedding_size, out_channels=hidden_size, kernel_size=5, padding=1)\n",
        "        self.conv3 = nn.Conv1d(in_channels=embedding_size, out_channels=hidden_size, kernel_size=7, padding=2)\n",
        "        \n",
        "        # Fully connect layer\n",
        "        self.output = nn.Linear(hidden_size, output_size)\n",
        "        \n",
        "        # Dropout (regularization)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        \n",
        "    def forward(self, x):  # batch_size x seq_length\n",
        "        # Embed\n",
        "        embedded = self.embed(x)  # batch_size x seq_length x embedding_size\n",
        "      \n",
        "        # Permute dimensions\n",
        "        embedded = embedded.permute(0, 2, 1)  # batch_size x embedding_size x seq_length\n",
        "        \n",
        "        # Convolutional layers\n",
        "        hidden_1 = self.dropout(F.relu(self.conv1(embedded)))  # batch_size x hidden_size x new_seq_length\n",
        "        hidden_2 = self.dropout(F.relu(self.conv2(embedded)))  # batch_size x hidden_size x new_seq_length\n",
        "        hidden_3 = self.dropout(F.relu(self.conv3(embedded)))  # batch_size x hidden_size x new_seq_length\n",
        "        \n",
        "        # Sum\n",
        "        hidden = hidden_1 + hidden_2 + hidden_3    # batch_size x hidden_size x new_seq_length\n",
        "        \n",
        "        # Max pooling across sequence\n",
        "        hidden, _ = hidden.max(dim=-1)    # batch_size x hidden_size\n",
        "        \n",
        "        # Output\n",
        "        logit = self.output(hidden)  # batch_size x output_size\n",
        "        \n",
        "        return logit"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XuaRaqdtBCrM",
        "colab_type": "text"
      },
      "source": [
        "## Build CNN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZD95lEXmAmGQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = CNN(vocab_size, padding_idx, embedding_size, hidden_size, output_size, dropout)\n",
        "\n",
        "print(model)\n",
        "print(f'Number of parameters = {param_count(model):,}')\n",
        "\n",
        "if use_cuda and torch.cuda.is_available():\n",
        "    model = model.cuda()\n",
        "\n",
        "optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay) "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sWwArpwcBEeo",
        "colab_type": "text"
      },
      "source": [
        "## Train CNN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eNN7Zrr2Ap9D",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "for epoch in range(1, epochs + 1):\n",
        "    train_epoch(model, train_loader, optimizer, epoch)\n",
        "    eval_epoch(model,  dev_loader, \"Dev\")\n",
        "    print(\"---\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3wROoo34BGBj",
        "colab_type": "text"
      },
      "source": [
        "## Test CNN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KxZmibBj-GDK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "eval_epoch(model,  test_loader, \"Test\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_BlOKfd0M_Kq",
        "colab_type": "text"
      },
      "source": [
        "# Improving Performance\n",
        "\n",
        "Now that you've built an MLP, RNN, and CNN, try your hand at maximizing the performance of each model. Experiment with different network architectures (e.g. different numbers of layers) and different model and training settings (see the Model and Training Settings section near the beginning). Which model performs best? How well does it do?"
      ]
    }
  ]
}