Implementing graph neural networks with TensorFlow-Keras

TL;DR

This paper introduces a flexible implementation of graph neural network layers using TensorFlow-Keras, enabling seamless integration and mini-batch processing with RaggedTensor, and presents the kgcnn Python package for practical use.

Contribution

It provides a new implementation of GNN layers in TensorFlow-Keras with a focus on transparency and ease of use, including a dedicated Python package.

Findings

Developed convolution and pooling layers for GNNs in TensorFlow-Keras.

Implemented mini-batch processing using RaggedTensor for graph data.

Created the kgcnn Python package for accessible GNN model development.

Abstract

Graph neural networks are a versatile machine learning architecture that received a lot of attention recently. In this technical report, we present an implementation of convolution and pooling layers for TensorFlow-Keras models, which allows a seamless and flexible integration into standard Keras layers to set up graph models in a functional way. This implies the usage of mini-batches as the first tensor dimension, which can be realized via the new RaggedTensor class of TensorFlow best suited for graphs. We developed the Keras Graph Convolutional Neural Network Python package kgcnn based on TensorFlow-Keras that provides a set of Keras layers for graph networks which focus on a transparent tensor structure passed between layers and an ease-of-use mindset.