Unverified Commit 812fa8f3 authored by Bharath Ramsundar's avatar Bharath Ramsundar Committed by GitHub
Browse files

Merge pull request #1215 from miaecle/ani1

ANI-1 perf fix
parents 6749d8d1 d95b5493
Loading
Loading
Loading
Loading
+12 −4
Original line number Diff line number Diff line
@@ -15,8 +15,8 @@ import tensorflow as tf

import deepchem as dc

from deepchem.models.tensorgraph.layers import Dense, Concat, WeightedError, Stack, Layer, ANIFeat
from deepchem.models.tensorgraph.layers import L2Loss, Label, Weights, Feature
from deepchem.models.tensorgraph.layers import Dense, Concat, WeightedError, Stack, Layer, ANIFeat, Exp
from deepchem.models.tensorgraph.layers import L2Loss, Label, Weights, Feature, Dropout, WeightDecay
from deepchem.models.tensorgraph.tensor_graph import TensorGraph
from deepchem.models.tensorgraph.graph_layers import DTNNEmbedding
from deepchem.models.tensorgraph.symmetry_functions import DistanceMatrix, \
@@ -115,6 +115,8 @@ class ANIRegression(TensorGraph):
  def __init__(self,
               n_tasks,
               max_atoms,
               exp_loss=False,
               activation_fn='ani',
               layer_structures=[128, 64],
               atom_number_cases=[1, 6, 7, 8, 16],
               **kwargs):
