Graph Neural Networks with TensorFlow and Keras. Focused on Molecular Machine Learning.

Quick start

Benchmark the performance of MolGraph here, and implement a complete model pipeline with MolGraph here.

Highlights

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([
    layers.GNNInput(type_spec=g.spec),
    layers.GATv2Conv(units=32),
    layers.GATv2Conv(units=32),
    layers.Readout(),
    keras.layers.Dense(units=1),
])

pred = model(g)

# Save and load Keras model
model.save('/tmp/gatv2_model.keras')
loaded_model = keras.models.load_model('/tmp/gatv2_model.keras')
loaded_pred = loaded_model(g)
assert pred == loaded_pred

Combine outputs of GNN layers to improve predictive performance:

model = keras.Sequential([
    layers.GNNInput(type_spec=g.spec),
    layers.GNN([
        layers.FeatureProjection(units=32),
        layers.GINConv(units=32),
        layers.GINConv(units=32),
        layers.GINConv(units=32),
    ]),
    layers.Readout(),
    keras.layers.Dense(units=128),
    keras.layers.Dense(units=1),
])

model.summary()

Installation

For CPU users:

pip install molgraph

For GPU users:

pip install molgraph[gpu]

Implementations

• Tensors
  • Graph tensor
    • A composite tensor holding graph data; compatible with tf.data.Dataset, keras.Sequential and much more.
• Layers
  • Graph convolutional layers
  • Graph attentional layers
  • Graph message passing layers
  • Graph readout layers
  • Preprocessing layers
  • Postprocessing layers
  • Positional encoding layers
• Models
  • Graph neural networks
  • Saliency mapping

Overview

Documentation

See readthedocs

Papers

• MolGraph: a Python package for the implementation of molecular graphs and graph neural networks with TensorFlow and Keras
• A hands-on tutorial on quantitative structure-activity relationships using fully expressive graph neural networks