Commit cb5cbf42 authored by Bharath Ramsundar's avatar Bharath Ramsundar
Browse files

Adding tg notebook

parent 74d47a73
Loading
Loading
Loading
Loading
+297 −0
Original line number Diff line number Diff line
%% Cell type:markdown id: tags:

In this notebook, we will explore the use of TensorGraph to create graph convolutional models with DeepChem. In particular, we will build a graph convolutional network on the Tox21 dataset.

Let's start with some basic imports.

%% Cell type:code id: tags:

``` python
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import numpy as np
import tensorflow as tf
import deepchem as dc
from deepchem.models.tensorgraph.models.graph_models import GraphConvTensorGraph
```

%% Output

    Warning: No xgboost installed on your system
    Attempting to run xgboost will throw runtime errors
    Warning: No pyGPGO.covfunc installed on your system
    Attempting to run pyGPGO.covfunc will throw runtime errors
    Warning: No pyGPGO.acquisition installed on your system
    Attempting to run pyGPGO.acquisition will throw runtime errors
    Warning: No pyGPGO.surrogates.GaussianProcess installed on your system
    Attempting to run pyGPGO.surrogates.GaussianProcess will throw runtime errors
    Warning: No pyGPGO.GPGO installed on your system
    Attempting to run pyGPGO.GPGO will throw runtime errors

%% Cell type:markdown id: tags:

Now, let's use MoleculeNet to load the Tox21 dataset. We need to make sure to process the data in a way that graph convolutional networks can use For that, we make sure to set the featurizer option to 'GraphConv'. The MoleculeNet call will return a training set, an validation set, and a test set for us to use. The call also returns `transformers`, a list of data transformations that were applied to preprocess the dataset. (Most deep networks are quite finicky and require a set of data transformations to ensure that training proceeds stably.)

%% Cell type:code id: tags:

``` python
# Load Tox21 dataset
tox21_tasks, tox21_datasets, transformers = dc.molnet.load_tox21(featurizer='GraphConv')
train_dataset, valid_dataset, test_dataset = tox21_datasets
```

%% Output

    Loading dataset from disk.
    Loading dataset from disk.
    Loading dataset from disk.

%% Cell type:markdown id: tags:

Let's now train a graph convolutional network on this dataset. DeepChem has the class `GraphConvTensorGraph` that wraps a standard graph convolutional architecture underneath the hood for user convenience. Let's instantiate an object of this class and train it on our dataset.

%% Cell type:code id: tags:

``` python
model = GraphConvTensorGraph(
    len(tox21_tasks), batch_size=50, mode='classification')
model.fit(train_dataset, nb_epoch=10)
```

%% Output

    /home/rbharath/anaconda2/lib/python2.7/site-packages/tensorflow/python/ops/gradients_impl.py:91: UserWarning: Converting sparse IndexedSlices to a dense Tensor of unknown shape. This may consume a large amount of memory.
      "Converting sparse IndexedSlices to a dense Tensor of unknown shape. "

    Starting epoch 0
    Starting epoch 1
    Starting epoch 2
    Starting epoch 3
    Starting epoch 4
    Starting epoch 5
    Starting epoch 6
    Starting epoch 7
    Ending global_step 999: Average loss 471.091
    Starting epoch 8
    Starting epoch 9
    Ending global_step 1290: Average loss 386.859
    TIMING: model fitting took 135.723 s

%% Cell type:markdown id: tags:

Let's try to evaluate the accuracy of the model we've trained. For this, we need to define a metric, a measure of model accuracy. `dc.metrics` holds a collection of metrics already. For this dataset, it is standard to use the ROC-AUC score, the area under the receiver operating characteristic curve (which measures the tradeoff between precision and recall). Luckily, the ROC-AUC score is already available in DeepChem.

To measure the accuracy of the model under this metric, we can use the convenience function `model.evaluate()`.

%% Cell type:code id: tags:

``` python
metric = dc.metrics.Metric(
    dc.metrics.roc_auc_score, np.mean, mode="classification")

print("Evaluating model")
train_scores = model.evaluate(train_dataset, [metric], transformers)
print("Training ROC-AUC Score: %f" % train_scores["mean-roc_auc_score"])
valid_scores = model.evaluate(valid_dataset, [metric], transformers)
print("Validation ROC-AUC Score: %f" % valid_scores["mean-roc_auc_score"])
```

%% Output

    Evaluating model
    computed_metrics: [0.90388790308447242, 0.95376302738091723, 0.90556048588854177, 0.916407286317377, 0.80895734251723961, 0.91104450761477085, 0.93863052288227733, 0.83045864558985349, 0.92538469587661831, 0.87919873473828425, 0.90509387643008332, 0.89722435415786572]
    Training ROC-AUC Score: 0.897968
    computed_metrics: [0.79571138800347718, 0.82797702220152147, 0.86280212842712833, 0.82468196274420702, 0.71598240796553037, 0.72798145806995362, 0.77118644067796605, 0.78516181076207647, 0.78875605177354013, 0.72683792815371762, 0.85786224821312551, 0.83015384615384624]
    Validation ROC-AUC Score: 0.792925

