Torch Geometric Pool¶

The graph pooling library made for PyTorch Geometric.

tgp (Torch Geometric Pool) is a library built on top of PyTorch Geometric that brings every major graph-pooling operator into a single, unified framework. Drop-in layers let you turn any vanilla GNN into a hierarchical one with just a few lines of code.

Want to learn more about graph pooling?

Read our series of blog posts on the topic.

Features¶

Unified API

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

The SRC(L) framework

All your pooling operators in one place.

Choose from a variety of pooling methods, including sparse techniques like TopkPooling, NDPPooling, GraclusPooling, and dense methods such as DiffPool and MinCutPooling.

Poolers

Tweak and create new pooling layers.

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

SRCPooling

Precomputed & On-the-Fly Pooling.

Accelerate training by precomputing the coarse graph (assignments and connectivity) for methods like NDPPooling and GraclusPooling. Alternatively, use on-the-fly pooling (e.g., TopkPooling or MinCutPooling) that computes assignments dynamically and supports end-to-end gradient flow.

Data

Not familiar with Select-Reduce-Connect? Read our introductory notes about The SRC(L) framework.

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

Before installation

tgp requires PyTorch>=1.8 and PyG>=2.6 to be present in your environment.

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 a deeper dive, head over to the Installation section, browse the Tutorials, and explore our APIs.

Citing¶

If you use tgp 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},
}