from collections.abc import Iterator
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Union
import torch
from torch import Tensor
from torch_geometric.data import Data
from torch_geometric.typing import Adj
from torch_geometric.utils import to_dense_adj, to_dense_batch
from tgp.connect import Connect
from tgp.lift import Lift
from tgp.reduce import Reduce
from tgp.select import Select, SelectOutput
from tgp.utils import Signature, connectivity_to_edge_index, foo_signature
from tgp.utils.ops import build_pooled_batch, dense_to_block_diag, is_dense_adj
[docs]
@dataclass
class PoolingOutput:
r"""The pooling output of a model of class :class:`~tgp.src.SRCPooling`.
Args:
x (~torch.Tensor): The pooled node features.
edge_index (~torch.Tensor): The edge indices of the pooled graph.
edge_weight (~torch.Tensor, optional): The edge features of the coarsened
graph. (default: :obj:`None`)
batch (~torch.Tensor, optional): The batch vector of the pooled nodes.
so (:class:`~tgp.select.SelectOutput`): The selection output. (default: :obj:`None`)
mask: Derived from ``so.out_mask`` when ``so`` is set.
This is a pooled-supernode validity mask of shape :math:`[B, K]`
(or :math:`[1, K]` for single-graph dense assignments).
loss (Optional[Dict], optional): The loss dictionary. (default: :obj:`None`)
"""
x: Optional[Tensor] = None
edge_index: Optional[Tensor] = None
edge_weight: Optional[Tensor] = None
batch: Optional[Tensor] = None
so: Optional[SelectOutput] = None
loss: Optional[Dict] = None
@property
def mask(self) -> Optional[Tensor]:
"""Pooled-supernode validity mask, derived from ``so.out_mask``."""
return self.so.out_mask if self.so is not None else None
def __repr__(self) -> str:
return (
f"PoolingOutput(so={[self.so.num_nodes, self.so.num_supernodes] if self.so is not None else None}, "
f"x={[*self.x.shape] if self.x is not None else None}, "
f"edge_index={[*self.edge_index.shape] if self.edge_index is not None else None}, "
f"edge_weight={[*self.edge_weight.shape] if self.edge_weight is not None else None}, "
f"batch={[*self.batch.shape] if self.batch is not None else None}, "
f"mask={[*self.mask.shape] if self.mask is not None else None}, "
f"loss={list(self.loss.keys()) if self.loss is not None else None})"
)
def __iter__(self) -> Iterator:
return iter(
(
self.x,
self.edge_index,
self.edge_weight,
self.batch,
self.mask,
self.so,
self.loss,
)
)
@property
def has_loss(self):
r"""Returns :obj:`True` if the pooling output has a loss."""
return bool(isinstance(self.loss, dict) and len(self.loss) > 0)
[docs]
def get_loss_value(self, name: str = None) -> Union[float, List[float]]:
r"""Returns the value of the loss with name ``name`` or all losses.
If the pooling output does not have a loss, it returns ``0``.
Args:
name (str, optional): The name of the loss to return. If :obj:`None`, returns all losses.
(default: :obj:`None`)
Returns:
Union[float, List[float]]: The value of the loss ``name`` or all losses.
"""
if not self.has_loss:
return 0
if name is None:
return [v for v in self.loss.values()]
return self.loss[name]
[docs]
def as_data(self):
r"""Converts the pooling output to a :class:`~torch_geometric.data.Data` object.
Returns:
~torch_geometric.data.Data: The pooling output as a Data object.
"""
num_nodes = None
if self.batch is not None:
num_nodes = self.batch.numel()
elif self.x is not None:
num_nodes = self.x.size(-2)
elif self.so is not None:
num_nodes = self.so.num_supernodes
return Data(
x=self.x,
edge_index=self.edge_index,
edge_weight=self.edge_weight,
batch=self.batch,
mask=self.mask,
so=self.so,
num_nodes=num_nodes,
)
[docs]
class SRCPooling(torch.nn.Module):
r"""A base class for pooling layers based on the paper
`"Understanding Pooling in Graph Neural Networks" <https://arxiv.org/abs/1905.05178>`_
(Grattarola et al., TNNLS 2022). Each pooler should inherit from this class.
:class:`~tgp.src.SRCPooling` decomposes a pooling layer into three components:
+ :class:`~tgp.select.Select` defines how input nodes map to supernodes.
+ :class:`~tgp.reduce.Reduce` defines how input node features are aggregated.
+ :class:`~tgp.lift.Lift` defines how pooled node features are un-pooled.
+ :class:`~tgp.connect.Connect` decides how the supernodes are connected to each other.
This class should return an object of type :obj:`~tgp.src.PoolingOutput`.
Args:
selector (:class:`~tgp.select.Select`): The node selection operator.
reducer (:class:`~tgp.reduce.Reduce`): The node feature aggregation operator.
lifter (:class:`~tgp.lift.Lift`): The node feature un-pooling operator.
connector (:class:`~tgp.connect.Connect`): The edge connection operator.
cached (bool, optional): If set to :obj:`True`, will cache the
:class:`~tgp.select.Select` output and the :class:`~tgp.connect.Connect`
output. (default: :obj:`False`)
"""
def __init__(
self,
selector: Select = None,
reducer: Reduce = None,
lifter: Lift = None,
connector: Connect = None,
cached: bool = False,
):
super().__init__()
self.selector = selector
self.reducer = reducer
self.lifter = lifter
self.connector = connector
self.cached = cached
self._so_cached = None
self._pooled_edge_index = None
self._pooled_edge_weight = None
def reset_parameters(self):
r"""Resets all learnable parameters of the module."""
self.selector.reset_parameters()
self.reducer.reset_parameters()
self.lifter.reset_parameters()
self.connector.reset_parameters()
[docs]
def select(
self,
**kwargs,
) -> SelectOutput:
r"""Calls the :class:`~tgp.select.Select` operator.
Returns:
An object of type :class:`~tgp.select.SelectOutput` containing the
mapping from nodes to supernodes :math:`\mathbf{S} \in \mathbb{R}^{N \times K}`.
"""
if self.selector is not None:
if self._so_cached is not None:
return self._so_cached
so = self.selector(
**kwargs,
)
if self.cached:
self._so_cached = so
return so
raise NotImplementedError
[docs]
def reduce(
self,
**kwargs,
) -> Tensor:
r"""Calls the :class:`~tgp.reduce.Reduce` operator.
Returns:
The pooled supernode features :math:`\mathbf{X}_{\text{pool}}`.
"""
if self.reducer is not None:
return self.reducer(**kwargs)
raise NotImplementedError
[docs]
def lift(self, **kwargs) -> Adj:
r"""Calls the :class:`~tgp.lift.Lift` operator.
Returns:
The un-pooled node features :math:`\mathbf{X}_{\text{lift}} \approx \mathbf{X}`.
"""
if self.lifter is not None:
return self.lifter(**kwargs)
raise NotImplementedError
[docs]
def connect(
self,
**kwargs,
) -> Tuple[Adj, Optional[Tensor]]:
r"""Calls the :class:`~tgp.connect.Connect` operator.
Returns:
The adjacency matrix of the coarse graph :math:`\mathbf{A}_{\text{pool}}`.
"""
if self.connector is not None:
if self._pooled_edge_index is not None:
return self._pooled_edge_index, self._pooled_edge_weight
pooled_edge_index, pooled_edge_weight = self.connector(**kwargs)
if self.cached:
self._pooled_edge_index = pooled_edge_index
self._pooled_edge_weight = pooled_edge_weight
return pooled_edge_index, pooled_edge_weight
raise NotImplementedError
[docs]
def preprocessing(
self, x: Tensor, edge_index: Adj, **kwargs
) -> Tuple[Tensor, Adj, Optional[Tensor]]:
"""Preprocess inputs, if needed."""
return x, edge_index, None
@property
def is_dense(self) -> bool:
"""Returns :obj:`True` if the pooler uses dense assignments."""
if self.selector is not None:
return self.selector.is_dense
raise NotImplementedError
@property
def is_sparse(self) -> bool:
"""Returns :obj:`True` if the pooler uses sparse assignments."""
return not self.is_dense
@property
def has_loss(self) -> bool:
r"""Returns :obj:`True` if the pooler overrides ``compute_loss``."""
return self.compute_loss.__qualname__.split(".")[0] != "SRCPooling"
@property
def is_trainable(self) -> bool:
r"""Returns :obj:`True` if any parameter belonging to the pooler or any
of its registered sub-modules is trainable.
"""
return any(p.requires_grad for p in self.parameters())
def compute_loss(self, *args, **kwargs) -> Optional[dict]:
"""Compute loss function."""
return None
[docs]
def clear_cache(self):
r"""Clear the caching done by :math:`\texttt{select}` and :math:`\texttt{connect}`."""
self._so_cached = None
self._pooled_edge_index = None
self._pooled_edge_weight = None
@property
def is_precoarsenable(self) -> bool:
r"""Returns :obj:`True` if the pooler is precoarsenable."""
if isinstance(self, Precoarsenable):
return not self.is_trainable
else:
return False
[docs]
@classmethod
def get_signature(cls) -> Signature:
"""Get signature of the pooler's ``__init__`` function."""
return foo_signature(cls)
[docs]
@classmethod
def get_forward_signature(cls) -> Signature:
"""Get signature of the pooler's ``forward`` function."""
return foo_signature(cls.forward)
def __repr__(self) -> str:
out = [f"{self.__class__.__name__}("]
out.append(f"\tselect={self.selector}")
out.append(f"\treduce={self.reducer}")
out.append(f"\tlift={self.lifter}")
out.append(f"\tconnect={self.connector}")
for k, v in self.extra_repr_args().items():
out.append(f"\t{k}={v}")
out.append(")")
return "\n".join(out)
def extra_repr_args(self) -> dict:
"""Add extra arguments to :meth:`~tgp.src.SRCPooling.__repr__`."""
return {}
[docs]
class DenseSRCPooling(SRCPooling):
r"""A base class for *dense* pooling layers that extends :class:`~tgp.src.SRCPooling`.
It provides a preprocessing function that transform a batch of graphs in
sparse representation into a batch of dense graphs.
When ``batched=True``, dense poolers accept either raw sparse inputs
(which are converted internally) or already-dense padded tensors. In the
latter case, an external input-node validity mask can be provided to mark
real (non-padded) nodes; otherwise a full-ones mask is assumed.
Args:
selector (:class:`~tgp.select.Select`): The *dense* :math:`\texttt{select}` operator.
reducer (:class:`~tgp.reduce.Reduce`): The *dense* :math:`\texttt{reduce}` operator.
lifter (:class:`~tgp.lift.Lift`): The *dense* :math:`\texttt{lift}` operator.
connector (:class:`~tgp.connect.Connect`): The *dense* :math:`\texttt{connect}` operator.
cached (bool, optional): If set to :obj:`True`, will cache the
:class:`~tgp.select.Select` output and the :class:`~tgp.connect.Connect`
output. (default: :obj:`False`)
cache_preprocessing (bool, optional):
If set to :obj:`True`, will cache the dense adjacency produced by
:meth:`~tgp.src.DenseSRCPooling.preprocessing`. This should only be
enabled when the same graph is reused across iterations (e.g.,
transductive single-graph tasks). (default: :obj:`False`)
adj_transpose (bool, optional):
If :obj:`True`, the preprocessing step and
the :class:`~tgp.connect.DenseConnect` operation returns transposed
adjacency matrices, so that they could be passed "as is" to the dense
message-passing layers.
(default: :obj:`True`)
batched (bool, optional):
If :obj:`True`, sparse inputs are converted to batched dense tensors
internally. If :obj:`False`, the pooler expects unbatched dense
assignments without padding. (default: :obj:`True`)
sparse_output (bool, optional):
If :obj:`True`, the pooled outputs are returned as block-diagonal
sparse representations. If :obj:`False`, outputs are returned in
batched dense form. (default: :obj:`False`)
"""
def __init__(
self,
selector: Select = None,
reducer: Reduce = None,
lifter: Lift = None,
connector: Connect = None,
cached: bool = False,
adj_transpose: bool = False,
batched: bool = True,
sparse_output: bool = False,
cache_preprocessing: bool = False,
):
super().__init__(
selector=selector,
reducer=reducer,
lifter=lifter,
connector=connector,
cached=cached,
)
self.batched = batched
self.sparse_output = sparse_output
self.adj_transpose = adj_transpose
self.cache_preprocessing = cache_preprocessing
self.preprocessing_cache = None
[docs]
def preprocessing(
self,
x: Tensor,
edge_index: Adj,
edge_weight: Optional[Tensor] = None,
batch: Optional[Tensor] = None,
max_num_nodes: Optional[int] = None,
batch_size: Optional[int] = None,
use_cache: bool = False,
**kwargs,
) -> Tuple[Tensor, Tensor, Tensor]:
r"""Preprocess inputs for dense pooling methods.
Transform a batch of graphs in sparse representation into a batch of graphs
with dense representation.
Args:
x (~torch.Tensor):
The node features.
A tensor of shape :math:`[N, F]`, where :math:`N` is the total number
of nodes in the batch and :math:`F` is the number of node features.
edge_index (~torch.Tensor):
The edge indices.
A tensor of of shape :math:`[2, E]`,
where :math:`E` is the number of edges in the batch.
edge_weight (~torch.Tensor, optional):
A vector of shape :math:`[E]` or :math:`[E, 1]` containing the weights
of the edges.
(default: :obj:`None`)
batch (~torch.Tensor, optional): The batch vector
:math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which indicates
to which graph in the batch each node belongs.
(default: :obj:`None`)
max_num_nodes (int, optional):
The maximum number of nodes of a graph in the batch.
(default: :obj:`None`)
batch_size (int, optional):
The number of graphs in the batch.
(default: :obj:`None`)
use_cache (bool, optional):
If :obj:`True`, it stores the preprocessed adjacency matrix in
``preprocessing_cache``. This is intended for static, single-graph
inputs. (default: :obj:`False`)
Returns:
tuple[~torch.Tensor, ~torch.Tensor, ~torch.Tensor]:
A tuple ``(x, adj, mask)`` where ``x`` contains batched node features
of shape :math:`[B, N_\text{max}, F]`, ``adj`` contains batched dense
adjacencies of shape :math:`[B, N_\text{max}, N_\text{max}]`, and
``mask`` marks real versus padded nodes with shape
:math:`[B, N_\text{max}]`.
"""
if use_cache and self.preprocessing_cache is not None:
adj = self.preprocessing_cache
else:
edge_index, edge_weight = connectivity_to_edge_index(
edge_index, edge_weight
)
adj = to_dense_adj(
edge_index=edge_index,
edge_attr=edge_weight,
max_num_nodes=max_num_nodes,
batch=batch,
batch_size=batch_size,
)
if self.adj_transpose:
adj = adj.transpose(-1, -2)
if use_cache:
self.preprocessing_cache = adj
x, mask = to_dense_batch(
x=x, batch=batch, max_num_nodes=max_num_nodes, batch_size=batch_size
)
return x, adj, mask
def _ensure_batched_inputs(
self,
x: Tensor,
edge_index: Optional[Adj],
edge_weight: Optional[Tensor],
batch: Optional[Tensor],
mask: Optional[Tensor],
use_cache: Optional[bool] = None,
) -> Tuple[Tensor, Tensor, Optional[Tensor]]:
if edge_index is None:
raise ValueError("edge_index cannot be None when batched=True.")
if use_cache is None:
use_cache = self.cache_preprocessing
if use_cache and batch is not None and batch.numel() > 0:
# Disable cache for multiple graphs in the batch
batch_min = int(batch.min().item())
batch_max = int(batch.max().item())
use_cache = batch_min == batch_max
if is_dense_adj(edge_index):
x = x.unsqueeze(0) if x.dim() == 2 else x
if mask is None:
mask = x.new_ones(x.size(0), x.size(1), dtype=torch.bool)
adj = edge_index
if use_cache:
self.preprocessing_cache = adj
return x, adj, mask
x, adj, mask = self.preprocessing(
x=x,
edge_index=edge_index,
edge_weight=edge_weight,
batch=batch,
use_cache=use_cache,
)
return x, adj, mask
[docs]
def clear_cache(self):
r"""Clear the caching done by :math:`\texttt{select}` and :math:`\texttt{connect}`
and the preprocessing cache.
"""
super().clear_cache()
self.preprocessing_cache = None
def _finalize_sparse_output(
self,
x_pool: Tensor,
adj_pool: Tensor,
batch: Optional[Tensor],
batch_pooled: Optional[Tensor],
so: SelectOutput,
) -> Tuple[Tensor, Tensor, Tensor, Optional[Tensor]]:
"""Convert batched dense outputs to block-diagonal sparse representation.
Uses ``so.out_mask`` when available so only valid supernodes and
edges between them are included. When ``so.out_mask`` is :obj:`None`,
all pooled nodes are kept.
"""
B, K = adj_pool.size(0), adj_pool.size(1)
x_flat = x_pool.reshape(-1, x_pool.size(-1))
out_mask = so.out_mask
if batch_pooled is None and batch is not None:
batch_pooled = self.reducer.reduce_batch(so, batch)
if batch_pooled is None and B > 1:
batch_pooled = build_pooled_batch(B, K, x_pool.device)
if batch_pooled is None and out_mask is not None:
# Single graph, no batch, dense path
batch_pooled = torch.zeros(B * K, dtype=torch.long, device=x_pool.device)
if out_mask is not None:
valid_flat = out_mask.reshape(-1)
valid_indices = valid_flat.nonzero(as_tuple=True)[0]
num_valid = valid_indices.size(0)
x_pool = x_flat[valid_indices]
batch_pooled = batch_pooled[valid_flat]
# Zero out edges at padded positions so dense_to_block_diag only sees valid edges
adj_masked = (
adj_pool
* out_mask.unsqueeze(-1).to(adj_pool.dtype)
* out_mask.unsqueeze(-2).to(adj_pool.dtype)
)
edge_index, edge_weight = dense_to_block_diag(adj_masked)
# Remap node indices from [0 .. B*K-1] to compact [0 .. num_valid-1]
old_to_new = torch.full(
(B * K,), -1, dtype=torch.long, device=x_pool.device
)
old_to_new[valid_indices] = torch.arange(num_valid, device=x_pool.device)
keep = (old_to_new[edge_index[0]] >= 0) & (old_to_new[edge_index[1]] >= 0)
edge_index = torch.stack(
[
old_to_new[edge_index[0][keep]],
old_to_new[edge_index[1][keep]],
],
dim=0,
)
edge_weight = edge_weight[keep]
else:
edge_index, edge_weight = dense_to_block_diag(adj_pool)
x_pool = x_flat
return x_pool, edge_index, edge_weight, batch_pooled
class Precoarsenable:
def precoarsening(
self,
**kwargs,
) -> PoolingOutput:
"""Precompute a coarsened graph from the original graph.
Must be implemented by the poolers that support precoarsening.
"""
raise NotImplementedError("Precoarsening is not supported by this pooler.")
def multi_level_precoarsening(
self,
levels: int,
edge_index: Optional[Adj] = None,
edge_weight: Optional[Tensor] = None,
*,
batch: Optional[Tensor] = None,
num_nodes: Optional[int] = None,
**kwargs,
) -> List[PoolingOutput]:
r"""Precompute multiple coarsening levels.
The default implementation performs a greedy rollout by repeatedly
calling :meth:`precoarsening`. Poolers can override this method to
implement method-specific multi-level strategies.
Notes:
The rollout clears pooling caches before each level when available.
This avoids stale :class:`~tgp.select.SelectOutput` reuse when
:obj:`cached=True` and the graph size changes between levels.
"""
if levels < 1:
raise ValueError(f"'levels' must be >= 1, got {levels}.")
clear_cache = getattr(self, "clear_cache", None)
pooled_levels = []
current_edge_index = edge_index
current_edge_weight = edge_weight
current_batch = batch
current_num_nodes = num_nodes
for _ in range(levels):
if callable(clear_cache):
clear_cache()
pooled = self.precoarsening(
edge_index=current_edge_index,
edge_weight=current_edge_weight,
batch=current_batch,
num_nodes=current_num_nodes,
**kwargs,
)
pooled_levels.append(pooled)
pooled_data = pooled.as_data()
current_edge_index = pooled_data.edge_index
current_edge_weight = pooled_data.edge_weight
current_batch = pooled_data.batch
current_num_nodes = pooled_data.num_nodes
if callable(clear_cache):
clear_cache()
return pooled_levels
class BasePrecoarseningMixin(Precoarsenable):
r"""A mixin class for pooling layers that implements the
pre-coarsening strategy.
"""
def _precoarsening_from_select_output(
self,
so: SelectOutput,
edge_index: Adj,
edge_weight: Optional[Tensor] = None,
*,
batch: Optional[Tensor] = None,
**kwargs,
) -> PoolingOutput:
# Reduce batch vector according to the select output.
if batch is None:
batch = so.batch if getattr(so, "batch", None) is not None else None
if batch is None:
n_nodes = so.num_nodes
batch = torch.zeros(n_nodes, dtype=torch.long, device=so.s.device)
so.batch = batch
batch_pooled = self.reducer.reduce_batch(select_output=so, batch=batch)
# Compute pooled adj matrix through the connect operator.
connector = getattr(self, "preconnector", self.connector)
edge_index_pooled, edge_weight_pooled = connector(
so=so,
edge_index=edge_index,
edge_weight=edge_weight,
batch=batch,
batch_pooled=batch_pooled,
**kwargs,
)
return PoolingOutput(
edge_index=edge_index_pooled,
edge_weight=edge_weight_pooled,
batch=batch_pooled,
so=so,
)
def precoarsening(
self,
edge_index: Optional[Adj] = None,
edge_weight: Optional[Tensor] = None,
*,
batch: Optional[Tensor] = None,
num_nodes: Optional[int] = None,
**kwargs,
) -> PoolingOutput:
if edge_index is None:
raise ValueError("edge_index cannot be None for precoarsening.")
so = self.select(
edge_index=edge_index,
edge_weight=edge_weight,
batch=batch,
num_nodes=num_nodes,
**kwargs,
)
return self._precoarsening_from_select_output(
so=so,
edge_index=edge_index,
edge_weight=edge_weight,
batch=batch,
**kwargs,
)
