def test_call(self, node_size, edge_size, globals_size, y_size, model_type,
n_layers):
"""Check that shapes for call are correct."""
graphs, n_graphs = self._setup_graphs()
model = experiments.GNN(node_size, edge_size, globals_size, y_size,
model_type, 'relu', templates.TargetType.globals,
n_layers)
y_out = model(graphs).numpy()
self.assertEqual((n_graphs, y_size), y_out.shape)
def test_predict(self, node_size, edge_size, globals_size, y_size, model_type,
n_layers):
"""Check that predict is 1-D."""
graphs, n_graphs = self._setup_graphs()
model = experiments.GNN(node_size, edge_size, globals_size, y_size,
model_type, 'relu', templates.TargetType.globals,
n_layers)
model(graphs)
y_out = model.predict(graphs).numpy()
self.assertEqual((n_graphs,), y_out.shape)
def test_get_graph_embedding(self, node_size, edge_size, globals_size, y_size,
model_type, n_layers):
"""Check that shapes for get_embedding are correct."""
graphs, n_graphs = self._setup_graphs()
model = experiments.GNN(node_size, edge_size, globals_size, y_size,
model_type, 'relu', templates.TargetType.globals,
n_layers)
model(graphs)
emb = model.get_graph_embedding(graphs).numpy()
self.assertEqual((n_graphs, globals_size), emb.shape)
def test_get_prediction_weights(self, node_size, edge_size, globals_size,
y_size, model_type, n_layers):
"""Check that we have enough weights as globals vector dim."""
graphs, _ = self._setup_graphs()
model = experiments.GNN(node_size, edge_size, globals_size, y_size,
model_type, 'relu', templates.TargetType.globals,
n_layers)
model(graphs)
w = model.get_prediction_weights()
self.assertEqual(w.shape, globals_size)
def test_get_gradient(self, node_size, edge_size, globals_size, y_size,
model_type, n_layers):
"""Check that shapes for gradients are correct."""
graphs, _ = self._setup_graphs()
model = experiments.GNN(node_size, edge_size, globals_size, y_size,
model_type, 'relu', templates.TargetType.globals,
n_layers)
model(graphs)
node_grads, edge_grads = model.get_gradient(graphs)
self.assertEqual(graphs.nodes.shape, node_grads.shape)
self.assertEqual(graphs.edges.shape, edge_grads.shape)
def _setup_graphs_model(self):
"""Setup graphs and smiles if needed."""
tensorizer = featurization.MolTensorizer()
smiles = ['CO', 'CCC', 'CN1C=NC2=C1C(=O)N(C(=O)N2C)C']
graphs = graph_utils.smiles_to_graphs_tuple(smiles, tensorizer)
# Fix seed so that initialization is deterministic.
tf.random.set_seed(0)
model = experiments.GNN(5, 3, 10, 1, models.BlockType('gcn'), 'relu',
templates.TargetType.globals, 3)
model(graphs)
return graphs, model, tensorizer
def test_get_gap_activations(self, node_size, edge_size, globals_size, y_size,
model_type, n_layers):
"""Check that shapes for gap activation are like global vectors."""
graphs, _ = self._setup_graphs()
n_nodes, n_edges = graphs.nodes.shape[0], graphs.edges.shape[0]
model = experiments.GNN(node_size, edge_size, globals_size, y_size,
model_type, 'relu', templates.TargetType.globals,
n_layers)
model(graphs)
node_act, edge_act = model.get_gap_activations(graphs)
self.assertEqual((n_nodes, globals_size), node_act.shape)
self.assertEqual((n_edges, globals_size), edge_act.shape)
def _setup_experiment(self):
"""Setup graphs and smiles if needed."""
smiles = ['CO', 'CCC', 'CN1C=NC2=C1C(=O)N(C(=O)N2C)C']
n = len(smiles)
smiles_to_mol = functools.partial(
featurization.smiles_to_mol, infer_hydrogens=True)
tensorizer = featurization.MolTensorizer(preprocess_fn=smiles_to_mol)
train_index, test_index = np.arange(n - 1), np.arange(n - 1, n)
mol_list = [smiles_to_mol(smi) for smi in smiles]
x = graph_utils.smiles_to_graphs_tuple(smiles, tensorizer)
task = tasks.get_task(tasks.Task.crippen)
y = task.get_true_predictions(mol_list)
atts = task.get_true_attributions(mol_list)
exp = experiments.ExperimentData.from_data_and_splits(
x, y, atts, train_index, test_index)
model = experiments.GNN(5, 3, 10, 1, models.BlockType.gcn, 'relu',
templates.TargetType.globals, 2)
model(x)
method = techniques.CAM()
return exp, model, task, method
def test_get_intermediate_activations_gradients(self, node_size, edge_size,
globals_size, y_size,
model_type, n_layers):
"""Check shapes for actvations/gradients, check that predictions match."""
graphs, _ = self._setup_graphs()
n_nodes, n_edges = graphs.nodes.shape[0], graphs.edges.shape[0]
model = experiments.GNN(node_size, edge_size, globals_size, y_size,
model_type, 'relu', templates.TargetType.globals,
n_layers)
model(graphs)
acts, grads, y = model.get_intermediate_activations_gradients(graphs)
y_expected = model.predict(graphs)
self.assertTrue(np.allclose(y, y_expected))
self.assertLen(acts, n_layers)
self.assertLen(grads, n_layers)
for (act_nodes, act_edges), (grad_nodes, grad_edges) in zip(acts, grads):
self.assertEqual((n_nodes, node_size), act_nodes.shape)
self.assertEqual((n_edges, edge_size), act_edges.shape)
self.assertEqual((n_nodes, node_size), grad_nodes.shape)
self.assertEqual((n_edges, edge_size), grad_edges.shape)