Commit 690025e7 authored by Bharath Ramsundar's avatar Bharath Ramsundar
Browse files

fixes

parent 8672f280

deepchem/models/graph_models.py

+292 −75
Original line number Diff line number Diff line
@@ -495,6 +495,115 @@ class DAGModel(KerasModel): 
        ], [y_b], [w_b])





class GraphConvKerasModel(tf.keras.Model):



  def __init__(self,

               n_tasks,

               graph_conv_layers,

               dense_layer_size=128,

               dropout=0.0,

               mode="classification",

               number_atom_features=75,

               n_classes=2,

               uncertainty=False,

               batch_size=100):

    """An internal keras model class.



    The graph convolutions use a nonstandard control flow so the

    standard Keras functional API can't support them. We instead

    use the imperative "subclassing" API to implement the graph

    convolutions.



    All arguments have the same meaning as in GraphConvModel.

    """

    super(GraphConvKerasModel, self).__init__()

    if mode not in ['classification', 'regression']:

      raise ValueError("mode must be either 'classification' or 'regression'")



    self.mode = mode

    self.uncertainty = uncertainty



    if not isinstance(dropout, collections.Sequence):

      dropout = [dropout] * (len(graph_conv_layers) + 1)

    if len(dropout) != len(graph_conv_layers) + 1:

      raise ValueError('Wrong number of dropout probabilities provided')

    if uncertainty:

      if mode != "regression":

        raise ValueError("Uncertainty is only supported in regression mode")

      if any(d == 0.0 for d in dropout):

        raise ValueError(

            'Dropout must be included in every layer to predict uncertainty')



    self.graph_convs = [

        layers.GraphConv(layer_size, activation_fn=tf.nn.relu)

        for layer_size in graph_conv_layers

    ]

    self.batch_norms = [

        BatchNormalization(fused=False)

        for _ in range(len(graph_conv_layers) + 1)

    ]

    self.dropouts = [

        layers.SwitchedDropout(rate=rate) if rate > 0.0 else None

        for rate in dropout

    ]

    self.graph_pools = [layers.GraphPool() for _ in graph_conv_layers]

    self.dense = Dense(dense_layer_size, activation=tf.nn.relu)

    self.graph_gather = layers.GraphGather(

        batch_size=batch_size, activation_fn=tf.nn.tanh)

    self.trim = TrimGraphOutput()

    if self.mode == 'classification':

      self.reshape_dense = Dense(n_tasks * n_classes)

      self.reshape = Reshape((n_tasks, n_classes))

      self.softmax = Softmax()

    else:

      self.regression_dense = Dense(n_tasks)

      if self.uncertainty:

        self.uncertainty_dense = Dense(n_tasks)

        self.uncertainty_trim = TrimGraphOutput()

        self.uncertainty_activation = Activation(tf.exp)



  def call(self, inputs):

    atom_features = inputs[0]

    degree_slice = tf.cast(inputs[1], dtype=tf.int32)

    membership = tf.cast(inputs[2], dtype=tf.int32)

    n_samples = tf.cast(inputs[3], dtype=tf.int32)

    dropout_switch = inputs[4]

    deg_adjs = [tf.cast(deg_adj, dtype=tf.int32) for deg_adj in inputs[5:]]



    in_layer = atom_features

    for i in range(len(self.graph_convs)):

      gc_in = [in_layer, degree_slice, membership] + deg_adjs

      gc1 = self.graph_convs[i](gc_in)

      batch_norm1 = self.batch_norms[i](gc1)

      if self.dropouts[i] is not None:

        batch_norm1 = self.dropouts[i]([batch_norm1, dropout_switch])

      gp_in = [batch_norm1, degree_slice, membership] + deg_adjs

      in_layer = self.graph_pools[i](gp_in)

    dense = self.dense(in_layer)

    batch_norm3 = self.batch_norms[-1](dense)

    if self.dropouts[-1] is not None:

      batch_norm3 = self.dropouts[1]([batch_norm3, dropout_switch])

    neural_fingerprint = self.graph_gather(

        [batch_norm3, degree_slice, membership] + deg_adjs)

    if self.mode == 'classification':

      logits = self.reshape(self.reshape_dense(neural_fingerprint))

      logits = self.trim([logits, n_samples])

      output = self.softmax(logits)

      outputs = [output, logits, neural_fingerprint]

    else:

      output = self.regression_dense(neural_fingerprint)

      output = self.trim([output, n_samples])

      if self.uncertainty:

        log_var = self.uncertainty_dense(neural_fingerprint)

        log_var = self.uncertainty_trim([log_var, n_samples])

        var = self.uncertainty_activation(log_var)

        outputs = [output, var, output, log_var, neural_fingerprint]

      else:

        outputs = [output, neural_fingerprint]



    return outputs





