Quick start

In the following, we will go through a few examples that showcase the main functionalities of tgp.

If you’re running this in Google Colab, you might need to restart the session after running the installation cell below.

import sys
if 'google.colab' in sys.modules:
    %pip install torch==2.4.1 --index-url https://download.pytorch.org/whl/cu124
    %pip install torch_geometric==2.6.1
    %pip install torch_scatter torch_sparse torch_cluster -f https://data.pyg.org/whl/torch-2.4.0+cu124.html
    %pip install pygsp==0.6.1
    %pip install -q git+https://github.com/tgp-team/torch-geometric-pool.git@main

Let’s start by importing the required libraries and checking the pooling operators that are available.

import torch
import torch.nn.functional as F
from torch_geometric.datasets import TUDataset
from torch_geometric.loader import DataLoader
from torch_geometric.nn import DenseGCNConv, GCNConv

from tgp.poolers import TopkPooling, get_pooler, pooler_classes, pooler_map
from tgp.reduce import GlobalReduce

torch.set_printoptions(threshold=2, edgeitems=2)

print("Available poolers:")
for i,pooler in enumerate(pooler_classes):
    print(f"{i+1}. {pooler}")

Available poolers:
1. ASAPooling
2. AsymCheegerCutPooling
3. BNPool
4. DiffPool
5. DMoNPooling
6. EdgeContractionPooling
7. EigenPooling
8. GraclusPooling
9. HOSCPooling
10. LaPooling
11. JustBalancePooling
12. KMISPooling
13. MaxCutPooling
14. MinCutPooling
15. NDPPooling
16. NMFPooling
17. NoPool
18. PANPooling
19. SAGPooling
20. SEPPooling
21. TopkPooling

For example, let’s create a TopkPooling object.

pooler = TopkPooling(in_channels=16)
print(f"Pooler: {pooler}")

Pooler: TopkPooling(
	select=TopkSelect(in_channels=16, ratio=0.5, act=Tanh(), s_inv_op=transpose)
	reduce=BaseReduce()
	lift=BaseLift(matrix_op=precomputed, reduce_op=sum)
	connect=SparseConnect(reduce_op=sum, remove_self_loops=True, edge_weight_norm=False, degree_norm=False)
	multiplier=1.0
)

Each pooler is associated with an alias that can be used to quickly instantiate a pooler.

print("Available poolers:")
for alias, cls in zip(pooler_map.keys(), pooler_map.values()):
    print(f"'{alias}' --> {cls.__name__}")

Available poolers:
'asap' --> ASAPooling
'acc' --> AsymCheegerCutPooling
'bnpool' --> BNPool
'diff' --> DiffPool
'dmon' --> DMoNPooling
'ec' --> EdgeContractionPooling
'eigen' --> EigenPooling
'graclus' --> GraclusPooling
'hosc' --> HOSCPooling
'lap' --> LaPooling
'jb' --> JustBalancePooling
'kmis' --> KMISPooling
'maxcut' --> MaxCutPooling
'mincut' --> MinCutPooling
'ndp' --> NDPPooling
'nmf' --> NMFPooling
'nopool' --> NoPool
'pan' --> PANPooling
'sag' --> SAGPooling
'sep' --> SEPPooling
'topk' --> TopkPooling

We can instantiate the same object of class TopkPooling by passing the alias and a dict with the parameters needed to initialize the pooler.

params = {
    "in_channels": 3,  # Number of input features
    "ratio": 0.25,  # Percentage of nodes to keep
}

pooler = get_pooler("topk", **params)  # Get the pooler by alias
print(pooler)

TopkPooling(
	select=TopkSelect(in_channels=3, ratio=0.25, act=Tanh(), s_inv_op=transpose)
	reduce=BaseReduce()
	lift=BaseLift(matrix_op=precomputed, reduce_op=sum)
	connect=SparseConnect(reduce_op=sum, remove_self_loops=True, edge_weight_norm=False, degree_norm=False)
	multiplier=1.0
)

