Keras Graph Convolution Neural Networks

A set of layers for graph convolutions in TensorFlow Keras that use RaggedTensors.

• General
• Requirements
• Installation
• Documentation
• Implementation details
• Literature
• Datasets
• Examples
• Issues
• Citing
• References

General

The package in kgcnn contains several layer classes to build up graph convolution models. Some models are given as an example. A documentation is generated in docs. Any comments, suggestions or help are very welcome!

Requirements

For kgcnn, usually the latest version of tensorflow is required, but is listed as extra requirements in the setup.py for simplicity. Additional python packages are placed in the setup.py requirements and are installed automatically. Packages which must be installed manually for full functionality:

• tensorflow>=2.4.1
• rdkit>=2020.03.4
• openbabel>=3.0.1
• pymatgen>=??.??.??

Installation

Clone repository https://github.com/aimat-lab/gcnn_keras and install with editable mode:

pip install -e ./gcnn_keras

or latest release via Python Package Index.

pip install kgcnn

Documentation

Auto-documentation is generated at https://kgcnn.readthedocs.io/en/latest/index.html .

Implementation details

Representation

The most frequent usage for graph convolutions is either node or graph classification. As for their size, either a single large graph, e.g. citation network or small (batched) graphs like molecules have to be considered. Graphs can be represented by an index list of connections plus feature information. Typical quantities in tensor format to describe a graph are listed below.

• nodes: Node-list of shape (batch, [N], F) where N is the number of nodes and F is the node feature dimension.
• edges: Edge-list of shape (batch, [M], F) where M is the number of edges and F is the edge feature dimension.
• indices: Connection-list of shape (batch, [M], 2) where M is the number of edges. The indices denote a connection of incoming or receiving node i and outgoing or sending node j as (i, j).
• state: Graph state information of shape (batch, F) where F denotes the feature dimension.

A major issue for graphs is their flexible size and shape, when using mini-batches. Here, for a graph implementation in the spirit of keras, the batch dimension should be kept also in between layers. This is realized by using RaggedTensors.

Input

Graph tensors for edge-indices or attributes for multiple graphs is passed to the model in form of ragged tensors of shape (batch, None, Dim) where Dim denotes a fixed feature or index dimension. Such a ragged tensor has ragged_rank=1 with one ragged dimension indicated by None and is build from a value plus partition tensor. For example, the graph structure is represented by an index-list of shape (batch, None, 2) with index of incoming or receiving node i and outgoing or sending node j as (i, j). Note, an additional edge with (j, i) is required for undirected graphs. A ragged constant can be easily created and passed to a model:

import tensorflow as tf
import numpy as np
idx = [[[0, 1], [1, 0]], [[0, 1], [1, 2], [2, 0]], [[0, 0]]]  # batch_size=3
# Get ragged tensor of shape (3, None, 2)
print(tf.ragged.constant(idx, ragged_rank=1, inner_shape=(2, )).shape)
print(tf.RaggedTensor.from_row_lengths(np.concatenate(idx), [len(i) for i in idx]).shape) 

Model

Models can be set up in a functional way. Example message passing from fundamental operations:

import tensorflow.keras as ks
from kgcnn.layers.gather import GatherNodes
from kgcnn.layers.modules import DenseEmbedding, LazyConcatenate  # ragged support
from kgcnn.layers.pooling import PoolingLocalMessages, PoolingNodes

n = ks.layers.Input(shape=(None, 3), name='node_input', dtype="float32", ragged=True)
ei = ks.layers.Input(shape=(None, 2), name='edge_index_input', dtype="int64", ragged=True)

n_in_out = GatherNodes()([n, ei])
node_messages = DenseEmbedding(10, activation='relu')(n_in_out)
node_updates = PoolingLocalMessages()([n, node_messages, ei])
n_node_updates = LazyConcatenate(axis=-1)([n, node_updates])
n_embedd = DenseEmbedding(1)(n_node_updates)
g_embedd = PoolingNodes()(n_embedd)

message_passing = ks.models.Model(inputs=[n, ei], outputs=g_embedd)

or via sub-classing of the message passing base layer. Where only message_function and update_nodes must be implemented:

from kgcnn.layers.conv.message import MessagePassingBase
from kgcnn.layers.modules import DenseEmbedding, LazyAdd

class MyMessageNN(MessagePassingBase):
  def __init__(self, units, **kwargs):
    super(MyMessageNN, self).__init__(**kwargs)
    self.dense = DenseEmbedding(units)
    self.add = LazyAdd(axis=-1)
  
  def message_function(self, inputs, **kwargs):
    n_in, n_out, edges = inputs
    return self.dense(n_out)
  
  def update_nodes(self, inputs, **kwargs):
    nodes, nodes_update = inputs
    return self.add([nodes, nodes_update])