class GraphConvModel(KerasModel):



  def __init__(self,
@@ -505,10 +614,15 @@ class GraphConvModel(KerasModel): 
               mode="classification",

               number_atom_features=75,

               n_classes=2,

               uncertainty=False,

               batch_size=100,

               uncertainty=False,

               **kwargs):

    """

    """The wrapper class for graph convolutions.



    Note that since the underlying GraphConvKerasModel class is

    specified using imperative subclassing style, this model

    cannout make predictions for arbitrary outputs. 



    Parameters

    ----------

    n_tasks: int
@@ -534,94 +648,40 @@ class GraphConvModel(KerasModel): 
      if True, include extra outputs and loss terms to enable the uncertainty

      in outputs to be predicted

    """

    if mode not in ['classification', 'regression']:

      raise ValueError("mode must be either 'classification' or 'regression'")

    self.n_tasks = n_tasks

    self.mode = mode

    self.dense_layer_size = dense_layer_size

    self.graph_conv_layers = graph_conv_layers

    self.number_atom_features = number_atom_features

    self.n_tasks = n_tasks

    self.n_classes = n_classes

    self.batch_size = batch_size

    self.uncertainty = uncertainty

    if not isinstance(dropout, collections.Sequence):

      dropout = [dropout] * (len(graph_conv_layers) + 1)

    if len(dropout) != len(graph_conv_layers) + 1:

      raise ValueError('Wrong number of dropout probabilities provided')

    self.dropout = dropout

    if uncertainty:

      if mode != "regression":

        raise ValueError("Uncertainty is only supported in regression mode")

      if any(d == 0.0 for d in dropout):

        raise ValueError(

            'Dropout must be included in every layer to predict uncertainty')



    # Build the model.



    atom_features = Input(shape=(self.number_atom_features,))

    degree_slice = Input(shape=(2,), dtype=tf.int32)

    membership = Input(shape=tuple(), dtype=tf.int32)

    n_samples = Input(shape=tuple(), dtype=tf.int32)

    dropout_switch = tf.keras.Input(shape=tuple())



    self.deg_adjs = []

    for i in range(0, 10 + 1):

      deg_adj = Input(shape=(i + 1,), dtype=tf.int32)

      self.deg_adjs.append(deg_adj)

    in_layer = atom_features

    for layer_size, dropout in zip(self.graph_conv_layers, self.dropout):

      gc1_in = [in_layer, degree_slice, membership] + self.deg_adjs

      gc1 = layers.GraphConv(layer_size, activation_fn=tf.nn.relu)(gc1_in)

      batch_norm1 = BatchNormalization(fused=False)(gc1)

      if dropout > 0.0:

        batch_norm1 = layers.SwitchedDropout(rate=dropout)(

            [batch_norm1, dropout_switch])

      gp_in = [batch_norm1, degree_slice, membership] + self.deg_adjs

      in_layer = layers.GraphPool()(gp_in)

    dense = Dense(self.dense_layer_size, activation=tf.nn.relu)(in_layer)

    batch_norm3 = BatchNormalization(fused=False)(dense)

    if self.dropout[-1] > 0.0:

      batch_norm3 = layers.SwitchedDropout(rate=self.dropout[-1])(

          [batch_norm3, dropout_switch])

    self.neural_fingerprint = layers.GraphGather(

        batch_size=batch_size,

        activation_fn=tf.nn.tanh)([batch_norm3, degree_slice, membership] +

                                  self.deg_adjs)



    n_tasks = self.n_tasks

    if self.mode == 'classification':

      n_classes = self.n_classes

      logits = Reshape((n_tasks, n_classes))(Dense(n_tasks * n_classes)(

          self.neural_fingerprint))

      logits = TrimGraphOutput()([logits, n_samples])

      output = Softmax()(logits)

      outputs = [output, logits]

      output_types = ['prediction', 'loss']

    model = GraphConvKerasModel(

        n_tasks,

        graph_conv_layers=graph_conv_layers,

        dense_layer_size=dense_layer_size,

        dropout=dropout,

        mode=mode,

        number_atom_features=number_atom_features,

        n_classes=n_classes,

        uncertainty=uncertainty,

        batch_size=batch_size)

    if mode == "classification":

      output_types = ['prediction', 'loss', 'embedding']

      loss = SoftmaxCrossEntropy()

    else:

      output = Dense(n_tasks)(self.neural_fingerprint)

      output = TrimGraphOutput()([output, n_samples])

      if self.uncertainty:

        log_var = Dense(n_tasks)(self.neural_fingerprint)

        log_var = TrimGraphOutput()([log_var, n_samples])

        var = Activation(tf.exp)(log_var)

        outputs = [output, var, output, log_var]

        output_types = ['prediction', 'variance', 'loss', 'loss']

        output_types = ['prediction', 'variance', 'loss', 'loss', 'embedding']



        def loss(outputs, labels, weights):

          diff = labels[0] - outputs[0]

          return tf.reduce_mean(diff * diff / tf.exp(outputs[1]) + outputs[1])

      else:

        outputs = [output]

        output_types = ['prediction']

        output_types = ['prediction', 'embedding']

        loss = L2Loss()

    model = tf.keras.Model(

        inputs=[

            atom_features, degree_slice, membership, n_samples, dropout_switch

        ] + self.deg_adjs,

        outputs=outputs)

    super(GraphConvModel, self).__init__(

        model, loss, output_types=output_types, batch_size=batch_size, **kwargs)



  def fit(self, *args, **kwargs):

    super(GraphConvModel, self).fit(*args, **kwargs)



  def default_generator(self,

                        dataset,

                        epochs=1,
@@ -651,6 +711,163 @@ class GraphConvModel(KerasModel): 
        yield (inputs, [y_b], [w_b])





#class GraphConvModel(KerasModel):

#

#  def __init__(self,

#               n_tasks,

#               graph_conv_layers=[64, 64],

#               dense_layer_size=128,

#               dropout=0.0,

#               mode="classification",

#               number_atom_features=75,

#               n_classes=2,

#               uncertainty=False,

#               batch_size=100,

#               **kwargs):

#    """

