chapter19_graph-neural-networks/Graph-Neural-Networks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deep Learning on Graphs with Message Passing Neural Networks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point in the Straight Dope, we've seen a wide variety of different types of data fed as input to our models.\n",
    "\n",
    "We started with linear regression models and MLPs, which take simple, 1-dimensional vectors of real numbers as input.  Then we met CNNs, which take images represented as 3-dimensional tensors as input.  Next we saw how RNNs can take sequence data, like time-series or natural-language sentences, or really anything we can represent as a sequence of tensors, as input.  And we even saw how to consume tree-structured data, like a parse tree of a natural-language sentence, using a Tree LSTM.\n",
    "\n",
    "In this chapter we'll see how to build models to handle yet another type of data: graph-structured data. We'll learn how to build Message Passing Neural Networks (MPNNs), which are a class of deep model that can take arbitrary graphs as input.\n",
    "\n",
    "**Wait, \"graphs\"?**\n",
    "\n",
    "When I say \"graph\", I mean that word the way a mathematician means it.  [Wikipedia explains the concept well][1], if you're not familiar.  Going forward I'll assume we're familiar with graph-lingo like \"directed edge\" and \"adjacency matrix\", so take a gander at that link if you need to.\n",
    "\n",
    "**So what exactly does \"taking graphs as input\" mean?**\n",
    "\n",
    "Good question!  Reading papers or blogs about this topic can be confusing, since (at least) two distinct learning scenarios both go by the name \"learning on graphs\":\n",
    "1. *We're trying to learn a model whose inputs are arbitrary graphs.*  Our dataset consists of (graph, label) pairs.  E.g. predicting the pharmacological activity of a molecule based on how its atoms are connected.\n",
    "    \n",
    "2. *We're trying to learn a model whose inputs are vertices in some graph.*  Our dataset is one big graph whose vertices are datapoints with edges between them, some labeled, some unlabeled.  E.g. predicting the impact factor of an article given a bag-of-words representation of the article and edges connecting it to its references.\n",
    "\n",
    "In this chapter, we're focusing only on scenario 1, but MPNNs can be used for scenario 2 as well.\n",
    "\n",
    "**Aren't sequences and trees just special types of graphs?  We already know models that handle those. (RNNs and Tree-RNNs.)**\n",
    "\n",
    "Yes they are!  In fact you can (and people do) even think of images as graphs where each pixel is a vertex with edges to all its adjacent pixels.  But MPNNs can operate on *any* type of graph: directed or not, cyclic or not, etc.  Be careful though: MPNNs likely won't perform as well on sequences, trees, or images as models designed specifically for these data types will.\n",
    "\n",
    "**But can't you basically represent anything as a graph if you try hard enough?**\n",
    "\n",
    "Yeah, that's partly why graphs are ubiquitous in math and computer science.  They're a super general concept.  \n",
    "\n",
    "This generality should make us veeeery suspicious that deep learning on graphs won't work as consistently well as, say, deep learning on real-world images does.  If it did, we could use deep networks to reason about nearly anything, and that would smell like a free lunch.\n",
    "\n",
    "But MPNNs are still worth learning about.  They're the best tool we have at the moment for understanding graph-structured data, and they're a hot area of research.\n",
    "\n",
    "[1]:https://en.wikipedia.org/wiki/Graph_(discrete_mathematics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Message Passing Neural Networks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Message Passing Neural Networks were introduced in [this paper](https://arxiv.org/pdf/1704.01212.pdf).  MPNNs are actually a family of models rather than a specific implementation, like how RNNs are a general model family, one implementation of which is an LSTM.  We'll first go over the general MPNN idea and then build a specific implementation.\n",
    "\n",
    "### The Setup\n",
    "\n",
    "We've got a dataset of `(graph, label)` pairs.  In each graph, each vertex $v$ has associated features $x_v$, and each edge has features $e_{vw}$.  For simplicity of explanation we'll assume each graph is undirected, but once you understand MPNNs it's easy to see how to extend them to directed graphs or multigraphs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Model\n",
    "\n",
    "The goal of an MPNN is to take in a `graph` and output the correct `label`.  They do this by the following procedure:\n",
    "1. Initialize a \"hidden state\" $h_v^0$ for each vertex $v$ in the graph as a function of the vertex's features: $$h_v^0 = \\text{init_hidden}(x_v).$$\n",
    "2. For each round $t$ out of $T$ total rounds:\n",
    "    3. Each vertex $v$ receives a \"message\" $m_v^{t+1}$, which is the sum of messages passed by $v$'s neighbors as functions of their current hidden states and the edge features: $$m_v^{t+1} = \\sum_{w \\in \\text{neighbors of }v} M_t(h_v^t, h_w^t, e_{vw}).$$\n",
    "    4. Each vertex $v$ updates its hidden state as a function of the message it received: $$h_v^{t+1} = U_t(h_v^t, m_v^{t+1}).$$\n",
    "5. The output is computed as the \"readout\" function of all the hidden states: $$\\hat{y} = R_t(\\{h_v^T \\vert v \\text{ is in the graph} \\}).$$\n",
    "\n",
    "Here's an base class for any type of MPNN that encapsulates this procedure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import mxnet as mx\n",
    "from mxnet import nd, autograd, gluon\n",
    "import sklearn.metrics as metrics\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "import math\n",
    "np.random.seed(1)\n",
    "mx.random.seed(1)\n",
    "ctx = mx.cpu() # NOTE: CHANGE THIS TO .gpu() IF YOU HAVE A GPU - MUCH TIME WILL BE SAVED!\n",
    "\n",
    "class MPNN(gluon.Block):\n",
    "    '''\n",
    "    General base class for all varieties of Message Passing Neural Network.\n",
    "    '''\n",
    "    def __init__(self, n_msg_pass_iters, *args, **kwargs):\n",
    "        super(MPNN, self).__init__(**kwargs)\n",
    "        self.n_msg_pass_iters = n_msg_pass_iters\n",
    "    \n",
    "    def init_hidden_states_and_edges(self, graph):\n",
    "        # Performs \"init_hidden\" from above and prepares adjacency information from the graph\n",
    "        # (This function is here so the model can be flexible about what format the graph is given to us in.)\n",
    "        raise NotImplementedError()\n",
    "    \n",
    "    def compute_messages(self, hidden_states, edges, t):\n",
    "        # Computes M_t from above and sums the messages\n",
    "        raise NotImplementedError()\n",
    "    \n",
    "    def update_hidden_states(self, hidden_states, messages, t):\n",
    "        # Performs U_t from above\n",
    "        raise NotImplementedError()\n",
    "    \n",
    "    def readout(self, hidden_states):\n",
    "        # Performs R_t from above\n",
    "        raise NotImplementedError()\n",
    "        \n",
    "    def forward(self, graph):\n",
    "        hidden_states, edges = self.init_hidden_states_and_edges(graph)\n",
    "        for t in range(self.n_msg_pass_iters):\n",
    "            messages = self.compute_messages(hidden_states, edges, t)\n",
    "            hidden_states = self.update_hidden_states(hidden_states, messages, t)\n",
    "            \n",
    "        return self.readout(hidden_states)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Different flavors of MPNN use different functions for $\\text{init_hidden}$, $M_t$, $U_t$, and $R_t$, and more often than not these functions are simpler than the fully general versions described above.  For example, in the GGSNN version of MPNN we'll discuss below, $M_t$ is the same function for each $t$, and it doesn't depend on the neighboring vertex's hidden state or any edge features."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Gated Graph Sequence Neural Networks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we've got the MPNN framework down, let's grab some real data and implement a particular type of MPNN, called a Gated Graph Sequence Neural Network (GGSNN), to learn on it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### An actual dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a demonstration task, we'll use the [Tox21 dataset][1].  The objective of this dataset is to take in the [chemical structure of a molecule][2], represented as an undirected graph with atoms as vertices and bonds as edges, and predict the toxicity of the molecule.  In particular, we'll try to predict whether a molecule might [activate a particular cellular response to pollutants in your body][3].\n",
    "\n",
    "To access the Tox21 dataset we use the DeepChem package, which you'll have to install either from its [website][4] or with this code:\n",
    "\n",
    "[1]:https://tripod.nih.gov/tox21/challenge/\n",
    "[2]:https://en.wikipedia.org/wiki/Structural_formula\n",
    "[3]:https://pubchem.ncbi.nlm.nih.gov/bioassay/743122#section=Top\n",
    "[4]:https://deepchem.io/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# WARNING: You may have to change this command to suit the python environment you're working in, e.g. if you're not using conda\n",
    "# You also may need to restart your jupyter notebook and/or kernel\n",
    "!conda install -y -c deepchem -c rdkit -c conda-forge -c omnia deepchem=2.0.0\n",
    "!conda update numpy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll load the data and convert it to graph format.  If you're not fluent in chemistry, don't worry about the details of the following preprocessing.  We're just transforming the data from a molecular format into the format we're used to seeing from above.\n",
    "\n",
    "What we'll end up with is a dataset of `(graph, label)` tuples where each `label` is a binary label (toxic or not), and each `graph` is an undirected graph represented as a vector of features for each vertex and an adjacency matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import deepchem as dc\n",
    "from deepchem.feat.mol_graphs import ConvMol\n",
    "\n",
    "tox21_tasks, tox21_datasets, transformers = dc.molnet.load_tox21(featurizer='GraphConv')\n",
    "train_mols, valid_mols, test_mols = tox21_datasets\n",
    "\n",
    "def molecules_to_dataset(molecules):\n",
    "    dataset = []\n",
    "    for ind, (mols, targets, _, _) in enumerate(molecules.iterbatches(1, deterministic=True)):\n",
    "        mol = mols[0]\n",
    "        target = targets[0]\n",
    "        adj_list = mol.get_adjacency_list()\n",
    "        if any(adj_list):\n",
    "            d = {}\n",
    "            # Grabbing \"NR-AhR\" endocrine system results, which is the 3rd assay in targets\n",
    "            d['label'] = target[2]\n",
    "            # Grabbing features for each atom, including the element type and some other chemical information\n",
    "            d['vertex_features'] = mol.get_atom_features()\n",
    "            # Grabbing the connectivity and converting it to sparse array\n",
    "            a = sp.sparse.dok_matrix((len(adj_list), len(adj_list)), dtype='float32')\n",
    "            for i in range(len(adj_list)):\n",
    "                for j in adj_list[i]:\n",
    "                    a[i,j] = 1\n",
    "            d['adj_mat'] = a.tocsr()\n",
    "            dataset.append(d)\n",
    "    return dataset\n",
    "\n",
    "train_dataset = molecules_to_dataset(train_mols)\n",
    "valid_dataset = molecules_to_dataset(valid_mols)\n",
    "test_dataset = molecules_to_dataset(test_mols)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The GGSNN model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll implement a Gated Graph Sequence Neural Network, introduced in [this paper][1], on this dataset.\n",
    "\n",
    "A GGSNN is an MPNN with the following customizations:\n",
    "1. Hidden states $h_v^0$ for each vertex are initialized with a single-layer MLP.\n",
    "2. The messages passed to $v$ by its neighbors are a simple matrix multiplication of each neighbor's hidden state: $$m_v^{t+1} = \\sum_{w \\in \\text{neighbors of }v} W_{\\texttt{msg_fxn}}h_w^t.$$\n",
    "3. Each vertex $v$ updates its hidden state to be the output of a [GRU cell](http://gluon.mxnet.io/chapter05_recurrent-neural-networks/gru-scratch.html) (a type of RNN cell) whose hidden state is the vertex's hidden state and whose input is the message the vertex received: $$h_v^{t+1} = \\text{GRU}(m_v^{t+1}, h_v^t).$$\n",
    "4. The \"readout\" function is this funny little beast: $$\\hat{y} = \\text{softmax}\\left(f_{\\text{out}}\\left(\\sum_{v} \\sigma\\left(f_1([h_v^T, h_v^0])\\right) \\odot f_2(h_v^T)\\right)\\right),$$ where the $f$s are MLPs, $\\sigma$ is the sigmoid function, and $\\odot$ is elementwise multiplication.  This acts like a sort of attention mechanism that depends on how much each vertex's hidden state changed during message passing.\n",
    "\n",
    "Here's an implementation of GGSNN that fills out the details of the MPNN base class from above:\n",
    "\n",
    "> *A key implementation note about what follows:* You'll notice below that the GGSNN is coded as though it takes in a single graph, rather than a minibatch of graphs as you might expect.  This is intentional.  We want to reserve the 0th/batch dimension of the tensors in our implementation to index over the vertices of the graph.  This makes the implementation more elegant, since MXNet operations are built to handle inputs that vary in size along the 0th dimension, and the number of vertices in each graph is usually different.\n",
    "\n",
    "> But of course, we DO want to process minibatches of data.  To do this, combine a minibatch of graphs into a single, large, disconnected graph, do all the message passing on this graph (no messages will get passed between minibatch elements, because their graphs are disconnected), and use the `batch_sizes` list to produce separate outputs for each graph in the minibatch in the `readout` step.\n",
    "\n",
    "[1]:https://arxiv.org/pdf/1511.05493.pdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class GGSNN(MPNN):\n",
    "    '''\n",
    "    GGSNN model for operating on the Tox21 dataset\n",
    "    '''\n",
    "    def __init__(self, vertex_feature_size, hidden_size, output_size, **kwargs):\n",
    "        super(GGSNN, self).__init__(**kwargs)\n",
    "        \n",
    "        # Initializing model components\n",
    "        with self.name_scope():\n",
    "            self.vertex_init = gluon.nn.Dense(hidden_size, in_units=vertex_feature_size)\n",
    "            self.message_fxn = gluon.nn.Dense(hidden_size, in_units=hidden_size, use_bias=False)\n",
    "            self.gru = gluon.rnn.GRUCell(hidden_size, input_size=hidden_size)\n",
    "            self.readout_1 = gluon.nn.Sequential()\n",
    "            with self.readout_1.name_scope():\n",
    "                self.readout_1.add(gluon.nn.Dense(hidden_size*2, activation='tanh'))\n",
    "                self.readout_1.add(gluon.nn.Dense(hidden_size))\n",
    "            self.readout_2 = gluon.nn.Sequential()\n",
    "            with self.readout_2.name_scope():\n",
    "                self.readout_2.add(gluon.nn.Dense(hidden_size, activation='tanh'))\n",
    "                self.readout_2.add(gluon.nn.Dense(hidden_size))\n",
    "            self.readout_final = gluon.nn.Dense(output_size, in_units=hidden_size)\n",
    "                \n",
    "    def init_hidden_states_and_edges(self, graph):\n",
    "        # vertex_features are a (num_vertices x num_features) NDarray \n",
    "        # edges is a (num_vertices x num_vertices) sparse NDarray\n",
    "        # batch_sizes is a list of the sizes of the graphs in the batch that were combined into the graph\n",
    "        vertex_features, edges, batch_sizes = graph\n",
    "        init_hidden_states = nd.tanh(self.vertex_init(vertex_features))\n",
    "        # Saving these for use in the readout function later - not every MPNN requires this, but GGSNNs do\n",
    "        self.init_hidden_states = init_hidden_states.copy()\n",
    "        self.batch_sizes = batch_sizes\n",
    "        return init_hidden_states, edges\n",
    "    \n",
    "    def compute_messages(self, hidden_states, edges, t):\n",
    "        passed_msgs = self.message_fxn(hidden_states)\n",
    "        summed_msgs = nd.sparse.dot(edges, passed_msgs)\n",
    "        return summed_msgs\n",
    "    \n",
    "    def update_hidden_states(self, hidden_states, messages, t):\n",
    "        hidden_states, _ = self.gru(messages, [hidden_states])\n",
    "        return hidden_states\n",
    "    \n",
    "    def readout(self, hidden_states):\n",
    "        readout_in_1 = nd.concat(hidden_states, self.init_hidden_states, dim=1)\n",
    "        readout_hid_1 = nd.sigmoid(self.readout_1(readout_in_1))\n",
    "        readout_hid_2 = self.readout_2(hidden_states)\n",
    "        readout_hid = readout_hid_1 * readout_hid_2\n",
    "        readout_attention = []\n",
    "        i = j = 0\n",
    "        while self.batch_sizes:\n",
    "            i = j\n",
    "            j += self.batch_sizes.pop(0)\n",
    "            readout_attention.append(nd.sum(readout_hid[i:j], axis=0, keepdims=True))\n",
    "        readout_attention = nd.concat(*readout_attention, dim=0)\n",
    "        return self.readout_final(readout_attention)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's train!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll create a new GGSNN instance and initialize our model parameters, loss function, and optimizer as usual:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "model = GGSNN(vertex_feature_size=75, hidden_size=100, output_size=2, n_msg_pass_iters=6)\n",
    "model.collect_params().initialize(mx.init.Normal(sigma=.01), ctx=ctx)\n",
    "softmax_cross_entropy = gluon.loss.SoftmaxCrossEntropyLoss()\n",
    "trainer = gluon.Trainer(model.collect_params(), 'adam', {'learning_rate': .0002})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we'll add a few helper functions to keep the training loop code clean:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def batchify_graphs(graphs):\n",
    "    '''\n",
    "    Args:\n",
    "        batch: List of graphs in {vertex_feature, adjacency_matrix, label} format\n",
    "        \n",
    "    Returns:\n",
    "        The combination of the input graphs into a big disconnected graph\n",
    "        The labels of each of the input graphs\n",
    "    '''\n",
    "    vertex_features = np.concatenate([g['vertex_features'] for g in graphs])\n",
    "    vertex_features = nd.array(vertex_features, dtype='float32', ctx=ctx)\n",
    "    adj_mat = sp.sparse.block_diag([g['adj_mat'] for g in graphs]).tocsr()\n",
    "    adj_mat = nd.sparse.csr_matrix((adj_mat.data, adj_mat.indices, adj_mat.indptr), dtype='float32', ctx=ctx)\n",
    "    batch_sizes = [g['vertex_features'].shape[0] for g in graphs]\n",
    "    labels = nd.array([g['label'] for g in graphs], dtype='float32', ctx=ctx)\n",
    "    return (vertex_features, adj_mat, batch_sizes), labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def evaluate_accuracy(dataset, model, n_batch):\n",
    "    '''\n",
    "    Measures the accuracy of the model on the provided dataset, in batches\n",
    "    '''\n",
    "    acc = mx.metric.Accuracy()\n",
    "    for i in range(0, math.ceil(len(dataset)/n_batch)):\n",
    "        data = dataset[n_batch*i:n_batch*(i+1)]\n",
    "        graph, label = batchify_graphs(data)\n",
    "        output = model(graph)\n",
    "        predictions = nd.argmax(output, axis=1)\n",
    "        acc.update(preds=predictions, labels=label)\n",
    "    return acc.get()[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def evaluate_roc_score(dataset, model, n_batch):\n",
    "    '''\n",
    "    Measures the area under the ROC curve of the model on the provided dataset, in batches\n",
    "    '''\n",
    "    pos_probs = []\n",
    "    labels = []\n",
    "    for i in range(0, math.ceil(len(dataset)/n_batch)):\n",
    "        data = dataset[n_batch*i:n_batch*(i+1)]\n",
    "        graph, label = batchify_graphs(data)\n",
    "        output = model(graph)\n",
    "        pos_probs.append(nd.softmax(output)[:,1])\n",
    "        labels.append(label)\n",
    "    labels = nd.concat(*labels, dim=0).asnumpy()\n",
    "    pos_probs = nd.concat(*pos_probs, dim=0).asnumpy()\n",
    "    return metrics.roc_auc_score(labels, pos_probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that the class balance in the dataset is heavily skewed toward the \"not toxic\" label:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "labels = np.array([i['label'] for i in train_dataset])\n",
    "print('Percentage of \"not toxic\" labels in training data = {}'.format(sum(labels == 0)/len(labels)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is why in the training loop below we're measuring the [ROC AUC](https://en.wikipedia.org/wiki/Receiver_operating_characteristic), in addition to just the accuracy.\n",
    "\n",
    "Now let's train!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "n_epochs = 30\n",
    "n_batch = 128\n",
    "\n",
    "for e in range(n_epochs):\n",
    "    cumulative_loss = 0\n",
    "    for i in range(0, math.ceil(len(train_dataset)/n_batch)):\n",
    "        data = train_dataset[n_batch*i:n_batch*(i+1)]\n",
    "        graph, label = batchify_graphs(data)\n",
    "        with autograd.record():\n",
    "            output = model(graph)\n",
    "            loss = softmax_cross_entropy(output, label)\n",
    "        loss.backward()\n",
    "        trainer.step(len(data))\n",
    "        cumulative_loss += nd.sum(loss).asscalar()\n",
    "    \n",
    "    valid_accuracy = evaluate_accuracy(valid_dataset, model, n_batch)\n",
    "    train_accuracy = evaluate_accuracy(train_dataset, model, n_batch)\n",
    "    valid_roc = evaluate_roc_score(valid_dataset, model, n_batch)\n",
    "    train_roc = evaluate_roc_score(train_dataset, model, n_batch)\n",
    "    print('Epoch {}. Loss: {}, \\n\\tTrain_acc {}, Valid_acc {}\\n\\tTrain_roc_auc {}, Valid_roc_auc {}'.format(\n",
    "            e, cumulative_loss/len(train_dataset), train_accuracy, valid_accuracy, train_roc, valid_roc))\n",
    "    \n",
    "print('Test Accuracy: {}'.format(evaluate_accuracy(test_dataset, model, n_batch)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alright!  Unsurprisingly, given the class imbalance, our accuracy didn't improve much; but our ROC score got much better, in line with the current state of the art on this dataset: see the physiology section [here](http://moleculenet.ai/latest-results).\n",
    "\n",
    "Now go forth and invent your own types of MPNNs!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}