Graph neural networks for molecular machine learning. Implemented and compatible with TensorFlow and Keras.
Graph Neural Networks with TensorFlow and Keras. Focused on Molecular Machine Learning.
Currently, Keras 3 does not support extension types. As soon as it does, it is hoped that MolGraph will migrate to Keras 3.
Build a Graph Neural Network with Keras' Sequential API:
from molgraph import GraphTensor
from molgraph import layers
from tensorflow import keras
g = GraphTensor(node_feature=[[4.], [2.]], edge_src=[0], edge_dst=[1])
model = keras.Sequential([
keras.layers.Input(type_spec=g.spec),
layers.GINConv(units=32),
layers.GINConv(units=32),
layers.Readout(),
keras.layers.Dense(units=1),
])
pred = model(g)
See arXiv
See readthedocs
For CPU users:
pip install molgraph
For GPU users:
pip install molgraph[gpu]
Now run your first program with MolGraph:
from tensorflow import keras
from molgraph import chemistry
from molgraph import layers
from molgraph import models
# Obtain dataset, specifically ESOL
esol = chemistry.datasets.get('esol')
# Define molecular graph encoder
atom_encoder = chemistry.Featurizer([
chemistry.features.Symbol(),
chemistry.features.Hybridization(),
# ...
])
bond_encoder = chemistry.Featurizer([
chemistry.features.BondType(),
# ...
])
encoder = chemistry.MolecularGraphEncoder(atom_encoder, bond_encoder)
# Obtain graphs and associated labels
x_train = encoder(esol['train']['x'])
y_train = esol['train']['y']
x_test = encoder(esol['test']['x'])
y_test = esol['test']['y']
# Build model via Keras API
gnn_model = keras.Sequential([
keras.layers.Input(type_spec=x_train.spec),
layers.GATConv(units=32, name='gat_conv_1'),
layers.GATConv(units=32, name='gat_conv_2'),
layers.Readout(),
keras.layers.Dense(units=1024, activation='relu'),
keras.layers.Dense(units=y_train.shape[-1])
])
# Compile, fit and evaluate
gnn_model.compile(optimizer='adam', loss='mae')
gnn_model.fit(x_train, y_train, epochs=50)
scores = gnn_model.evaluate(x_test, y_test)
# Compute gradient activation maps
gam_model = models.GradientActivationMapping(
model=gnn_model, layer_names=['gat_conv_1', 'gat_conv_2'])
maps = gam_model(x_train.separate())