Rethinking pooling in graph neural networks

This repo contains code for Rethinking pooling in graph neural networks (NeurIPS 2020).

In our work, we build upon representative GNNs (Graclus, DiffPool, GMN, MinCutPool) and introduce variants that challenge the need for locality-preserving representations, either using randomization or clustering on the complement graph. Results show that using these variants does not result in any decrease in performance.

Authors:

• Diego Mesquita (diego.mesquita@aalto.fi) --- corresponding author.
• Amauri H. Souza
• Sami Kaski

Requirements

We have built this implementation using:

• PyTorch v1.4.0
• PyTorch Geometric v1.4.3
• OGB v1.2.1

Usage

There is a folder and a main file for each model (diffpool, gmn, graclus and mincut). As an example, the code below runs Graclus and its non-local variant (Complement) on NCI1:

$ python main_graclus.py --dataset NCI1 --reproduce --pooling_type graclus
$ python main_graclus.py --dataset NCI1 --reproduce --pooling_type complement

Each main file includes optional arguments.

Cite

@inproceedings{rethinkpooling2020,
  title={Rethinking pooling in Graph Neural Networks},
  author={D. Mesquita and A. H. Souza and S. Kaski},
  booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
  year={2020}
}

Acknowledgement

To avoid implementation biases, we have closely followed available code. In doing so, we have either re-used or modified code from gnn-comparison, pytorch-geometric, and benchmarking-gnns.