We see that each pooling layer implements a specific select (\(\texttt{SEL}\)), reduce (\(\texttt{RED}\)), connect \(\texttt{CON}\) operations, as defined by the SRC framework.

• The \(\texttt{SEL}\) operation is what sets most pooling methods apart and defines how the nodes are assigned to the supernodes of the pooled graph.

• The \(\texttt{RED}\) operation specifies how to compute the features of the supernodes in the pooled graph.

• Finally, \(\texttt{CON}\) creates the connectivity matrix of the pooled graph.

The pooling operators also have a \(\texttt{LIFT}\) function, which is used by some GNN architectures to map the pooled node features back to the node space of the original graph. See here for an introduction to the SRC(L) framework.

Calling a pooling layer

A pooling layer can be called similarly to a message-passing layer in PyG. Let’s start by loading some data and creating a data batch.

dataset = TUDataset(root="/tmp/ENZYMES", name="ENZYMES")
print(f"Dataset: {dataset}")
loader = DataLoader(dataset, batch_size=32, shuffle=True)
data_batch = next(iter(loader))
print(f"Data batch: {data_batch}")

Downloading https://www.chrsmrrs.com/graphkerneldatasets/ENZYMES.zip

Processing...

Dataset: ENZYMES(600)
Data batch: DataBatch(edge_index=[2, 4198], x=[1070, 3], y=[32], batch=[1070], ptr=[33])

Done!

Pooling operators support edge weights, i.e., scalar values stored in a edge_weight attribute. However, some dataset have edge features stored in the edge_attr field. In tgp we assume that the edge attributes are processed by the message-passing layers before pooling, which embed the attributes into the node features that reach the pooling operators.

pooling_output = pooler(
    x=data_batch.x,
    adj=data_batch.edge_index,
    edge_weight=data_batch.edge_weight,
    batch=data_batch.batch,
)
print(pooling_output)

PoolingOutput(so=[1070, 279], x=[279, 3], edge_index=[2, 450], edge_weight=None, batch=[279], mask=None, loss=None)

The output of a pooling layer is an object of class PoolingOutput that contains different fields:

• the node features of the pooled graph (x),

• the indices and weights of the pooled adjacency matrix (edge_index, edge_weight),

• the batch indices of the pooled graphs (batch).

In addition, so is an object of class SelectOutput, i.e., the output of the \(\texttt{SEL}\) operation that describes how the nodes of the original graph are assigned to the supernodes of the pooled graph.

print(pooling_output.so)

SelectOutput(num_nodes=1070, num_supernodes=279)

Some pooling operators save additional data structures in the PoolingOutput, to be used downstream the \(\texttt{RES}\) and \(\texttt{CON}\).

The pooling layer can also be used to perform \(\texttt{LIFT}\), i.e., to map the pooled features back to the original node space.

x_lift = pooler(
    x=pooling_output.x, so=pooling_output.so, batch=pooling_output.batch, lifting=True
)

print(f"original x shape: {data_batch.x.shape}")
print(f"pooled x shape: {pooling_output.x.shape}")
print(f"x_lift shape: {x_lift.shape}")

original x shape: torch.Size([1070, 3])
pooled x shape: torch.Size([279, 3])
x_lift shape: torch.Size([1070, 3])

\(\texttt{LIFT}\) is typically used by GNNs with an autoencoder architecture that perform node-level tasks (e.g., node classification).

Types of pooling operator

On of the main differnces between the pooling operators in tgp is if they are dense or sparse. TopkPooling that we just saw is a sparse method. Let’s now look at a dense pooler: MinCutPooling.

params = {
    "in_channels": 3,  # Number of input features
    "k": 10,  # Number of supernodes in the pooled graph
}

dense_pooler = get_pooler("mincut", **params)
print(dense_pooler)

MinCutPooling(
	select=MLPSelect(in_channels=[3], k=10, act=None, dropout=0.0, s_inv_op=transpose)
	reduce=BaseReduce()
	lift=BaseLift(matrix_op=precomputed, reduce_op=sum)
	connect=DenseConnect(remove_self_loops=True, degree_norm=True, adj_transpose=True, edge_weight_norm=False, sparse_output=False)
	batched=True
	cut_loss_coeff=1.0
	ortho_loss_coeff=1.0
)

