torch-geometric-pool 1.0.2

The Graph Pooling library for PyTorch Geometric.

The library for pooling in Graph Neural Networks

---

📚 Documentation - 🚀 Getting Started - 💻 Introductory notebooks

tgp (Torch Geometric Pool) is a library that provides a broad suite of graph pooling layers to be inserted into Graph Neural Network architectures built with PyTorch Geometric . With tgp, you can effortlessly construct hierarchical GNNs by interleaving message-passing layers with any pooling operations.

Features

• Unified API. All pooling layers in tgp are implemented following the SRC (Select, Reduce, Connect) framework, introduced in Understanding Pooling in Graph Neural Networks, which ensures a consistent API across all methods and seamless interoperability.

• All your pooling operators in one place. Choose from a variety of pooling methods, including sparse techniques like Top-K, NDPPooling, GraclusPooling, and dense methods such as DiffPool and MinCutPooling. Each operator adheres to the modular SRC framework, allowing you to quickly exchange methods within the same GNN architecture and combine them with standard message-passing layers.

• Precomputed & On-the-Fly Pooling. Accelerate training by precomputing the coarse graph (assignments and connectivity) for methods like NDPPooling or GraclusPooling. Alternatively, use on-the-fly pooling (e.g., Top-K or MinCut) that computes assignments dynamically and supports end-to-end gradient flow.

• Alias-Based Instantiation. Quickly create any pooler by name (e.g., "topk", "ndp", "diff", "mincut"). Pass a configuration dict for hyperparameters, and receive a fully initialized pooling layer that conforms to the unified SRC interface.

• Tweak and create new pooling layers. Thanks to the modular SRC framework, the components of different pooling layers in tgp can be easily combined, replaced with existing modules, or swapped for entirely new ones.

Getting Started

If you are unfamiliar with graph pooling, we recommend checking this introduction to the SRC framework and this blog for a deeper dive into pooling in GNNs. Before you dive into using tgp, we recommend browsing the Documentation to familiarize yourself with the API.

If you prefer a notebook-based introduction, check out the following tutorials:

• Basic usage of common pooling operators, including how to pass arguments via aliases and inspect intermediate outputs.

• Demonstrates how to apply precomputed pooling methods and associated data transforms for faster training.

• Deep dive into the SRC framework, showing how each component interacts and how to use the select, reduce, connect and lift operations to modify the graph topology and the graph features.

In addition, check the example folder for a collection of minimalistic python script showcasing the usage of the pooling operators of tgp in all the most common downstream tasks, such as graph classification, node classification, and node clustering.

Installation

tgp is compatible with Python>=3.9. We recommend installation on a Anaconda or Miniconda environment or a virtual env. tgp is conveniently available as a Python package on PyPI and can be easily installed using pip.

pip install torch-geometric-pool

For the latest version, consider installing from source:

pip install git+https://github.com/tgp-team/torch-geometric-pool.git

[!CAUTION] tgp is built upon PyTorch>=1.8 and PyG>=2.6. Make sure you have both installed in your environment before installation. Check the installation guide for more details.

Quick Example

import torch
from torch_geometric.data import Data
from torch_geometric.nn import GCNConv
from tgp.poolers import get_pooler

# Create a simple graph (5 nodes, 5 edges)
edge_index = torch.tensor([[0, 1, 2, 3, 4, 0],
                           [1, 0, 3, 2, 0, 4]], dtype=torch.long)
x = torch.randn((5, 16))  # node features

data = Data(x=x, edge_index=edge_index)

# Instantiate a Top-K pooling layer via its alias
pool = get_pooler("topk", in_channels=32, ratio=0.5)

# Forward pass with pooling
out = GCNConv(in_channels=16, out_channels=32)(data.x, data.edge_index)
out = pool(out, data.edge_index, batch=None)  # PoolingOutput(so=[5, 3], x=[3, 32], edge_index=[2, 2])
out = GCNConv(in_channels=32, out_channels=32)(out.x, out.edge_index)

For more detailed examples, refer to the Tutorials and the coding Examples.

Contributing

Contributions are welcome! For major changes or new features, please open an issue first to discuss your ideas. See the Contributing guidelines for more details on how to get involved. Help us build a better tgp!

Thanks to all contributors 🤝

Citing

If you use Torch Geometric Pool for your research, please consider citing the paper:

@misc{bianchi2025torchgeometricpoolpytorch,
      title={Torch Geometric Pool: the Pytorch library for pooling in Graph Neural Networks}, 
      author={Filippo Maria Bianchi and Carlo Abate and Ivan Marisca},
      year={2025},
      eprint={2512.12642},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2512.12642}, 
}

By Filippo Maria Bianchi, Ivan Marisca and Carlo Abate.

License

This project is licensed under the terms of the MIT license. See the LICENSE file for details.