%% Cell type:markdown id: tags:

What's going on under the hood? Could we build `GraphConvTensorGraph` ourselves? Of course! The first step is to create a `TensorGraph` object. This object will hold the "computational graph" that defines the computation that a graph convolutional network will perform.

%% Cell type:code id: tags:

``` python
from deepchem.models.tensorgraph.tensor_graph import TensorGraph

tg = TensorGraph(use_queue=False)
```

%% Cell type:markdown id: tags:

Let's now define the inputs to our model. Conceptually, graph convolutions just requires a the structure of the molecule in question and a vector of features for every atom that describes the local chemical environment. However in practice, due to TensorFlow's limitations as a general programming environment, we have to have some auxiliary information as well preprocessed.

`atom_features` holds a feature vector of length 75 for each atom. The other feature inputs are required to support minibatching in TensorFlow. `degree_slice` is an indexing convenience that makes it easy to locate atoms from all molecules with a given degree. `membership` determines the membership of atoms in molecules (atom `i` belongs to molecule `membership[i]`). `deg_adjs` is a list that contains adjacency lists grouped by atom degree For more details, check out the [code](https://github.com/deepchem/deepchem/blob/master/deepchem/feat/mol_graphs.py).

To define feature inputs in `TensorGraph`, we use the `Feature` layer. Conceptually, a `TensorGraph` is a mathematical graph composed of layer objects. `Features` layers have to be the root nodes of the graph since they consitute inputs.

%% Cell type:code id: tags:

``` python
from deepchem.models.tensorgraph.layers import Feature

atom_features = Feature(shape=(None, 75))
degree_slice = Feature(shape=(None, 2), dtype=tf.int32)
membership = Feature(shape=(None,), dtype=tf.int32)

deg_adjs = []
for i in range(0, 10 + 1):
    deg_adj = Feature(shape=(None, i + 1), dtype=tf.int32)
    deg_adjs.append(deg_adj)
```

%% Cell type:markdown id: tags:

Let's now implement the body of the graph convolutional network. `TensorGraph` has a number of layers that encode various graph operations. Namely, the `GraphConv`, `GraphPool` and `GraphGather` layers. We will also apply standard neural network layers such as `Dense` and `BatchNorm`.

The layers we're adding effect a "feature transformation" that will create one vector for each molecule.

%% Cell type:code id: tags:

``` python
from deepchem.models.tensorgraph.layers import Dense, GraphConv, BatchNorm
from deepchem.models.tensorgraph.layers import GraphPool, GraphGather

batch_size = 50

gc1 = GraphConv(
    64,
    activation_fn=tf.nn.relu,
    in_layers=[atom_features, degree_slice, membership] + deg_adjs)
batch_norm1 = BatchNorm(in_layers=[gc1])
gp1 = GraphPool(in_layers=[batch_norm1, degree_slice, membership] + deg_adjs)
gc2 = GraphConv(
    64,
    activation_fn=tf.nn.relu,
    in_layers=[gp1, degree_slice, membership] + deg_adjs)
batch_norm2 = BatchNorm(in_layers=[gc2])
gp2 = GraphPool(in_layers=[batch_norm2, degree_slice, membership] + deg_adjs)
dense = Dense(out_channels=128, activation_fn=tf.nn.relu, in_layers=[gp2])
batch_norm3 = BatchNorm(in_layers=[dense])
readout = GraphGather(
    batch_size=batch_size,
    activation_fn=tf.nn.tanh,
    in_layers=[batch_norm3, degree_slice, membership] + deg_adjs)
```

%% Cell type:markdown id: tags:

Let's now make predictions from the `TensorGraph` model. Tox21 is a multitask dataset. That is, there are 12 different datasets grouped together, which share many common molecules, but with different outputs for each. As a result, we have to add a separate output layer for each task. We will use a `for` loop over the `tox21_tasks` list to make this happen. We need to add labels for each

We also have to define a loss for the model which tells the network the objective to minimize during training.

We have to tell `TensorGraph` which layers are outputs with `TensorGraph.add_output(layer)`. Similarly, we tell the network its loss with `TensorGraph.set_loss(loss)`.

%% Cell type:code id: tags:

``` python
from deepchem.models.tensorgraph.layers import Dense, SoftMax, \
    SoftMaxCrossEntropy, WeightedError, Concat
from deepchem.models.tensorgraph.layers import Label, Weights

costs = []
labels = []
for task in range(len(tox21_tasks)):
    classification = Dense(
        out_channels=2, activation_fn=None, in_layers=[readout])

    softmax = SoftMax(in_layers=[classification])
    tg.add_output(softmax)

    label = Label(shape=(None, 2))
    labels.append(label)
    cost = SoftMaxCrossEntropy(in_layers=[label, classification])
    costs.append(cost)
all_cost = Concat(in_layers=costs, axis=1)
weights = Weights(shape=(None, len(tox21_tasks)))
loss = WeightedError(in_layers=[all_cost, weights])
tg.set_loss(loss)
```

%% Cell type:markdown id: tags:

Now that we've successfully defined our graph convolutional model in `TensorGraph`, we need to train it. We can call `fit()`, but we need to make sure that each minibatch of data populates all four `Feature` objects that we've created. For this, we need to create a Python generator that given a batch of data generates a dictionary whose keys are the `Feature` layers and whose values are Numpy arrays we'd like to use for this step of training.

%% Cell type:code id: tags:

``` python
from deepchem.metrics import to_one_hot
from deepchem.feat.mol_graphs import ConvMol

def data_generator(dataset, epochs=1, predict=False, pad_batches=True):
  for epoch in range(epochs):
    if not predict:
        print('Starting epoch %i' % epoch)
    for ind, (X_b, y_b, w_b, ids_b) in enumerate(
        dataset.iterbatches(
            batch_size, pad_batches=True, deterministic=True)):
      d = {}
      for index, label in enumerate(labels):
        d[label] = to_one_hot(y_b[:, index])
      d[weights] = w_b
      multiConvMol = ConvMol.agglomerate_mols(X_b)
      d[atom_features] = multiConvMol.get_atom_features()
      d[degree_slice] = multiConvMol.deg_slice
      d[membership] = multiConvMol.membership
      for i in range(1, len(multiConvMol.get_deg_adjacency_lists())):
        d[deg_adjs[i - 1]] = multiConvMol.get_deg_adjacency_lists()[i]
      yield d
```

%% Cell type:markdown id: tags:

Now, we can train the model using `TensorGraph.fit_generator(generator)` which will use the generator we've defined to train the model.

%% Cell type:code id: tags:

``` python
tg.fit_generator(data_generator(train_dataset, epochs=10))
```

%% Output

    Starting epoch 0
    Starting epoch 1
    Starting epoch 2
    Starting epoch 3
    Starting epoch 4
    Starting epoch 5
    Starting epoch 6
    Starting epoch 7
    Ending global_step 999: Average loss 463.467
    Starting epoch 8
    Starting epoch 9
    Ending global_step 1290: Average loss 379.657
    TIMING: model fitting took 127.183 s

%% Cell type:markdown id: tags:

Now that we have trained our graph convolutional method, let's evaluate its performance. We again have to use our defined generator to evaluate model performance.

%% Cell type:code id: tags:

``` python
metric = dc.metrics.Metric(
    dc.metrics.roc_auc_score, np.mean, mode="classification")

print("Evaluating model")
train_scores = tg.evaluate_generator(data_generator(train_dataset, predict=True),
                                     [metric], labels=labels, weights=[weights])
print("Training ROC-AUC Score: %f" % train_scores["mean-roc_auc_score"])
valid_scores = tg.evaluate_generator(data_generator(valid_dataset, predict=True),
                                     [metric], labels=labels, weights=[weights])
print("Valid ROC-AUC Score: %f" % valid_scores["mean-roc_auc_score"])
```

%% Output

    Evaluating model
    computed_metrics: [0.90881144847528916, 0.96717149308328643, 0.9086312663123588, 0.91383160046160983, 0.80181888881149466, 0.91708002032436586, 0.94475492006679751, 0.83987084026121117, 0.92682785041187299, 0.87585845052450606, 0.90337844626637731, 0.9030230399539072]
    Training ROC-AUC Score: 0.900922
    computed_metrics: [0.78962619530570843, 0.77231796304921585, 0.86122384559884568, 0.81247160381644712, 0.67195767195767198, 0.74385447394296955, 0.76527991782229066, 0.81175856505646771, 0.85332477027961673, 0.76660401002506262, 0.85124849159936877, 0.80626923076923074]
    Valid ROC-AUC Score: 0.792161

%% Cell type:markdown id: tags:

Success! The model we've constructed behaves nearly identically to `GraphConvTensorGraph`. If you're looking to build your own custom models, you can follow the example we've provided here to do so. We hope to see exciting constructions from your end soon!
+5 −0
Original line number Diff line number Diff line
@@ -68,3 +68,8 @@ def test_quantum():
def test_pong():
  nb, errors = _notebook_read("pong.ipynb")
  assert errors == []


def test_graph_conv():
  nb, errors = _notebook_read("graph_convolutional_networks_for_tox21.ipynb")
  assert errors == []