Google Open-Sources Its TensorFlow GNN Framework to Encourage Graph Neural Network Productization and Experimentation

Graph Neural Networks (GNNs) that operate on graph-based data bring multimodal capabilities to machine learning models and have practical applications in areas as diverse as the modelling of physics systems, learning molecular fingerprints, predicting protein interfaces, classifying social networks and more. A key component for pushing GNN development is better software frameworks for learning from graph-structured data.

In the new paper TF-GNN: Graph Neural Networks in TensorFlow, a research team from Google Core ML, Google Research, and DeepMind open-sources the TensorFlow GNN (TF-GNN) scalable library, which leverages heterogeneous relational data to create GNN models and enable GNN training and inference on arbitrary graph-structured data.

The research team summarizes their main contributions as follows:

1. We present TF-GNN, an open-source Python library, to create graph neural network models that can leverage heterogeneous relational data.
2. TF-GNN enables training and inference of Graph Neural Networks (GNNs) on arbitrary graph-structured data.
3. TF-GNN’s four API levels allow developers of all skill levels access to powerful GNN models.
4. As a native citizen of the TensorFlow ecosystem, TF-GNN shares its benefits, including pre-trained models for various modalities (e.g., an NLP model) and support for fast mathematical hardware such as Tensor Processing Units (TPUs).

TF-GNN includes four API components of varying abstraction levels to assist developers with varying machine learning expertise in creating graph models: 1) a data level for representing heterogeneous graphs and loading them into TensorFlow, which will appeal to proficient users; 2) a data exchange level for sending information between its nodes, edges, and the graph context, aimed at intermediate users; 3) a model level that offers trainable transformations of the data exchanged across the graphs; and 4) a minimal-code experience level for beginners, where an “Orchestrator” toolkit that includes popular graph learning objectives, distributed training capabilities and accelerator support — and can handle some of the vexing TensorFlow idiosyncrasies — enables simpler data input, feature processing, graph objectives, training and validation.

TF-GNN models are currently being used by many Google teams, and the company hopes the library’s open-sourcing will facilitate the creation of GNNs for developers of all levels and push the industrial adaptation of these promising models at more organizations.

The paper TF-GNN: Graph Neural Networks in TensorFlow is on arXiv.

---

Author: Hecate He | Editor: Michael Sarazen

---

We know you don’t want to miss any news or research breakthroughs. Subscribe to our popular newsletter Synced Global AI Weekly to get weekly AI updates.