Something that sets sparse and dense pooling methods apart is the format of the data they use internally. In tgp, all poolers are called with the same interface: pass node features, edge index, and batch vector. Poolers accept ((x, adj, batch)) directly; dense poolers handle conversion to dense format internally.

dense_pooling_output = dense_pooler(
    x=data_batch.x,
    adj=data_batch.edge_index,
    batch=data_batch.batch,
)
print(dense_pooling_output)

PoolingOutput(so=[52, 10], x=[32, 10, 3], edge_index=[32, 10, 10], edge_weight=None, batch=None, mask=[32, 10], loss=['cut_loss', 'ortho_loss'])

The connectivity of the coarsened graphs generated by a dense pooling operator is also a dense tensor.

print(dense_pooling_output.edge_index[0])

tensor([[0.0000, 0.1528,  ..., 0.1284, 0.1276],
        [0.1528, 0.0000,  ..., 0.1232, 0.1277],
        ...,
        [0.1284, 0.1232,  ..., 0.0000, 0.1040],
        [0.1276, 0.1277,  ..., 0.1040, 0.0000]], grad_fn=<SelectBackward0>)

Another difference w.r.t. TopkPooling is the presence of one or more loss terms in the pooling output.

for key, value in dense_pooling_output.loss.items():
    print(f"{key}: {value:.3f}")

cut_loss: -0.971
ortho_loss: 1.158

These are auxiliary losses that must be minimized along with the other task’s losses used to train the GNN. Most dense pooling methods have an auxiliary loss. A few sparse methods have an auxiliary loss too.

GNN model with pooling layers

Let’s create a simple GNN for graph classification with the following architecture:

\[[\texttt{MP}-\texttt{Pool}-\texttt{MP}-\texttt{GlobalPool}-\texttt{Linear}]\]

Initialization

First, in the __init__() we specify the architecture, instatiating the MP layers, the pooling layer from its alias and parameters, and the readout.

class GNN(torch.nn.Module):
    def __init__(
        self, in_channels, out_channels, pooler_type, pooler_kwargs, hidden_channels=64
    ):
        super().__init__()

        # First MP layer
        self.conv1 = GCNConv(in_channels=in_channels, out_channels=hidden_channels)

        # Pooling
        self.pooler = pooler_kwargs.update({"in_channels": hidden_channels})
        self.pooler = get_pooler(pooler_type, **pooler_kwargs)

        # Second MP layer
        if self.pooler.is_dense:
            self.conv2 = DenseGCNConv(
                in_channels=hidden_channels, out_channels=hidden_channels
            )
        else:
            self.conv2 = GCNConv(
                in_channels=hidden_channels, out_channels=hidden_channels
            )

        # Global readout
        self.readout = GlobalReduce(reduce_op="sum")

        # Classification layer
        self.lin = torch.nn.Linear(hidden_channels, out_channels)

Note that the type of pooling operator determines what kind of MP layer is used after pooling. A sparse pooler is followed by a sparse MP operator such as GCNConv. On the other hand, a dense pooling operator that returns a dense connectivity matrix must be followed by a dense MP layer such as DenseGCNConv. The type of pooling operator can be checked by the property is_dense.

Forward pass

Next, we define the forward pass of the GNN.

def forward(self, x, edge_index, edge_weight, batch=None):
    # First MP layer
    x = self.conv1(x, edge_index, edge_weight)
    x = F.relu(x)

    # Pooling
    out = self.pooler(x=x, adj=edge_index, edge_weight=edge_weight, batch=batch)
    x_pool, adj_pool = out.x, out.edge_index

    # Second MP layer
    x = self.conv2(x_pool, adj_pool)
    x = F.relu(x)

    # Global pooling
    x = self.readout(x, batch=out.batch)

    # Readout layer
    x = self.lin(x)

    if out.loss is not None:
        return F.log_softmax(x, dim=-1), sum(out.get_loss_value())
    else:
        return F.log_softmax(x, dim=-1), torch.tensor(0.0)


