A library for easy construction of message-passing networks in tensorflow keras. It is inspired largely by this DeepMind paper and the corresponding open-sourced library (original graph_nets library)

The tf_gnns library has no external dependencies except tensorflow 2.x (there is no support for tf 1.x graphs/sessions-based computation). It implements some alternative design constraints from graph_nets taking advantage of some facilities keras provides to make complex models easily and without large drops in performance.

tf_gnns is built to support arbitrary node/edge/global attributes and update functions. A set of convenience functions providing MLP construction with keras are also provided (i.e., handling input/output sizes for valid networks) that replaces sonnet.

The main motivation for this library was the absense of a relatively short and efficient implementation of GNNs that was explicitly created to take advantage of keras's functionalities. GNN implementations which take advantage of tensorflow_probability functionality are to be released in the future (as the one used in [2]).

Install with pip:

pip install tf_gnns

You may inspect some basic functionality on the following colab notebook:

(Example from the original deepmind/graph_nets library) If you are familiar with the original graph_nets library, this example will help you understand how you can transition to tf_gnns.

Sort a list of elements. This notebook and the accompanying code demonstrates how to use the Graph Nets library to learn to sort a list of elements.

A list of elements is treated as a fully connected graph between the elements. The network is trained to label the start node, and which (directed) edges correspond to the links to the next largest element, for each node.

After training, the prediction ability is tested by comparing its output to true sorted lists. Then the network's ability to generalise is tested, by using it to sort larger lists.

From some initial tests the performance of the tf_gnns library seems to be at least as good as deepmind/graph_nets when using tensor dictionaries.

The library has been used so far in the following publications:

[1] Bayesian graph neural networks for strain-based crack localization

[2] Remaining Useful Life Estimation Under Uncertainty with Causal GraphNets [3] Relational VAE: A Continuous Latent Variable Model for Graph Structured Data