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 MinCutPool. Each operator adheres to the modular SRC framework, allowing to quickly echange them within the same GNN architecture and combine 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", "diffpool", "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 tgp can be easily combined with each other, replaced with existing modules or with completely 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/regression, node classification/regression, 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 library

@software{Bianchi_Torch_Geometric_Pool_2025,
    author = {Bianchi, Filippo Maria and Marisca, Ivan},
    license = {MIT},
    month = {3},
    title = {{Torch Geometric Pool}},
    url = {https://github.com/tgp-team/tgp},
    year = {2025}
}

By Filippo Maria Bianchi and Ivan Marisca.

License

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