TensorFlow GNN is a library to build Graph Neural Networks on the TensorFlow platform. It provides...
- a
tfgnn.GraphTensortype to represent graphs with a heterogeneous schema, that is, multiple types of nodes and edges; - tools for data preparation, notably a graph sampler to convert a huge database into a stream of reasonably-sized subgraphs for training and inference;
- a collection of ready-to-use models and Keras layers to do your own GNN modeling;
- a high-level API for training orchestration.
This library is an OSS port of a Google-internal library used in a broad variety of contexts, on homogeneous and heterogeneous graphs, and in conjunction with other scalable graph mining tools.
For background, please see our blog post and the TF-GNN paper (full citation below).
Google Colab lets you run TF-GNN demos from your browser, no installation required:
- Molecular Graph Classification with the MUTAG dataset.
- Solving OGBN-MAG end-to-end trains a model on heterogeneous sampled subgraphs from the popular OGBN-MAG benchmark.
- Learning shortest paths with GraphNetworks demonstrates an advanced Encoder/Process/Decoder architecture for predicting the edges of a shortest path.
For all colabs and user guides, please see the Documentation overview page, which also links to the API docs.
The latest stable release of TensorFlow GNN is available from
pip install tensorflow-gnn
For installation from source, see our Developer Guide.
Key platform requirements:
- TensorFlow 2.12 or higher, and any GPU drivers it needs [instructions].
- Keras v2, as traditionally included with TensorFlow 2.x.
TF-GNN does not work with the new multi-backend Keras v3.
Users of TF2.16+ must alsopip install tf-kerasand set TF_USE_LEGACY_KERAS=1, see our Keras version guide for details. - Apache Beam for distributed graph sampling.
- For some tests or scripts that requires tensorflow.lite it is required to
install ai-edge-litert by using
pip install ai-edge-litert
TF-GNN is developed and tested on Linux. Running on other platforms supported by TensorFlow may be possible.
When referencing this library in a paper, please cite the TF-GNN paper:
@article{tfgnn,
author = {Oleksandr Ferludin and Arno Eigenwillig and Martin Blais and
Dustin Zelle and Jan Pfeifer and Alvaro Sanchez{-}Gonzalez and
Wai Lok Sibon Li and Sami Abu{-}El{-}Haija and Peter Battaglia and
Neslihan Bulut and Jonathan Halcrow and
Filipe Miguel Gon{\c{c}}alves de Almeida and Pedro Gonnet and
Liangze Jiang and Parth Kothari and Silvio Lattanzi and
Andr{\'{e}} Linhares and Brandon Mayer and Vahab Mirrokni and
John Palowitch and Mihir Paradkar and Jennifer She and
Anton Tsitsulin and Kevin Villela and Lisa Wang and David Wong and
Bryan Perozzi},
title = {{TF-GNN:} Graph Neural Networks in TensorFlow},
journal = {CoRR},
volume = {abs/2207.03522},
year = {2023},
url = {http://arxiv.org/abs/2207.03522},
}