Literature

A version of the following models is implemented in literature:

• GCN: Semi-Supervised Classification with Graph Convolutional Networks by Kipf et al. (2016)
• INorp: Interaction Networks for Learning about Objects,Relations and Physics by Battaglia et al. (2016)
• Megnet: Graph Networks as a Universal Machine Learning Framework for Molecules and Crystals by Chen et al. (2019)
• NMPN: Neural Message Passing for Quantum Chemistry by Gilmer et al. (2017)
• Schnet: SchNet – A deep learning architecture for molecules and materials by Schütt et al. (2017)
• Unet: Graph U-Nets by H. Gao and S. Ji (2019)
• GNNExplainer: GNNExplainer: Generating Explanations for Graph Neural Networks by Ying et al. (2019)
• GraphSAGE: Inductive Representation Learning on Large Graphs by Hamilton et al. (2017)
• GAT: Graph Attention Networks by Veličković et al. (2018)
• GATv2: How Attentive are Graph Attention Networks? by Brody et al. (2021)
• DimeNetPP: Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules by Klicpera et al. (2020)
• AttentiveFP: Pushing the Boundaries of Molecular Representation for Drug Discovery with the Graph Attention Mechanism by Xiong et al. (2019)
• GIN: How Powerful are Graph Neural Networks? by Xu et al. (2019)
• PAiNN: Equivariant message passing for the prediction of tensorial properties and molecular spectra by Schütt et al. (2020)
• DMPNN: Analyzing Learned Molecular Representations for Property Prediction by Yang et al. (2019)

Datasets

The base class is MemoryGraphDataset which holds a lists of graph properties in tensor-like numpy arrays. Each property that starts with node_, edge_, graph_ holds a list with length of the dataset that fit into memory. Furthermore, file information on disk can be provided in the constructor, that points to a data_directory for the dataset.

├── data_directory
    ├── file_directory
    │   ├── *.*
    │   └── ... 
    ├── file_name
    └── dataset_name.pickle

Create and store a general dataset via:

from kgcnn.data.base import MemoryGraphDataset
import numpy as np
dataset = MemoryGraphDataset(data_directory="ExampleDir", dataset_name="Example")
dataset.assign_property("edge_indices", [np.array([[1, 0], [0, 1]])])
dataset.assign_property("edge_labels", [np.array([[0], [1]])])

The subclasses QMDataset, MoleculeNetDataset and GraphTUDataset further have functions required for the specific dataset to convert and load files such as ".txt", ".sdf", ".xyz" etc. via prepare_data() and read_in_memory().

from kgcnn.data.qm import QMDataset
dataset = QMDataset(data_directory="ExampleDir", dataset_name="methane", 
                    file_name="geom.xyz", file_directory=None)
dataset.prepare_data()  # Also make .sdf
dataset.read_in_memory()

In data.datasets there are graph learning datasets as subclasses which are being downloaded from e.g. TUDatasets or MoleculeNet and directly processed and loaded. They are downloaded and stored at ~/.kgcnn/datasets on your computer.

from kgcnn.data.datasets.MUTAGDataset import MUTAGDataset
dataset = MUTAGDataset()
print(dataset.edge_indices[0])

Examples

A set of example training can be found in training. Training scripts are configurable with a hyperparameter config file and command line arguments regarding model and dataset.

Issues

Some known issues to be aware of, if using and making new models or layers with kgcnn.

• RaggedTensor can not yet be used as a keras model output (https://github.com/tensorflow/tensorflow/issues/42320), which means only padded tensors can be used for batched node embedding tasks.
• Using RaggedTensor's for arbitrary ragged rank apart from kgcnn.layers.modules can cause significant performance decrease. This is due to shape check during add, multiply or concatenate (we think). We therefore use lazy add and concat in the kgcnn.layers.modules layers or directly operate on the value tensor for possible rank.
• With tensorflow version <=2.5 there is a problem with numpy version >=1.20 also affect kgcnn (https://github.com/tensorflow/tensorflow/issues/47691)

Citing

If you want to cite this repo, refer to our paper:

@article{REISER2021100095,
title = {Graph neural networks in TensorFlow-Keras with RaggedTensor representation (kgcnn)},
journal = {Software Impacts},
pages = {100095},
year = {2021},
issn = {2665-9638},
doi = {https://doi.org/10.1016/j.simpa.2021.100095},
url = {https://www.sciencedirect.com/science/article/pii/S266596382100035X},
author = {Patrick Reiser and Andre Eberhard and Pascal Friederich}
}

References

• https://www.tensorflow.org/api_docs/python/tf/RaggedTensor