PyTorch Graph Networks

PyTorch implementation of DeepMind Graph Nets. The original code depends on Tensorflow and Sonnet.

This implementation is based on PyTorch Geometric which is a geometric deep learning extension library for PyTorch

Graph Networks

Graph Networks are a general framework that generalizes graph neural networks. It unifies Message Passing Neural Networks (MPNNs) and Non-Local Neural Networks (NLNNs), as well as other variants like Interaction Networks (INs) orRelation Networks (RNs).

You can have a look at Graph Networks in their arXiV paper: Battaglia, Peter W., et al. "Relational inductive biases, deep learning, and graph networks." arXiv preprint arXiv:1806.01261 (2018)

Available Models

The following models are available:

• Interaction Network
• Graph Independent

You can also build your own models using the Blocks:

• Node Model
• Edge Model

Requirements

PyTorch 1.8.0 and PyTorch Geometric.

Example

We provide an example that tests the output against DeepMind's graph_nets.