#    Parameters

#    ----------

#    n_tasks: int

#      Number of tasks

#    graph_conv_layers: list of int

#      Width of channels for the Graph Convolution Layers

#    dense_layer_size: int

#      Width of channels for Atom Level Dense Layer before GraphPool

#    dropout: list or float

#      the dropout probablity to use for each layer.  The length of this list should equal

#      len(graph_conv_layers)+1 (one value for each convolution layer, and one for the

#      dense layer).  Alternatively this may be a single value instead of a list, in which

#      case the same value is used for every layer.

#    mode: str

#      Either "classification" or "regression"

#    number_atom_features: int

#        75 is the default number of atom features created, but

#        this can vary if various options are passed to the

#        function atom_features in graph_features

#    n_classes: int

#      the number of classes to predict (only used in classification mode)

#    uncertainty: bool

#      if True, include extra outputs and loss terms to enable the uncertainty

#      in outputs to be predicted

#    """

#    if mode not in ['classification', 'regression']:

#      raise ValueError("mode must be either 'classification' or 'regression'")

#    self.n_tasks = n_tasks

#    self.mode = mode

#    self.dense_layer_size = dense_layer_size

#    self.graph_conv_layers = graph_conv_layers

#    self.number_atom_features = number_atom_features

#    self.n_classes = n_classes

#    self.uncertainty = uncertainty

#    if not isinstance(dropout, collections.Sequence):

#      dropout = [dropout] * (len(graph_conv_layers) + 1)

#    if len(dropout) != len(graph_conv_layers) + 1:

#      raise ValueError('Wrong number of dropout probabilities provided')

#    self.dropout = dropout

#    if uncertainty:

#      if mode != "regression":

#        raise ValueError("Uncertainty is only supported in regression mode")

