Repository files navigation

TensorFlow GNN is a library to buildGraph Neural Networks on the TensorFlow platform.It provides...

• atfgnn.GraphTensor type to representgraphs with aheterogeneous schema, that is,multiple types of nodes and edges;
• tools fordata preparation,notably agraph samplerto convert a huge database into a stream of reasonably-sized subgraphs fortraining and inference;
• a collection ofready-to-use modelsand Keras layers to do your ownGNN modeling;
• a high-level API for trainingorchestration.

This library is an OSS port of a Google-internal library used in a broad varietyof contexts, on homogeneous and heterogeneous graphs, and in conjunction withother scalable graph mining tools.

For background, please see ourblog postand theTF-GNN paper (full citation below).

Google Colab lets you run TF-GNN demos from your browser, no installationrequired:

• Molecular GraphClassificationwith the MUTAG dataset.
• Solving OGBN-MAGend-to-endtrains a model on heterogeneous sampled subgraphs from the popularOGBN-MAG benchmark.
• Learning shortest paths withGraphNetworksdemonstrates an advanced Encoder/Process/Decoder architecture for predictingthe edges of a shortest path.

For all colabs and user guides, please see theDocumentation overviewpage, which also links to theAPI docs.

The latest stable release of TensorFlow GNN is available from

pip install tensorflow-gnn

For installation from source, see ourDeveloperGuide.

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-keras and setTF_USE_LEGACY_KERAS=1,see ourKeras version guide for details.
• Apache Beam for distributed graph sampling.
• For some tests or scripts that requires tensorflow.lite it is required toinstall ai-edge-litert by usingpip install ai-edge-litert

TF-GNN is developed and tested on Linux. Running on other platforms supportedby TensorFlow may be possible.

When referencing this library in a paper, please cite theTF-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},}