DEV Community

ekka
ekka

Posted on • Edited on

ASAP: Pooling layer for Graph Neural Nets

My Final Project

I am interested in Deep Learning and have been trying to keep pace with the rapidly evolving field. In my last internship, I worked on designing new pooling layers for Graph Neural Networks. My work finds application in tasks involving graphs, e.g., from finding toxicity of Molecule/Drug to learning attributes about users on a social network.

Link to Code

GitHub logo malllabiisc / ASAP

AAAI 2020 - ASAP: Adaptive Structure Aware Pooling for Learning Hierarchical Graph Representations

ASAP: Adaptive Structure Aware Pooling for Learning Hierarchical Graph Representations

Conference Paper

Source code for AAAI 2020 paper: ASAP: Adaptive Structure Aware Pooling for Learning Hierarchical Graph Representation

Overview of ASAP: ASAP initially considers all possible local clusters with a fixed receptive field for a given input graph. It then computes the cluster membership of the nodes using an attention mechanism. These clusters are then scored using a GNN. Further, a fraction of the top scoring clusters are selected as nodes in the pooled graph and new edge weights are computed between neighboring clusters. Please refer to Section 4 of the paper for details.

File Descriptions

  • main.py - contains the driver code for the whole project
  • asap_pool.py - source code for ASAP pooling operator proposed in the paper
  • le_conv.py - source code for LEConv GNN used in the paper
  • asap_pool_model.py - a network which uses ASAP pooling as pooling operator

Dependencies

  • Python…

How I built it (what's the stack? did I run into issues or discover something new along the way?)

The Machine Learning domain is quite lucky in the sense that the tools needed for research and deployment are open-sourced. Major libraries like Pytorch, Tensorflow, and sklearn are all maintained actively by the open-source community.

My research stack primarily includes Pytorch library. For the purpose of working in the graph domain, I had to use Pytorch_Geometric which is a library made on top of Pytorch. Its maintainer has done tremendous work in building and keeping it up to date. So whenever I ran into problems or implemented a new feature for my research work, I made it sure to package it properly and merge it to Pytorch_Geometric so that others can use it too.

For managing sessions on the GPU server, I used tmux client so that if my personal machine turns off, the session is still alive on the server to which I can connect later.

To collaborate with my mentor I used Github.

Additional Thoughts / Feelings / Stories

It took 3.5 months of consistent work and frequent iterations to submit our work which was accepted as a long paper at AAAI 2020 (a top-tier ML conference). This would not be possible without all the open-source tools available to me for research. I am grateful to the open-source community for sharing their work. Realizing the impact of open-source tools on my workflow motivates me to spend my free time in contributing to open-source projects.

Top comments (0)