Commit c10f34fc authored by leswing's avatar leswing
Browse files

WIP

parent 0e2cb8d9
Loading
Loading
Loading
Loading
+47 −22
Original line number Diff line number Diff line
@@ -372,8 +372,8 @@ class Conv1D(Layer):
  """

  def __init__(self,
               width,
               out_channels,
               kernel_size,
               filter_size,
               stride=1,
               padding='SAME',
               activation_fn=tf.nn.relu,
@@ -381,6 +381,7 @@ class Conv1D(Layer):
               weights_initializer=tf.random_normal_initializer,
               **kwargs):
    """Create a Conv1D layer.
    TODO(LESWING) match function prototype to tf.keras.layers

    Parameters
    ----------
@@ -400,8 +401,8 @@ class Conv1D(Layer):
    weights_initializer: callable object
      the initializer for weight values
    """
    self.width = width
    self.out_channels = out_channels
    self.kernel_size = kernel_size
    self.filter_size = filter_size
    self.stride = stride
    self.padding = padding
    self.activation_fn = activation_fn
@@ -411,7 +412,7 @@ class Conv1D(Layer):
    super(Conv1D, self).__init__(**kwargs)
    try:
      parent_shape = self.in_layers[0].shape
      self._shape = (parent_shape[0], parent_shape[1] // stride, out_channels)
      self._shape = (parent_shape[0], parent_shape[1] // stride, filter_size)
    except:
      pass

@@ -424,18 +425,11 @@ class Conv1D(Layer):
      parent = tf.expand_dims(parent, 2)
    elif len(parent.get_shape()) != 3:
      raise ValueError("Parent tensor must be (batch, width, channel)")
    parent_shape = parent.get_shape()
    parent_channel_size = parent_shape[2].value
    f = tf.Variable(self.weights_initializer()(
        [self.width, parent_channel_size, self.out_channels]))
    t = tf.nn.conv1d(parent, f, stride=self.stride, padding=self.padding)
    if self.biases_initializer is not None:
      b = tf.Variable(self.biases_initializer()([self.out_channels]))
      t = tf.nn.bias_add(t, b)
    if self.activation_fn is None:
      out_tensor = t
    else:
      out_tensor = self.activation_fn(t)
    out_tensor = tf.keras.layers.Conv1D(
      filters=self.filter_size,
      kernel_size=self.kernel_size,
      activation=self.activation_fn,
    )(parent)
    if set_tensors:
      self._record_variable_scope(self.name)
      self.out_tensor = out_tensor
@@ -644,6 +638,38 @@ class Reshape(Layer):
    return out_tensor


class Cast(Layer):
  """
  Wrapper around tf.cast.  Changes the dtype of a single layer
  """

  def __init__(self, in_layers=None, dtype=None, **kwargs):
    """
    Parameters
    ----------
    dtype: tf.DType
      the dtype to cast the in_layer to
      e.x. tf.int32
    """
    if dtype is None:
      raise ValueError("Must cast to a dtype")
    self.dtype = dtype
    super(Cast, self).__init__(in_layers, **kwargs)
    try:
      parent_shape = self.in_layers[0].shape
      self._shape = parent_shape
    except:
      pass

  def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
    inputs = self._get_input_tensors(in_layers)
    parent_tensor = inputs[0]
    out_tensor = tf.cast(parent_tensor, self.dtype)
    if set_tensors:
      self.out_tensor = out_tensor
    return out_tensor


class Squeeze(Layer):

  def __init__(self, in_layers=None, squeeze_dims=None, **kwargs):
@@ -655,7 +681,8 @@ class Squeeze(Layer):
        self._shape = [i for i in parent_shape if i != 1]
      else:
        self._shape = [
            parent_shape[i] for i in range(len(parent_shape))
          parent_shape[i]
          for i in range(len(parent_shape))
          if i not in squeeze_dims
        ]
    except:
@@ -2817,15 +2844,13 @@ class VinaFreeEnergy(Layer):

  def hydrophobic(self, d):
    """Computes Autodock Vina's hydrophobic interaction term."""
    out_tensor = tf.where(d < 0.5,
                          tf.ones_like(d),
    out_tensor = tf.where(d < 0.5, tf.ones_like(d),
                          tf.where(d < 1.5, 1.5 - d, tf.zeros_like(d)))
    return out_tensor

  def hydrogen_bond(self, d):
    """Computes Autodock Vina's hydrogen bond interaction term."""
    out_tensor = tf.where(d < -0.7,
                          tf.ones_like(d),
    out_tensor = tf.where(d < -0.7, tf.ones_like(d),
                          tf.where(d < 0, (1.0 / 0.7) * (0 - d),
                                   tf.zeros_like(d)))
    return out_tensor
+73 −0
Original line number Diff line number Diff line
@@ -399,3 +399,76 @@ class SeqToSeq(TensorGraph):
      for initial, zero in zip(self.rnn_initial_states, self.rnn_zero_states):
        feed_dict[initial] = zero
      yield feed_dict


class AspuruGuzikAutoEncoder(SeqToSeq):
  def __init__(self,
               input_tokens,
               output_tokens,
               max_output_length,
               encoder_layers=4,
               decoder_layers=4,
               embedding_dimension=512,
               dropout=0.0,
               reverse_input=True,
               variational=True,
               annealing_start_step=5000,
               annealing_final_step=10000,
               **kwargs):
    """
    TODO(LESWING) have qm9/zinc hyper params and construct via static method
    """
    self.filter_sizes = [9, 9, 10]
    self.kernel_sizes = [9, 9, 11]
    super(AspuruGuzikAutoEncoder, self).__init__(
      input_tokens,
      output_tokens,
      max_output_length,
      encoder_layers,
      decoder_layers,
      embedding_dimension,
      dropout,
      reverse_input,
      variational,
      annealing_start_step,
      annealing_final_step,
      **kwargs
    )


  def _create_encoder(self, n_layers, dropout):
    """Create the encoder layers."""
    prev_layer = self._features
    for i in range(len(self.filter_sizes)):
      filter_size = self.filter_sizes[i]
      kernel_size = self.kernel_sizes[i]
      if dropout > 0.0:
        prev_layer = layers.Dropout(dropout, in_layers=prev_layer)
      prev_layer = layers.Conv1D(
        kernel_size, filter_size, in_layers=prev_layer, activation_fn=tf.nn.relu)
    prev_layer = layers.Gather(in_layers=[prev_layer, self._gather_indices])
    if self._variational:
      self._embedding_mean = layers.Dense(
        self._embedding_dimension, in_layers=prev_layer)
      self._embedding_stddev = layers.Dense(
        self._embedding_dimension, in_layers=prev_layer)
      prev_layer = layers.CombineMeanStd(
        [self._embedding_mean, self._embedding_stddev], training_only=True)
    return prev_layer

  def _create_decoder(self, n_layers, dropout):
    """
    TODO(LESWING): Teacher forcing
    """
    """Create the decoder layers."""
    prev_layer = layers.Repeat(
      self._max_output_length, in_layers=self.embedding)
    for i in range(3):
      if dropout > 0.0:
        prev_layer = layers.Dropout(dropout, in_layers=prev_layer)
      prev_layer = layers.GRU(
        488, self.batch_size, in_layers=prev_layer)
    return layers.Dense(
      len(self._output_tokens),
      in_layers=prev_layer,
      activation_fn=tf.nn.softmax)
+86 −57
Original line number Diff line number Diff line
%% Cell type:markdown id: tags:

SeqToSeq Fingerprint
--------------------

In this example, we will use a `SeqToSeq` model to generate fingerprints for classifying molecules.  This is based on the following paper, although some of the implementation details are different: Xu et al., "Seq2seq Fingerprint: An Unsupervised Deep Molecular Embedding for Drug Discovery" (https://doi.org/10.1145/3107411.3107424).

Many types of models require their inputs to have a fixed shape.  Since molecules can vary widely in the numbers of atoms and bonds they contain, this makes it hard to apply those models to them.  We need a way of generating a fixed length "fingerprint" for each molecule.  Various ways of doing this have been designed, such as Extended-Connectivity Fingerprints (ECFPs).  But in this example, instead of designing a fingerprint by hand, we will let a `SeqToSeq` model learn its own method of creating fingerprints.

A `SeqToSeq` model performs sequence to sequence translation.  For example, they are often used to translate text from one language to another.  It consists of two parts called the "encoder" and "decoder".  The encoder is a stack of recurrent layers.  The input sequence is fed into it, one token at a time, and it generates a fixed length vector called the "embedding vector".  The decoder is another stack of recurrent layers that performs the inverse operation: it takes the embedding vector as input, and generates the output sequence.  By training it on appropriately chosen input/output pairs, you can create a model that performs many sorts of transformations.

In this case, we will use SMILES strings describing molecules as the input sequences.  We will train the model as an autoencoder, so it tries to make the output sequences identical to the input sequences.  For that to work, the encoder must create embedding vectors that contain all information from the original sequence.  That's exactly what we want in a fingerprint, so perhaps those embedding vectors will then be useful as a way to represent molecules in other models!

Let's start by loading the data.  We will use the MUV dataset.  It includes 74,501 molecules in the training set, and 9313 molecules in the validation set, so it gives us plenty of SMILES strings to work with.

%% Cell type:code id: tags:

``` python
import deepchem as dc
tasks, datasets, transformers = dc.molnet.load_muv()
train_dataset, valid_dataset, test_dataset = datasets
train_smiles = train_dataset.ids
valid_smiles = valid_dataset.ids
```

%% Output

    /home/leswing/miniconda3/envs/deepchem/lib/python3.5/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.
      from ._conv import register_converters as _register_converters

    About to load MUV dataset.
    Loading dataset from disk.
    Loading dataset from disk.
    Loading dataset from disk.

%% Cell type:markdown id: tags:

We need to define the "alphabet" for our `SeqToSeq` model, the list of all tokens that can appear in sequences.  (It's also possible for input and output sequences to have different alphabets, but since we're training it as an autoencoder, they're identical in this case.)  Make a list of every character that appears in any training sequence.

%% Cell type:code id: tags:

``` python
tokens = set()
for s in train_smiles:
  tokens = tokens.union(set(c for c in s))
tokens = sorted(list(tokens))
```

%% Cell type:markdown id: tags:

Create the model and define the optimization method to use.  In this case, learning works much better if we gradually decrease the learning rate.  We use an `ExponentialDecay` to multiply the learning rate by 0.9 after each epoch.

%% Cell type:code id: tags:

``` python
from deepchem.models import SeqToSeq
from deepchem.models.tensorgraph import layers
from deepchem.models.tensorgraph.layers import Layer
import tensorflow as tf

class AspuruGuzikAutoEncoder(SeqToSeq):
  def __init__(self,
               input_tokens,
               output_tokens,
               max_output_length,
               encoder_layers=4,
               decoder_layers=4,
               embedding_dimension=512,
               dropout=0.0,
               reverse_input=True,
               variational=False,
               annealing_start_step=5000,
               annealing_final_step=10000,
               **kwargs):
    self.filter_sizes = [9, 9, 10]
    self.kernel_sizes = [9, 9, 11]
    super(AspuruGuzikAutoEncoder, self).__init__(
      input_tokens,
      output_tokens,
      max_output_length,
      encoder_layers,
      decoder_layers,
      embedding_dimension,
      dropout,
      reverse_input,
      variational,
      annealing_start_step,
      annealing_final_step,
    )


  def _create_encoder(self, n_layers, dropout):
    """Create the encoder layers."""
    prev_layer = self._features
    for i in range(len(self.filter_sizes)):
      filter_size = self.filter_sizes[i]
      kernel_size = self.kernel_sizes[i]
      if dropout > 0.0:
        prev_layer = layers.Dropout(dropout, in_layers=prev_layer)
      prev_layer = layers.Conv1D(
        kernel_size, filter_size, in_layers=prev_layer, activation_fn=tf.nn.relu)
    prev_layer = layers.Gather(in_layers=[prev_layer, self._gather_indices])
    if self._variational:
      self._embedding_mean = layers.Dense(
        196, in_layers=prev_layer)
      self._embedding_stddev = layers.Dense(
        196, in_layers=prev_layer)
      prev_layer = layers.CombineMeanStd(
        [self._embedding_mean, self._embedding_stddev], training_only=True)
    return prev_layer

  def _create_decoder(self, n_layers, dropout):
    """Create the decoder layers."""
    prev_layer = layers.Repeat(
      self._max_output_length, in_layers=self.embedding)
    for i in range(3):
      if dropout > 0.0:
        prev_layer = layers.Dropout(dropout, in_layers=prev_layer)
      prev_layer = layers.GRU(
        488, self.batch_size, in_layers=prev_layer)
    retval = layers.Dense(
      len(self._output_tokens),
      in_layers=prev_layer,
      activation_fn=tf.nn.softmax)
    return retval
```

%% Cell type:code id: tags:

``` python
from deepchem.models.tensorgraph.optimizers import Adam, ExponentialDecay


max_length = max(len(s) for s in train_smiles)
model = dc.models.SeqToSeq(tokens,
model = AspuruGuzikAutoEncoder(tokens,
                           tokens,
                           max_length,
                           encoder_layers=2,
                           decoder_layers=2,
                           embedding_dimension=256,
                           variational=False,
                           model_dir='fingerprint')
batches_per_epoch = len(train_smiles)/model.batch_size
model.set_optimizer(Adam(learning_rate=ExponentialDecay(0.004, 0.9, batches_per_epoch)))
```

%% Cell type:markdown id: tags:

Let's train it!  The input to `fit_sequences()` is a generator that produces input/output pairs.  On a good GPU, this should take a few hours or less.

%% Cell type:code id: tags:

``` python
def generate_sequences(epochs):
  for i in range(epochs):
    for s in train_smiles:
      yield (s, s)

model.fit_sequences(generate_sequences(40))
model.fit_sequences(generate_sequences(100000))
```

%% Output

    Ending global_step 999: Average loss 72.0029
    Ending global_step 1999: Average loss 40.7221
    Ending global_step 2999: Average loss 31.5364
    Ending global_step 3999: Average loss 26.4576
    Ending global_step 4999: Average loss 22.814
    Ending global_step 5999: Average loss 19.5248
    Ending global_step 6999: Average loss 16.4594
    Ending global_step 7999: Average loss 18.8898
    Ending global_step 8999: Average loss 13.476
    Ending global_step 9999: Average loss 11.5528
    Ending global_step 10999: Average loss 10.1594
    Ending global_step 11999: Average loss 10.6434
    Ending global_step 12999: Average loss 6.57057
    Ending global_step 13999: Average loss 6.46177
    Ending global_step 14999: Average loss 7.53559
    Ending global_step 15999: Average loss 4.95809
    Ending global_step 16999: Average loss 4.35039
    Ending global_step 17999: Average loss 3.39137
    Ending global_step 18999: Average loss 3.5216
    Ending global_step 19999: Average loss 3.08579
    Ending global_step 20999: Average loss 2.80738
    Ending global_step 21999: Average loss 2.92217
    Ending global_step 22999: Average loss 2.51032
    Ending global_step 23999: Average loss 1.86265
    Ending global_step 24999: Average loss 1.67088
    Ending global_step 25999: Average loss 1.87016
    Ending global_step 26999: Average loss 1.61166
    Ending global_step 27999: Average loss 1.40708
    Ending global_step 28999: Average loss 1.4488
    Ending global_step 29801: Average loss 1.33917
    TIMING: model fitting took 5619.924 s
    Ending global_step 999: Average loss 96.9049
    Ending global_step 1999: Average loss 95.9348

%% Cell type:markdown id: tags:

Let's see how well it works as an autoencoder.  We'll run the first 500 molecules from the validation set through it, and see how many of them are exactly reproduced.

%% Cell type:code id: tags:

``` python
predicted = model.predict_from_sequences(valid_smiles[:500])
count = 0
for s,p in zip(valid_smiles[:500], predicted):
  print(''.join(p))
  if ''.join(p) == s:
    count += 1
print('reproduced', count, 'of 500 validation SMILES strings')
```

%% Output

    reproduced 363 of 500 validation SMILES strings

%% Cell type:markdown id: tags:

Now we'll trying using the encoder as a way to generate molecular fingerprints.  We compute the embedding vectors for all molecules in the training and validation datasets, and create new datasets that have those as their feature vectors.  The amount of data is small enough that we can just store everything in memory.

%% Cell type:code id: tags:

``` python
train_embeddings = model.predict_embeddings(train_smiles)
train_embeddings_dataset = dc.data.NumpyDataset(train_embeddings,
                                                train_dataset.y,
                                                train_dataset.w,
                                                train_dataset.ids)

valid_embeddings = model.predict_embeddings(valid_smiles)
valid_embeddings_dataset = dc.data.NumpyDataset(valid_embeddings,
                                                valid_dataset.y,
                                                valid_dataset.w,
                                                valid_dataset.ids)
```

%% Cell type:markdown id: tags:

For classification, we'll use a simple fully connected network with one hidden layer.

%% Cell type:code id: tags:

``` python
classifier = dc.models.MultiTaskClassifier(n_tasks=len(tasks),
                                                      n_features=256,
                                                      layer_sizes=[512])
classifier.fit(train_embeddings_dataset, nb_epoch=10)
```

%% Output

    Ending global_step 999: Average loss 829.805
    Ending global_step 1999: Average loss 450.42
    Ending global_step 2999: Average loss 326.079
    Ending global_step 3999: Average loss 265.199
    Ending global_step 4999: Average loss 246.724
    Ending global_step 5999: Average loss 224.64
    Ending global_step 6999: Average loss 202.624
    Ending global_step 7460: Average loss 213.885
    TIMING: model fitting took 19.780 s

%% Cell type:markdown id: tags:

Find out how well it worked.  Compute the ROC AUC for the training and validation datasets.

%% Cell type:code id: tags:

``` python
import numpy as np
metric = dc.metrics.Metric(dc.metrics.roc_auc_score, np.mean, mode="classification")
train_score = classifier.evaluate(train_embeddings_dataset, [metric], transformers)
valid_score = classifier.evaluate(valid_embeddings_dataset, [metric], transformers)
print('Training set ROC AUC:', train_score)
print('Validation set ROC AUC:', valid_score)
```

%% Output

    computed_metrics: [0.97828427249789751, 0.98705973960125326, 0.966007068438685, 0.9874401066031584, 0.97794394675150698, 0.98021719680962449, 0.95318452689781941, 0.97185747562764213, 0.96389538770053473, 0.96798988621997473, 0.9690779239145807, 0.98544402211472004, 0.97762497271338133, 0.96843239633294886, 0.97753648081489997, 0.96504683675485614, 0.93547151958366914]
    computed_metrics: [0.90790686952512678, 0.79891461649782913, 0.61900937081659968, 0.75241212956581671, 0.58678903240426017, 0.72765072765072758, 0.34929006085192693, 0.83986814712005553, 0.82379943502824859, 0.61844636844636847, 0.863620199146515, 0.68106930272108857, 0.98020477815699669, 0.85073580939032944, 0.781015678254942, 0.75399733510992673, nan]
    Training set ROC AUC: {'mean-roc_auc_score': 0.97132433878689139}
    Validation set ROC AUC: {'mean-roc_auc_score': 0.74592061629292239}