#      if any(d == 0.0 for d in dropout):

#        raise ValueError(

#            'Dropout must be included in every layer to predict uncertainty')

#

#    # Build the model.

#

#    atom_features = Input(shape=(self.number_atom_features,))

#    degree_slice = Input(shape=(2,), dtype=tf.int32)

#    membership = Input(shape=tuple(), dtype=tf.int32)

#    n_samples = Input(shape=tuple(), dtype=tf.int32)

#    dropout_switch = tf.keras.Input(shape=tuple())

#

#    self.deg_adjs = []

#    for i in range(0, 10 + 1):

#      deg_adj = Input(shape=(i + 1,), dtype=tf.int32)

#      self.deg_adjs.append(deg_adj)

#    in_layer = atom_features

#    for layer_size, dropout in zip(self.graph_conv_layers, self.dropout):

#      gc1_in = [in_layer, degree_slice, membership] + self.deg_adjs

#      gc1 = layers.GraphConv(layer_size, activation_fn=tf.nn.relu)(gc1_in)

#      batch_norm1 = BatchNormalization(fused=False)(gc1)

#      if dropout > 0.0:

#        batch_norm1 = layers.SwitchedDropout(rate=dropout)(

#            [batch_norm1, dropout_switch])

#      gp_in = [batch_norm1, degree_slice, membership] + self.deg_adjs

#      in_layer = layers.GraphPool()(gp_in)

#    dense = Dense(self.dense_layer_size, activation=tf.nn.relu)(in_layer)

#    batch_norm3 = BatchNormalization(fused=False)(dense)

#    if self.dropout[-1] > 0.0:

#      batch_norm3 = layers.SwitchedDropout(rate=self.dropout[-1])(

#          [batch_norm3, dropout_switch])

#    self.neural_fingerprint = batch_norm3

#    self.neural_fingerprint = layers.GraphGather(

#        batch_size=batch_size,

#        activation_fn=tf.nn.tanh)([batch_norm3, degree_slice, membership] +

#                                  self.deg_adjs)

#

#    n_tasks = self.n_tasks

#    if self.mode == 'classification':

#      n_classes = self.n_classes

#      logits = Reshape((n_tasks, n_classes))(Dense(n_tasks * n_classes)(

#          self.neural_fingerprint))

#      logits = TrimGraphOutput()([logits, n_samples])

#      output = Softmax()(logits)

#      outputs = [output, logits]

#      output_types = ['prediction', 'loss']

#      loss = SoftmaxCrossEntropy()

#    else:

#      output = Dense(n_tasks)(self.neural_fingerprint)

#      output = TrimGraphOutput()([output, n_samples])

#      if self.uncertainty:

#        log_var = Dense(n_tasks)(self.neural_fingerprint)

#        log_var = TrimGraphOutput()([log_var, n_samples])

#        var = Activation(tf.exp)(log_var)

#        outputs = [output, var, output, log_var]

#        output_types = ['prediction', 'variance', 'loss', 'loss']

#

#        def loss(outputs, labels, weights):

#          diff = labels[0] - outputs[0]

#          return tf.reduce_mean(diff * diff / tf.exp(outputs[1]) + outputs[1])

#      else:

#        outputs = [output]

#        output_types = ['prediction']

#        loss = L2Loss()

#    model = tf.keras.Model(

#        inputs=[

#            atom_features, degree_slice, membership, n_samples, dropout_switch

#        ] + self.deg_adjs,

#        outputs=outputs)

#    super(GraphConvModel, self).__init__(

#        model, loss, output_types=output_types, batch_size=batch_size, **kwargs)

#

#  def default_generator(self,

#                        dataset,

#                        epochs=1,

#                        mode='fit',

#                        deterministic=True,

#                        pad_batches=True):

#    for epoch in range(epochs):

#      for (X_b, y_b, w_b, ids_b) in dataset.iterbatches(

#          batch_size=self.batch_size,

#          deterministic=deterministic,

#          pad_batches=pad_batches):

#        if self.mode == 'classification':

#          y_b = to_one_hot(y_b.flatten(), self.n_classes).reshape(

#              -1, self.n_tasks, self.n_classes)

#        multiConvMol = ConvMol.agglomerate_mols(X_b)

