def test_gat_regression():
# load datasets
featurizer = MolGraphConvFeaturizer()
tasks, dataset, transformers, metric = get_dataset('regression',
featurizer=featurizer)
# initialize models
n_tasks = len(tasks)
model = GATModel(mode='regression', n_tasks=n_tasks, batch_size=10)
# overfit test
# GAT's convergence is a little slow
model.fit(dataset, nb_epoch=300)
scores = model.evaluate(dataset, [metric], transformers)
assert scores['mean_absolute_error'] < 0.5
def test_gat_classification():
# load datasets
featurizer = MolGraphConvFeaturizer()
tasks, dataset, transformers, metric = get_dataset('classification',
featurizer=featurizer)
# initialize models
n_tasks = len(tasks)
model = GATModel(mode='classification',
n_tasks=n_tasks,
number_atom_features=30,
batch_size=10,
learning_rate=0.001)
# overfit test
model.fit(dataset, nb_epoch=100)
scores = model.evaluate(dataset, [metric], transformers)
assert scores['mean-roc_auc_score'] >= 0.85
# test on a small MoleculeNet dataset
from deepchem.molnet import load_bace_classification
tasks, all_dataset, transformers = load_bace_classification(
featurizer=featurizer)
train_set, _, _ = all_dataset
model = dc.models.GATModel(mode='classification',
n_tasks=len(tasks),
graph_attention_layers=[2],
n_attention_heads=1,
residual=False,
predictor_hidden_feats=2)
model.fit(train_set, nb_epoch=1)
def test_gat_classification():
# load datasets
featurizer = MolGraphConvFeaturizer()
tasks, dataset, transformers, metric = get_dataset('classification',
featurizer=featurizer)
# initialize models
n_tasks = len(tasks)
model = GATModel(mode='classification',
n_tasks=n_tasks,
number_atom_features=30,
batch_size=10,
learning_rate=0.001)
# overfit test
model.fit(dataset, nb_epoch=100)
scores = model.evaluate(dataset, [metric], transformers)
assert scores['mean-roc_auc_score'] >= 0.85
def test_gat_regression():
# load datasets
featurizer = MolGraphConvFeaturizer()
tasks, dataset, transformers, metric = get_dataset('regression',
featurizer=featurizer)
# initialize models
n_tasks = len(tasks)
model = GATModel(mode='regression',
n_tasks=n_tasks,
number_atom_features=30,
batch_size=10,
learning_rate=0.001)
# overfit test
model.fit(dataset, nb_epoch=400)
scores = model.evaluate(dataset, [metric], transformers)
assert scores['mean_absolute_error'] < 0.5
def test_gat_classification():
# load datasets
featurizer = MolGraphConvFeaturizer()
tasks, dataset, transformers, metric = get_dataset('regression',
featurizer=featurizer)
# initialize models
n_tasks = len(tasks)
model = GATModel(n_tasks=n_tasks,
loss=losses.L2Loss(),
batch_size=4,
learning_rate=0.001)
# overfit test
model.fit(dataset, nb_epoch=100)
scores = model.evaluate(dataset, [metric], transformers)
# TODO: check this asseration is correct or not
assert scores['mean_absolute_error'] < 1.0
def test_gat_reload():
# load datasets
featurizer = MolGraphConvFeaturizer()
tasks, dataset, transformers, metric = get_dataset('classification',
featurizer=featurizer)
# initialize models
n_tasks = len(tasks)
model_dir = tempfile.mkdtemp()
model = GATModel(mode='classification',
n_tasks=n_tasks,
number_atom_features=30,
model_dir=model_dir,
batch_size=10,
learning_rate=0.001)
model.fit(dataset, nb_epoch=100)
scores = model.evaluate(dataset, [metric], transformers)
assert scores['mean-roc_auc_score'] >= 0.85
reloaded_model = GATModel(mode='classification',
n_tasks=n_tasks,
number_atom_features=30,
model_dir=model_dir,
batch_size=10,
learning_rate=0.001)
reloaded_model.restore()
pred_mols = ["CCCC", "CCCCCO", "CCCCC"]
X_pred = featurizer(pred_mols)
random_dataset = dc.data.NumpyDataset(X_pred)
original_pred = model.predict(random_dataset)
reload_pred = reloaded_model.predict(random_dataset)
assert np.all(original_pred == reload_pred)