The TensorFlow GNN library makes it easy to build Graph Neural Networks, that is, neural networks on graph data (nodes and edges with arbitrary features). It provides TensorFlow code for building GNN models as well as tools for preparing their input data and running the training.
Throughout, TF-GNN supports heterogeneous graphs, that is, graphs consisting of multiple sets of nodes and multiple sets of edges, each with their own set of features. These come up naturally when modeling different types of objects (nodes) and their different types of relations (edges).
Start with our introductory guides:
-
Introduction to Graph Neural Networks. This page introduces the concept of graph neural networks with a focus on their application at scale.
-
The GraphTensor type. This page introduces the
tfgnn.GraphTensorclass, which defines our representation of graph data in TensorFlow. We recommend that every user of our library understands its basic data model. -
Describing your graph. This page explains how to declare the node sets and edge sets of your graph, including their respective features, with the
GraphSchemaprotocol message. This defines the interface between data preparation (which creates such graphs) and the GNN model written in TensorFlow (which consumes these graphs as training data). -
Data preparation and sampling. Training data for GNN models are graphs. When operating on very large graphs, we produce local subgraph samples from the full dataset which are serialized along with their features and streamed for training. This document describes the sampler tool as well as the data format we support for storing large graphs.
-
The TF-GNN Runner lets you train GNN models on the prepared input data for a variety of tasks (e.g., node prediction). We recommend using the Runner to get started quickly with a first model for the data at hand, and then customize it as needed.
The following docs go deeper into particular topics.
-
The Input pipeline guide explains how to set up a
tf.data.Datasetfor bulk input of the training and validation datasets produced by the data preparation step. The TF-GNN Runner already takes care of this for its users. -
TF-GNN modeling explains how to build a Graph Neural Network with TensorFlow and Keras, using the GraphTensor data from the previous steps. The TF-GNN library provides both a collection of standard models and a toolbox for writing your own. Users of the TF-GNN Runner are encouraged to consult this page to define custom models in the Runner.
These Colab notebooks run complete examples of building and training a TF-GNN model on a Google server from within your browser.
- Molecular Graph Classification trains a model for the MUTAG dataset (from the TUDatasets collection) that consists of 188 small, homogeneous graphs representing molecules. This is a good first read to get acquainted with GNNs.
- Solving OGBN-MAG end-to-end trains a model on heterogeneous sampled subgraphs from the OGBN-MAG dataset (from Stanford's Open Graph Benchmark) that contains 1 million research papers, their authors, and other relations. This is a more advanced example that introduces node classification from sampled subgraphs as well as the nuts and bolts of training in parallel on multiple accelerators (GPU, TPU).
- Learning shortest paths with GraphNetworks demonstrates an Encoder/Process/Decoder architecture for predicting the edges of a shortest path, using an Graph Network with edge states. Take a look if you are interested in advanced modeling.
TF-GNN has reference documentation for the API, extracted from the source code.
- See here for the core library
(
import tensorflow_gnn) . - See
this list
for the bundled models (
from tensorflow_gnn.models import foo).
How to contribute to the TF-GNN library.
- CONTRIBUTING.md describes the process for open-source contributions.
- The Developer guide describes how to clone our github repo and install the tools and libraries required to build and run TF-GNN code.
The following research paper describes the design of this library:
- O. Ferludin et al.: TF-GNN: Graph Neural Networks in TensorFlow, 2022.
You can find related talks at our NeurIPS 2020 Expo workshop: https://gm-neurips-2020.github.io/