#        n_samples = np.array(X_b.shape[0])

#        if mode == 'predict':

#          dropout = np.array(0.0)

#        else:

#          dropout = np.array(1.0)

#        inputs = [

#            multiConvMol.get_atom_features(), multiConvMol.deg_slice,

#            np.array(multiConvMol.membership), n_samples, dropout

#        ]

#        for i in range(1, len(multiConvMol.get_deg_adjacency_lists())):

#          inputs.append(multiConvMol.get_deg_adjacency_lists()[i])

#        yield (inputs, [y_b], [w_b])





class MPNNModel(KerasModel):

  """ Message Passing Neural Network,

      default structures built according to https://arxiv.org/abs/1511.06391 """

deepchem/models/keras_model.py

+60 −5
Original line number Diff line number Diff line
@@ -76,6 +76,8 @@ class KerasModel(Model): 
    Also be aware that if a model supports uncertainty, it MUST use dropout on

    every layer, and dropout most be enabled during uncertainty prediction.

    Otherwise, the uncertainties it computes will be inaccurate.

  - 'embedding': This output is an embedding that the model

  generates internally which should be returned to users.

  """



  def __init__(self,
@@ -136,10 +138,12 @@ class KerasModel(Model): 
      self._prediction_outputs = None

      self._loss_outputs = None

      self._variance_outputs = None

      self._embedding_outputs = None

    else:

      self._prediction_outputs = []

      self._loss_outputs = []

      self._variance_outputs = []

      self._embedding_outputs = []

      for i, type in enumerate(output_types):

        if type == 'prediction':

          self._prediction_outputs.append(i)
@@ -147,6 +151,8 @@ class KerasModel(Model): 
          self._loss_outputs.append(i)

        elif type == 'variance':

          self._variance_outputs.append(i)

        elif type == 'embedding':

          self._embedding_outputs.append(i)

        else:

          raise ValueError('Unknown output type "%s"' % type)

      if len(self._loss_outputs) == 0:
@@ -321,6 +327,7 @@ class KerasModel(Model): 


      if len(inputs) == 1:

        inputs = inputs[0]



      batch_loss = apply_gradient_for_batch(inputs, labels, weights, loss)

      current_step = self._global_step.numpy()
@@ -415,7 +422,7 @@ class KerasModel(Model): 
        loss=loss,

        callbacks=callbacks)



  def _predict(self, generator, transformers, outputs, uncertainty):

  def _predict(self, generator, transformers, outputs, uncertainty, embedding):

    """

    Predict outputs for data provided by a generator.
