Code for the paper Graph Element Networks: adaptive, structured computation and memory. See here for the arxiv link.

You can see the talk summarizing the paper here.

Since the paper, we have simplified the code, now relying on pytorch_geometric, which also optimizes the runtime of GNNs.

Now the code is faster and easy to customize. In particular, the GEN file contains the class that can be customized, as shown in GENSoftNN.

Despite a couple minor changes to make the code faster and runnable inside a Docker container (see the Dockerfile), we checked that results are very similar to the original ones.

You can download the image files to visualize the node optimization experiments here.

You can find the self-contained code for the scene representation experiments in the scene_representation folder folder and the Graph Element Network code in the composer and the representation function here. Note that this part of the codebase contains the original for GENs.