Google releases TF-GNN for creating graph neural networks in TensorFlow

Let the OSS Enterprise newsletter guide your open source journey! Sign up here.

Google today released TensorFlow Graph Neural Networks (TF-GNN) in alpha, a library designed to make it easier to work with graph structured data using TensorFlow, its machine learning framework. Used in production at Google for spam and anomaly detection, traffic estimation, and YouTube content labeling, Google says that TF-GNN is designed to “encourage collaborations with researchers in industry.”

Graphs are a set of objects, places, or people and the connections between them. A graph represents the relations (edges) between a collection of entities (nodes or vertices), all of which can store data. Directionality can be ascribed to the edges to describe information, traffic flow, and more.

More often than not, the data in machine learning problems is structured or relational and thus can be described with a graph. Fundamental research on GNNs is decades old, but recent advances have led to great achievements in many domains, like modeling the transition of glass from a liquid to a solid and predicting pedestrian, cyclist, and driver behavior on the road.

Above: Graphs can model the relationships between many different types of data, including web pages (left), social connections (center), or molecules (right).

Image Credit: Google

Indeed, GNNs can be used to answer questions about multiple characteristics of graphs. By working at the graph level, they can try to predict aspects of the entire graph, for example identifying the presence of certain “shapes” like circles in a graph that might represent close social relationships. GNNs can also be used on node-level tasks to classify the nodes of a graph or at the edge level to discover connections between entities.

TF-GNN

TF-GNN provides building blocks for implementing GNN models in TensorFlow. Beyond the modeling APIs, the library also delivers tooling around the task of working with graph data, including a data-handling pipeline and example models.

Also included with TF-GNN is an API to create GNN models that can be composed with other types of AI models. In addition to this, TF-GNN ships with a schema to declare the topology of a graph (and tools to validate it), helping to describe the shape of training data.

“Graphs are all around us, in the real world and in our engineered systems … In particular, given the myriad types of data at Google, our library was designed with heterogeneous graphs in mind,” Google’s Sibon Li, Jan Pfeifer, Bryan Perozzi, and Douglas Yarrington wrote in the blog post introducing TF-GNN.

TF-GNN adds to Google’s growing collection of TensorFlow libraries, which spans TensorFlow Privacy, TensorFlow Federated, and TensorFlow.Text. More recently, the company open-sourced TensorFlow Similarity, which trains models that search for related items — for example, finding similar-looking clothes and identifying currently playing songs.