@@ -130,6 +132,8 @@ class ANIRegression(TensorGraph):
    """
    self.n_tasks = n_tasks
    self.max_atoms = max_atoms
    self.exp_loss = exp_loss
    self.activation_fn = activation_fn
    self.layer_structures = layer_structures
    self.atom_number_cases = atom_number_cases
    super(ANIRegression, self).__init__(**kwargs)
@@ -291,7 +295,9 @@ class ANIRegression(TensorGraph):
        jac=self.grad_one,
        method="BFGS",
        tol=1e-6,
        options={'disp': True})
        options={
            'disp': True
        })

    return res.x.reshape((num_atoms, 3))

@@ -312,7 +318,7 @@ class ANIRegression(TensorGraph):
          self.max_atoms,
          n_hidden,
          self.atom_number_cases,
          activation='tanh',
          activation=self.activation_fn,
          in_layers=[previous_layer, self.atom_numbers])
      Hiddens.append(Hidden)
      previous_layer = Hiddens[-1]
@@ -333,6 +339,8 @@ class ANIRegression(TensorGraph):
    all_cost = Stack(in_layers=costs, axis=1)
    self.weights = Weights(shape=(None, self.n_tasks))
    loss = WeightedError(in_layers=[all_cost, self.weights])
    if self.exp_loss:
      loss = Exp(in_layers=[loss])
    self.set_loss(loss)

  def default_generator(self,
+9 −2
Original line number Diff line number Diff line
@@ -399,7 +399,7 @@ class AtomicDifferentiatedDense(Layer):
               out_channels,
               atom_number_cases=[1, 6, 7, 8],
               init='glorot_uniform',
               activation='relu',
               activation='ani',
               **kwargs):
    self.init = init  # Set weight initialization
    self.activation = activation  # Get activations
@@ -409,9 +409,16 @@ class AtomicDifferentiatedDense(Layer):

    super(AtomicDifferentiatedDense, self).__init__(**kwargs)

  @staticmethod
  def ani_activate(X):
    return tf.exp(-1 * tf.pow(X, 2))

  def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
    """ Generate Radial Symmetry Function """
    init_fn = initializations.get(self.init)  # Set weight initialization
    if self.activation == 'ani':
      activation_fn = self.ani_activate
    else:
      activation_fn = activations.get(self.activation)
    if in_layers is None:
      in_layers = self.in_layers
+67 −35
Original line number Diff line number Diff line
@@ -10,24 +10,43 @@ np.random.seed(123)
import tensorflow as tf
tf.set_random_seed(123)
import deepchem as dc
import os
import tempfile

HARTREE_TO_KCAL_PER_MOL = 627.509

# Load Tox21 dataset
tasks, datasets, transformers = dc.molnet.load_qm7_from_mat(
    featurizer='BPSymmetryFunction')
train_dataset, valid_dataset, test_dataset = datasets
    featurizer='BPSymmetryFunction', split='index', reload=False)
all_dataset = dc.data.DiskDataset.merge(datasets)
invalid_inds = []
X = all_dataset.X
for i in range(X.shape[0]):
  # Exclude all molecules having S
  if 16 in X[i, :, 0]:
    invalid_inds.append(i)
valid_inds = np.delete(np.arange(all_dataset.y.shape[0]), invalid_inds)
dataset = all_dataset.select(valid_inds)

splitter = dc.splits.RandomSplitter()
train, valid, test = splitter.train_valid_test_split(dataset)

y = dc.trans.undo_transforms(train.y, transformers) / HARTREE_TO_KCAL_PER_MOL
train = dc.data.DiskDataset.from_numpy(
    train.X, y, w=train.w, ids=train.ids, tasks=train.tasks)

y = dc.trans.undo_transforms(valid.y, transformers) / HARTREE_TO_KCAL_PER_MOL
valid = dc.data.DiskDataset.from_numpy(
    valid.X, y, w=valid.w, ids=valid.ids, tasks=valid.tasks)

y = dc.trans.undo_transforms(test.y, transformers) / HARTREE_TO_KCAL_PER_MOL
test = dc.data.DiskDataset.from_numpy(
    test.X, y, w=test.w, ids=test.ids, tasks=test.tasks)

# Batch size of models
max_atoms = 23
batch_size = 128
layer_structures = [128, 128, 64]
atom_number_cases = [1, 6, 7, 8, 16]

ANItransformer = dc.trans.ANITransformer(
    max_atoms=max_atoms, atom_cases=atom_number_cases)
train_dataset = ANItransformer.transform(train_dataset)
valid_dataset = ANItransformer.transform(valid_dataset)
test_dataset = ANItransformer.transform(test_dataset)
n_feat = ANItransformer.get_num_feats() - 1
layer_structures = [64, 64, 32]
atom_number_cases = [1, 6, 7, 8]

# Fit models
metric = [
@@ -35,26 +54,39 @@ metric = [
    dc.metrics.Metric(dc.metrics.pearson_r2_score, mode="regression")
]

model_dir = tempfile.mkdtemp()

lr_scedule = [1e-3, 1e-4, 1e-5, 3e-6, 1e-6, 3e-7, 1e-7, 3e-8, 1e-8, 3e-9, 1e-9]

valid_best = 100.
for lr in lr_scedule:
  model = dc.models.ANIRegression(
      len(tasks),
      max_atoms,
    n_feat,
      exp_loss=False,
      layer_structures=layer_structures,
      atom_number_cases=atom_number_cases,
      batch_size=batch_size,
    learning_rate=0.001,
      learning_rate=lr,
      use_queue=False,
    mode="regression")

# Fit trained model
model.fit(train_dataset, nb_epoch=300, checkpoint_interval=100)

print("Evaluating model")
train_scores = model.evaluate(train_dataset, metric, transformers)
valid_scores = model.evaluate(valid_dataset, metric, transformers)

print("Train scores")
print(train_scores)

print("Validation scores")
print(valid_scores)
      mode="regression",
      model_dir=model_dir)
  if lr < 1e-3:
    model.restore()
  model.fit(train, nb_epoch=10)
  local_ct = 0
  while local_ct < 100:
    local_ct += 1
    model.fit(train, nb_epoch=1)
    train_scores = model.evaluate(train, metric)
    valid_scores = model.evaluate(valid, metric)
    print("Train MAE(kcal/mol): " +
          str(train_scores['mean_absolute_error'] * HARTREE_TO_KCAL_PER_MOL))
    print("Valid MAE(kcal/mol): " +
          str(valid_scores['mean_absolute_error'] * HARTREE_TO_KCAL_PER_MOL))
    if valid_scores['mean_absolute_error'] < valid_best:
      local_ct = 0
      valid_best = valid_scores['mean_absolute_error']
      test_scores = model.evaluate(test, metric)
      print("Test MAE(kcal/mol): " +
            str(test_scores['mean_absolute_error'] * HARTREE_TO_KCAL_PER_MOL))
+13 −8
Original line number Diff line number Diff line
@@ -10,8 +10,9 @@ np.random.seed(123)
import tensorflow as tf
tf.set_random_seed(123)
import deepchem as dc
from deepchem.models.tensorgraph.optimizers import ExponentialDecay

# Load Tox21 dataset
# Load QM7 dataset
tasks, datasets, transformers = dc.molnet.load_qm7_from_mat()
train_dataset, valid_dataset, test_dataset = datasets

@@ -29,6 +30,7 @@ distance_min = -1.
distance_max = 9.2
n_hidden = 15

rate = ExponentialDecay(0.0001, 0.97, 5000)
model = dc.models.DTNNModel(
    len(tasks),
    n_embedding=n_embedding,
@@ -38,19 +40,22 @@ model = dc.models.DTNNModel(
    distance_max=distance_max,
    output_activation=False,
    batch_size=batch_size,
    learning_rate=0.0001,
    learning_rate=rate,
    use_queue=False,
    mode="regression")
#model.restore()

# Fit trained model
model.fit(train_dataset, nb_epoch=1000)
model.fit(train_dataset, nb_epoch=3000)

print("Evaluating model")
train_scores = model.evaluate(train_dataset, metric, transformers)
valid_scores = model.evaluate(valid_dataset, metric, transformers)

print("Train scores")
print("Train scores [kcal/mol]")
print(train_scores)

print("Validation scores")
valid_scores = model.evaluate(valid_dataset, metric, transformers)
print("Valid scores [kcal/mol]")
print(valid_scores)

test_scores = model.evaluate(test_dataset, metric, transformers)
print("Test scores [kcal/mol]")
print(test_scores)
+10 −7
Original line number Diff line number Diff line
@@ -9,19 +9,22 @@ import os
import deepchem as dc
import numpy as np
from deepchem.molnet import load_qm7_from_mat
from deepchem.models.tensorgraph.optimizers import ExponentialDecay

np.random.seed(123)
qm7_tasks, datasets, transformers = load_qm7_from_mat(split='stratified')
train_dataset, valid_dataset, test_dataset = datasets
fit_transformers = [dc.trans.CoulombFitTransformer(train_dataset)]
regression_metric = [
metric = [
    dc.metrics.Metric(dc.metrics.mean_absolute_error, mode="regression"),
    dc.metrics.Metric(dc.metrics.pearson_r2_score, mode="regression")
]

rate = ExponentialDecay(0.001, 0.95, 1000)
model = dc.models.MultiTaskFitTransformRegressor(
    n_tasks=1,
    n_features=[23, 23],
    learning_rate=0.001,
    learning_rate=rate,
    momentum=.8,
    batch_size=25,
    weight_init_stddevs=[1 / np.sqrt(400), 1 / np.sqrt(100), 1 / np.sqrt(100)],
@@ -31,19 +34,19 @@ model = dc.models.MultiTaskFitTransformRegressor(
    fit_transformers=fit_transformers,
    n_evals=10,
    seed=123)
#model.restore()

# Fit trained model
model.fit(train_dataset, nb_epoch=50)
model.save()
model.fit(train_dataset, nb_epoch=3000)

train_scores = model.evaluate(train_dataset, regression_metric, transformers)
train_scores = model.evaluate(train_dataset, metric, transformers)
print("Train scores [kcal/mol]")
print(train_scores)

valid_scores = model.evaluate(valid_dataset, regression_metric, transformers)
valid_scores = model.evaluate(valid_dataset, metric, transformers)
print("Valid scores [kcal/mol]")
print(valid_scores)

test_scores = model.evaluate(test_dataset, regression_metric, transformers)
test_scores = model.evaluate(test_dataset, metric, transformers)
print("Test scores [kcal/mol]")
print(test_scores)