Precoarsening and transforms¶
Some pooling operators such as NDPPooling, GraclusPooling, NMFPooling (and some configurations of KMISPooling and LaPooling), compute the \(\texttt{SEL}\) only based on the topology of the adacency matrix.
As opposed to the node features, which are modified by each layer of the GNN and evolve during training, the adjacency matrix is and remains fixed.
Therefore, the \(\texttt{SEL}\) and the \(\texttt{CON}\) operations of these poolers is always the same and can be precomputed bofeore starting to train the GNN.
This, allows us to save a lot of time during training because the only operation that we need to compute is the \(\texttt{RED}\) to compute the features of the supernodes.
Let’s start by loading some data.
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
from torch_geometric.datasets import TUDataset
dataset = TUDataset(root="/tmp/MUTAG", name="MUTAG", force_reload=True)
Processing...
Done!
Let’s now take the first graph.
data = dataset[0]
print(data)
Data(edge_index=[2, 38], x=[17, 7], edge_attr=[38, 4], y=[1])
Let’s consider NDPPooling: its \(\texttt{SEL}\) operation only looks at the graph connectivity.
This means that we can compute the SelectOutput without having to pass the node features.
from tgp.connect import KronConnect
from tgp.select import NDPSelect
selector = NDPSelect()
connector = KronConnect()
# Compute pooled graph
so = selector(data.edge_index)
print(so)
SelectOutput(num_nodes=17, num_supernodes=9, extra={'L'})
This also means that we can compute the coarsened graph connectivity witht the \(\texttt{CON}\) operation.
edge_index_pool = connector(data.edge_index, so)
print(edge_index_pool)
(tensor([[0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 6, 6, 6, 6,
7, 7, 8, 8],
[1, 2, 0, 2, 3, 4, 5, 0, 1, 3, 1, 2, 4, 1, 3, 5, 6, 1, 4, 6, 4, 5, 7, 8,
6, 8, 6, 7]]), tensor([0.5000, 0.5000, 0.5000, 0.3333, 0.3333, 0.3333, 0.3333, 0.5000, 0.3333,
0.3333, 0.3333, 0.3333, 0.5000, 0.3333, 0.5000, 0.3333, 0.5000, 0.3333,
0.3333, 0.5000, 0.5000, 0.5000, 0.3333, 0.3333, 0.3333, 0.3333, 0.3333,
0.3333]))
Note
NDPPooling uses the Kron reduction implemented by KronConnect to compute the \(\texttt{connect}\) operation. However, once the SelectOutput is computed other \(\texttt{CON}\) opertions, e.g., SparseConnect, can be used.
At this point, we can apply the \(\texttt{SEL}\) and the \(\texttt{CON}\) operation one more time on the pooled graph. This is useful if we want to use a GNN architecture that applies pooling multiple times.
so2 = selector(edge_index_pool[0], edge_index_pool[1])
print(so2)
edge_index_pool2 = connector(edge_index_pool[0], so2, edge_index_pool[1])
print(edge_index_pool2)
SelectOutput(num_nodes=9, num_supernodes=5, extra={'L'})
(tensor([[0, 0, 0, 0, 1, 1, 2, 2, 2, 2, 3, 3, 3, 4, 4, 4],
[1, 2, 3, 4, 0, 2, 0, 1, 3, 4, 0, 2, 4, 0, 2, 3]]), tensor([0.6786, 0.6347, 0.0328, 0.0328, 0.6786, 0.1429, 0.6347, 0.1429, 0.1475,
0.1475, 0.0328, 0.1475, 0.4098, 0.0328, 0.1475, 0.4098]))
We can repeat the procedure iteratively for all the pooling levels that we want to have in our GNN.
The Precoarsening transform¶
Precomputing pooling allows us to save a lot of time because we only need to do it once before starting to train our GNN.
However, for each sample in our dataset we end up having an instance of SelectOutput and a pooled connectivity for each pooling level.
Handling all of them during training, while keeping the correct association between data structures when we shuffle the data, is cumbersome.
tgp provides a couple of tools to handle precomputed pooled graphs efficiently.
The first is the PreCoarsening transform, which can be directly applied to the dataset like all the other PyG transforms.
from tgp.poolers import NDPPooling
from tgp.data import PreCoarsening
dataset = TUDataset(
root="/tmp/MUTAG",
name="MUTAG",
pre_transform=PreCoarsening(
poolers=[NDPPooling(), NDPPooling()]
),
force_reload=True,
)
data = dataset[0]
print(data)
Processing...
Done!
/home/docs/checkouts/readthedocs.org/user_builds/torch-geometric-pool/envs/latest/lib/python3.10/site-packages/torch_geometric/datasets/tu_dataset.py:132: UserWarning: Weights only load failed. Please file an issue to make `torch.load(weights_only=True)` compatible in your case. Please use `torch.serialization.add_safe_globals([torch_geometric.data.data.DataEdgeAttr])` to allowlist this global.
out = fs.torch_load(self.processed_paths[0])
Data(edge_index=[2, 38], x=[17, 7], edge_attr=[38, 4], y=[1], pooled_data=[2])
Once again we look at the first element of the dataset and this time we see that, compared to the standard Data structure, there is an additional field, pooled_data, which is a list of length recursive_depth.
The elements in the list are the hierarchy of pooled graphs computed with the selector and connector that we defined in the PreCoarsening transform.
Each pooled graph is a Data structure containing the SelectOutput and the pooled connectivity matrix.
Since we are using NDPPooling, which internally calls NDPSelect, we get an extra argument L representing the Laplacian matrix used by KronConnect.
for pooled_data in data.pooled_data:
print(pooled_data)
Data(edge_index=[2, 26], edge_weight=[26], batch=[8], so=SelectOutput(num_nodes=17, num_supernodes=8, extra={'L'}), num_nodes=8)
Data(edge_index=[2, 12], edge_weight=[12], batch=[4], so=SelectOutput(num_nodes=8, num_supernodes=4, extra={'L'}), num_nodes=4)
This new Data strcture is very convenient as it carries all the information that the GNN needs to perform pooling at each coarsening level. With it, we do not need to keep track manually of the association between data samples and their pooled graph.
The PoolDataLoader¶
The field pooled_data in these custom Data structures is not handled properly by the standard DataLoader of 
PyG.
While the node features, x, the edge indices, edge attributes, etc… are batched correctly, the pooled graphs are just concatenated in a list rather than being combined into a single batched graph for each pooling level.
from torch_geometric.loader import DataLoader
pyg_loader = DataLoader(dataset, batch_size=4, shuffle=True)
next_batch = next(iter(pyg_loader))
print(next_batch)
print(next_batch.pooled_data[0])
DataBatch(edge_index=[2, 168], x=[73, 7], edge_attr=[168, 4], y=[4], pooled_data=[4], batch=[73], ptr=[5])
[Data(edge_index=[2, 20], edge_weight=[20], batch=[7], so=SelectOutput(num_nodes=12, num_supernodes=7, extra={'L'}), num_nodes=7), Data(edge_index=[2, 6], edge_weight=[6], batch=[3], so=SelectOutput(num_nodes=7, num_supernodes=3, extra={'L'}), num_nodes=3)]
To obtain well-formed batches with precomputed pooled graphs tgp provides the PoolDataLoader.
Now, the field pooled_data in the batch is a list containing a single batched graph for each coarsening level (2 in our case).
from tgp.data import PoolDataLoader
tgp_loader = PoolDataLoader(dataset, batch_size=16, shuffle=True)
next_batch = next(iter(tgp_loader))
print(next_batch)
print(next_batch.pooled_data[0])
DataPooledBatch(edge_index=[2, 574], x=[262, 7], edge_attr=[574, 4], y=[16], pooled_data=[2], batch=[262], ptr=[17])
Data(edge_index=[2, 396], edge_weight=[396], batch=[129], so=SelectOutput(num_nodes=262, num_supernodes=129, extra={'L'}), num_nodes=129, ptr=[17])
A complete example of usage can be found here.
Advanced usage: multiple poolers and different configs per level¶
You can build a hierarchical network with more than one pooler and different poolers or configs at each level. Pass a list of level specs to PreCoarsening(poolers=...): each element is either a string alias (e.g. "ndp", "graclus") or a tuple (alias, kwargs) for that level.
Examples:
Same pooler, same config:
["ndp", "ndp"]Same pooler, different config:
[("nmf", {"k": 8}), ("nmf", {"k": 4})]Mixed poolers:
["ndp", ("eigen", {"k": 4, "num_modes": 3})]
In the model, use a ModuleList of reducers from pre_transform.poolers, one conv per level (accounting for num_modes for EigenPooling), and in forward loop over data.pooled_data and call reducer(x=x, so=pooled.so) then the next conv. Below is a full runnable example (mixed poolers: NDP then Eigen). For more schedules and a standalone script, see pre_coarsening.py.
# Full runnable example: mixed poolers (NDP -> Eigen), 2 levels
import torch
import torch.nn.functional as F
from torch_geometric.datasets import TUDataset
from torch_geometric.nn import ARMAConv
from tgp.data import PreCoarsening, PoolDataLoader
from tgp.reduce import GlobalReduce
level_specs = ["ndp", ("eigen", {"k": 4, "num_modes": 3})]
pre_transform_adv = PreCoarsening(poolers=level_specs)
level_poolers = pre_transform_adv.poolers
num_levels = len(level_poolers)
dataset_adv = TUDataset(
root="/tmp/MUTAG_ndp_eigen",
name="MUTAG",
pre_transform=pre_transform_adv,
force_reload=True,
)
train_loader_adv = PoolDataLoader(dataset_adv[:150], batch_size=32, shuffle=True)
test_loader_adv = PoolDataLoader(dataset_adv[150:], batch_size=32)
level_num_modes = [getattr(p, "num_modes", 1) for p in level_poolers]
class PrecoarsenedNet(torch.nn.Module):
def __init__(self, hidden_channels=64):
super().__init__()
self.conv1 = ARMAConv(
dataset_adv.num_features, hidden_channels, num_layers=2
)
self.reducers = torch.nn.ModuleList([p.reducer for p in level_poolers])
self.next_conv = torch.nn.ModuleList()
for num_modes in level_num_modes:
in_ch = hidden_channels * num_modes
self.next_conv.append(
ARMAConv(in_ch, hidden_channels, num_layers=2)
)
self.lin = torch.nn.Linear(hidden_channels, dataset_adv.num_classes)
def forward(self, data):
x = self.conv1(data.x, data.edge_index, data.edge_weight)
x = F.relu(x)
for pooled, conv, reducer in zip(data.pooled_data, self.next_conv, self.reducers):
x, _ = reducer(x=x, so=pooled.so)
x = conv(x, pooled.edge_index, pooled.edge_weight)
x = F.relu(x)
readout_layer = GlobalReduce(reduce_op="sum")
x = readout_layer(x, batch=pooled.batch)
return F.log_softmax(self.lin(x), dim=-1)
device_adv = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_adv = PrecoarsenedNet(hidden_channels=64).to(device_adv)
optimizer_adv = torch.optim.Adam(model_adv.parameters(), lr=1e-4)
def train_adv():
model_adv.train()
total = 0
for data in train_loader_adv:
data = data.to(device_adv)
optimizer_adv.zero_grad()
loss = F.nll_loss(model_adv(data), data.y.view(-1))
loss.backward()
optimizer_adv.step()
total += data.y.size(0) * float(loss)
return total / len(dataset_adv)
@torch.no_grad()
def test_adv(loader):
model_adv.eval()
correct = 0
for data in loader:
data = data.to(device_adv)
correct += int(model_adv(data).argmax(dim=-1).eq(data.y.view(-1)).sum())
return correct / len(loader.dataset)
for epoch in range(1, 11):
loss = train_adv()
acc = test_adv(test_loader_adv)
print(f"Epoch {epoch:2d}, Loss: {loss:.4f}, Test Acc: {acc:.4f}")
print(f"Final test accuracy: {test_adv(test_loader_adv):.4f}")
Downloading https://www.chrsmrrs.com/graphkerneldatasets/MUTAG.zip
Processing...
Done!
/home/docs/checkouts/readthedocs.org/user_builds/torch-geometric-pool/envs/latest/lib/python3.10/site-packages/torch_geometric/datasets/tu_dataset.py:132: UserWarning: Weights only load failed. Please file an issue to make `torch.load(weights_only=True)` compatible in your case. Please use `torch.serialization.add_safe_globals([torch_geometric.data.data.DataEdgeAttr])` to allowlist this global.
out = fs.torch_load(self.processed_paths[0])
Epoch 1, Loss: 0.5326, Test Acc: 0.6842
Epoch 2, Loss: 0.5013, Test Acc: 0.6842
Epoch 3, Loss: 0.4779, Test Acc: 0.6842
Epoch 4, Loss: 0.4614, Test Acc: 0.6842
Epoch 5, Loss: 0.4485, Test Acc: 0.6842
Epoch 6, Loss: 0.4395, Test Acc: 0.6842
Epoch 7, Loss: 0.4348, Test Acc: 0.6842
Epoch 8, Loss: 0.4295, Test Acc: 0.6842
Epoch 9, Loss: 0.4255, Test Acc: 0.6842
Epoch 10, Loss: 0.4203, Test Acc: 0.6842
Final test accuracy: 0.6842
Other data transforms¶
Some pooling layers come with custom transforms that should be applied to the data before starting to train the GNN.
For example, JustBalancePooling transforms the connectivity matrix \(\mathbf{A}\) as follows:
The transforms associated with a given pooling operator are stored in the field data_transforms().
They can be accessed and passed to the dataset as any other PyG transform.
from tgp.poolers import JustBalancePooling
pooler = JustBalancePooling(in_channels=dataset.num_features, k=10)
print(pooler.data_transforms())
NormalizeAdj()
dataset = TUDataset(
root="/tmp/MUTAG",
name="MUTAG",
force_reload=True,
pre_transform=pooler.data_transforms(), # transform specific for the pooler
)
Processing...
Done!
tgp