search

google-deepmind / graph_nets

Build Graph Nets in Tensorflow
https://arxiv.org/abs/1806.01261
Apache License 2.0
5.34k stars 783 forks source link

What's the difference between graph_nets and jraph? #143

Closed cfangplus closed 3 years ago

cfangplus commented 3 years ago

hi,

The arXiv paper said that Graph networks are part of the broader family of "graph neural networks" and I have two questions as follow. 1, Could I use graph_nets to train a GNN model, such as GCN/GraphSage/GAT ? 2, I ran the code https://github.com/deepmind/deepmind-research/blob/master/ogb_lsc/mag/README.md and it seems that graph_nets is used to define graph-structured data and the GNN model is defined and trained by jraph, is that right? Or what's the real difference between graph_nets and jraph? Thx.

cfangplus commented 3 years ago

I also read some demo code from jraph and it seems that the model there obay the Alogorith1 from https://arxiv.org/abs/1806.01261.

alvarosg commented 3 years ago

Thanks for the comment:

What's the difference between graph_nets and jraph? The main difference is that jraph works with JAX, and graph_nets works with TensorFlow. We also provided an improved API for jraph which offers native support for nested node, edge and global features.

1, Could I use graph_nets to train a GNN model, such as GCN/GraphSage/GAT ? Yes, although in principle you would have to write the model you want based on the building blocks. For example, for a GCN, you can simply broadcast senders nodes to edges, and then aggregate to receivers. And for GATs, you may be able to reuse the (or parts of) the self attention module: https://github.com/deepmind/graph_nets/blob/64771dff0d74ca8e77b1f1dcd5a7d26634356d61/graph_nets/modules.py#L657

it seems that graph_nets is used to define graph-structured data and the GNN model is defined and trained by jraph, is that right?

In that project we used TensorFlow for the data pipelines (to be able to benefit from fast prefetching), and JAX for the model, hence the use of the two graph nets libraries.

Hope this helps!

cfangplus commented 3 years ago

great thx@ @alvarosg