GNN.forward = forward

There are a few things to discuss.

Global pooling

The global pooling operation combines all the features in the current graph and is implemented differently depending if the pooler is sparse or dense. In the sparse case, we have a batch tensor indicating to which graph each node belongs to. In this case, global pooling should combine the features of the nodes belonging to the same graph. The output is a tensor of shape \([B, F]\).

# Sparse case
print(f"Input shape: {data_batch.x.shape}")
reducer_sparse = GlobalReduce(reduce_op="sum")
out_global_sparse = reducer_sparse(
    data_batch.x, batch=data_batch.batch
)
print(f"Output shape: {out_global_sparse.shape}")

Input shape: torch.Size([1070, 3])
Output shape: torch.Size([32, 3])

In the dense case, the features of the pooled graph are stored in a tensor of shape \([B, K, F]\) and global pooling can be done e.g., by summing or taking the average across the nodes dimension, yielding a tensor of shape \([B, F]\). In this case, batch is not needed.

# Dense case
print(f"Input shape: {dense_pooling_output.x.shape}")
reducer_dense = GlobalReduce(reduce_op="sum")
out_global_dense = reducer_dense(dense_pooling_output.x, batch=None)
print(f"Output shape: {out_global_dense.shape}")

Input shape: torch.Size([32, 10, 3])
Output shape: torch.Size([32, 3])

Note that in both cases the output is the same.

torch.allclose(out_global_sparse, out_global_dense)

True

Auxiliary losses

As we saw earlier, some pooling operators return an auxiliary loss, while others do not. In the forward pass we check if out.loss is not None and, in case, return the sum of all the auxiliary losses to be passed to the optimizer.

Testing the model

Let’s first test our GNN when configured with a sparse pooler.

num_features = dataset.num_features
num_classes = dataset.num_classes

sparse_params = {
    "ratio": 0.25,  # Percentage of nodes to keep
}

sparse_pool_gnn = GNN(
    in_channels=num_features,
    out_channels=num_classes,
    pooler_type="topk",
    pooler_kwargs=sparse_params,
)

sparse_gnn_out = sparse_pool_gnn(
    x=data_batch.x,
    edge_index=data_batch.edge_index,
    edge_weight=data_batch.edge_weight,
    batch=data_batch.batch,
)
print(f"Sparse GNN output shape: {sparse_gnn_out[0].shape}")
print(f"Sparse GNN loss: {sparse_gnn_out[1]:.3f}")

Sparse GNN output shape: torch.Size([32, 6])
Sparse GNN loss: 0.000

Since there is no auxiliary loss, the second output of the GNN is simply a constant zero-valued tensor that will not affect the gradients computation.

Next, we create the GNN with the dense pooling layer.

dense_params = {
    "k": 10,  # Number of supernodes in the pooled graph
}
dense_pool_gnn = GNN(
    in_channels=num_features,
    out_channels=num_classes,
    pooler_type="mincut",
    pooler_kwargs=dense_params,
)
dense_gnn_out = dense_pool_gnn(
    x=data_batch.x,
    edge_index=data_batch.edge_index,
    edge_weight=data_batch.edge_weight,
    batch=data_batch.batch,
)
print(f"Dense GNN output shape: {dense_gnn_out[0].shape}")
print(f"Dense GNN loss: {dense_gnn_out[1]:.3f}")

Dense GNN output shape: torch.Size([32, 6])
Dense GNN loss: 0.169

This time, we get an auxiliary loss that should be added to the other losses, e.g. the classification loss of the downstream task.

total_loss = F.nll_loss(dense_gnn_out[0], data_batch.y.view(-1)) + dense_gnn_out[1]
print(f"Loss: {total_loss:.3f}")

Loss: 2.313

And that’s it! We can train this GNN as any other that we normally build with PyG.

You can check the complete graph classification example here.