Torch Geometric Pool: the Pytorch library for pooling in Graph Neural Networks

Filippo Maria Bianchi, Carlo Abate, Ivan Marisca

Keywords: graph neural networks, graph pooling, graph coarsening

TL;DR: The graph pooling library made for PyTorch Geometric

Abstract: We introduce Torch Geometric Pool (tgp), a library for hierarchical pooling in Graph Neural Networks. Built upon PyTorch Geometric, tgp provides a wide variety of pooling operators, unified under a consistent API and a modular design based on the Select-Reduce-Connect-Lift framework. The library emphasizes usability and extensibility, and includes features like precomputed pooling, which significantly accelerate training for deterministic operators. In this paper, we present tgp's architecture and systematically compare the performance of the implemented poolers in different downstream tasks. The results, showing that the choice of the optimal pooling operator depends on tasks and data at hand, support the need for a library that enables fast prototyping.

Primary Area: infrastructure, software libraries, hardware, systems, etc.

Submission Number: 11041