@@ -439,12 +446,19 @@ class KerasModel(Model): 
      specifies whether this is being called as part of estimating uncertainty.

      If True, it sets the training flag so that dropout will be enabled, and

      returns the values of the uncertainty outputs.

    embedding: bool

      specifies whether this is being called as part of generating embeddings.

    Returns:

      a NumPy array of the model produces a single output, or a list of arrays

      if it produces multiple outputs

    """

    results = None

    variances = None

    embeddings = None

    if uncertainty and embedding:

      raise ValueError(

          'This model cannot compute uncertainties and embeddings simultaneously. Please invoke one at a time.'

      )

    if uncertainty:

      assert outputs is None

      if self._variance_outputs is None or len(self._variance_outputs) == 0:
@@ -452,9 +466,14 @@ class KerasModel(Model): 
      if len(self._variance_outputs) != len(self._prediction_outputs):

        raise ValueError(

            'The number of variances must exactly match the number of outputs')

    if outputs is not None and len(self.model.inputs) == 0:

    if embedding:

      assert outputs is None

      if self._embedding_outputs is None or len(self._embedding_outputs) == 0:

        raise ValueError('This model cannot compute embneddings.')

    if (outputs is not None and self.model.inputs is not None and

        len(self.model.inputs) == 0):

      raise ValueError(

          "Cannot use 'outputs' argument with a model that does not specify its inputs"

          "Cannot use 'outputs' argument with a model that does not specify its inputs. Note models defined in imperative subclassing style cannot specify outputs"

      )

    if isinstance(outputs, tf.Tensor):

      outputs = [outputs]
@@ -474,6 +493,7 @@ class KerasModel(Model): 
          self._output_functions[key] = tf.keras.backend.function(

              self.model.inputs, outputs)

        output_values = self._output_functions[key](inputs)

        output_values = self._compute_model(inputs)

      else:

        output_values = self._compute_model(inputs)

        if isinstance(output_values, tf.Tensor):
@@ -489,6 +509,16 @@ class KerasModel(Model): 
        else:

          for i, t in enumerate(var):

            variances[i].append(t)

      if embedding:

        ##################################

        print("len(self._embedding_outputs)")

        print(len(self._embedding_outputs))

        ##################################

        embeddings = [output_values[i] for i in self._embedding_outputs]

        ##################################

        print("embeddings[0].shape")

        print(embeddings[0].shape)

        ##################################

      if self._prediction_outputs is not None:

        output_values = [output_values[i] for i in self._prediction_outputs]

      if len(transformers) > 0:
@@ -507,12 +537,19 @@ class KerasModel(Model): 


    final_results = []

    final_variances = []

    final_embeddings = []

    for r in results:

      final_results.append(np.concatenate(r, axis=0))

    if uncertainty:

      for v in variances:

        final_variances.append(np.concatenate(v, axis=0))

      return zip(final_results, final_variances)

    if embedding:

      final_embeddings = embeddings

      if len(final_embeddings) == 1:

        return final_embeddings[0]

      else:

        return final_embeddings

    # If only one output, just return array

    if len(final_results) == 1:

      return final_results[0]
@@ -543,7 +580,7 @@ class KerasModel(Model): 
      a NumPy array of the model produces a single output, or a list of arrays

      if it produces multiple outputs

    """

    return self._predict(generator, transformers, outputs, False)

    return self._predict(generator, transformers, outputs, False, False)



  def predict_on_batch(self, X, transformers=[], outputs=None):

    """Generates predictions for input samples, processing samples in a batch.
@@ -622,6 +659,24 @@ class KerasModel(Model): 
        dataset, mode='predict', pad_batches=False)

    return self.predict_on_generator(generator, transformers, outputs)



  def predict_embedding(self, dataset):

    """

    Predicts embeddings created by underlying model 



    Parameters

    ----------

    dataset: dc.data.Dataset

      Dataset to make prediction on



    Returns

    -------

    a NumPy array of the embeddings model produces, or a list

    of arrays if it produces multiple embeddings

    """

    generator = self.default_generator(

        dataset, mode='predict', pad_batches=False)

    return self._predict(generator, [], None, False, True)



  def predict_uncertainty(self, dataset, masks=50):

    """

    Predict the model's outputs, along with the uncertainty in each one.
@@ -652,7 +707,7 @@ class KerasModel(Model): 
    for i in range(masks):

      generator = self.default_generator(

          dataset, mode='uncertainty', pad_batches=False)

      results = self._predict(generator, [], None, True)

      results = self._predict(generator, [], None, True, False)

      if len(sum_pred) == 0:

        for p, v in results:

          sum_pred.append(p)

deepchem/models/tests/test_graph_models.py

+1 −2
Original line number Diff line number Diff line
@@ -68,8 +68,7 @@ class TestGraphModels(unittest.TestCase): 
        mode='classification')



    model.fit(dataset, nb_epoch=1)

    neural_fingerprints = model.predict(

        dataset, outputs=model.neural_fingerprint)

    neural_fingerprints = model.predict_embedding(dataset)

    neural_fingerprints = np.array(neural_fingerprints)[:len(dataset)]

    self.assertEqual((len(dataset), fp_size * 2), neural_fingerprints.shape)

deepchem/models/tests/test_layers.py

+4 −0
Original line number Diff line number Diff line
@@ -88,6 +88,10 @@ class TestLayers(test_util.TensorFlowTestCase): 
    membership = multi_mol.membership

    deg_adjs = multi_mol.get_deg_adjacency_lists()[1:]

    args = [atom_features, degree_slice, membership] + deg_adjs

    print("HIIIIII")

    print("deg_adjs")

    print(deg_adjs)

    assert 0 == 1

    layer = layers.GraphConv(out_channels)

    result = layer(args)

    assert result.shape == (n_atoms, out_channels)

CRA Git | Maintained and supported by SUSTech CRA and CCSE