from typing import List, Optional, Union
import torch
from torch import Tensor
from torch_geometric.typing import Adj
from torch_geometric.utils import scatter
from tgp.connect import DenseConnect
from tgp.lift import BaseLift
from tgp.reduce import BaseReduce
from tgp.select import MLPSelect, SelectOutput
from tgp.src import DenseSRCPooling, PoolingOutput
from tgp.utils.losses import (
hosc_orthogonality_loss,
mincut_loss,
orthogonality_loss,
sparse_ho_mincut_loss,
sparse_mincut_loss,
unbatched_hosc_orthogonality_loss,
unbatched_orthogonality_loss,
)
from tgp.utils.ops import connectivity_to_edge_index, postprocess_adj_pool_dense
from tgp.utils.typing import LiftType, SinvType
[docs]
class HOSCPooling(DenseSRCPooling):
r"""The high-order pooling operator from the paper
`"Higher-order clustering and pooling for Graph Neural Networks"
<http://arxiv.org/abs/2209.03473>`_ (Duval & Malliaros, CIKM 2022)..
+ The :math:`\texttt{select}` operator is implemented with :class:`~tgp.select.MLPSelect`.
+ The :math:`\texttt{reduce}` operator is implemented with :class:`~tgp.reduce.BaseReduce`.
+ The :math:`\texttt{connect}` operator is implemented with :class:`~tgp.connect.DenseConnect`.
+ The :math:`\texttt{lift}` operator is implemented with :class:`~tgp.lift.BaseLift`.
This layer optimizes a combination of the following auxiliary losses:
+ the mincut loss (:class:`~tgp.utils.losses.mincut_loss`),
+ the orthogonality loss (:class:`~tgp.utils.losses.orthogonality_loss`),
+ the hosc orthogonality loss (:class:`~tgp.utils.losses.hosc_orthogonality_loss`),
Args:
in_channels (int, list of int):
Number of hidden units for each hidden layer in the MLP
of the :math:`\texttt{select}` operator.
The first integer must match the size of the node features.
k (int):
Number of clusters or supernodes in the pooler graph.
act (str or Callable, optional):
Activation function in the hidden layers of the MLP
of the :math:`\texttt{select}` operator.
dropout (float, optional):
Dropout probability in the MLP of the :math:`\texttt{select}` operator.
(default: ``0.0``)
mu (float, optional):
A scalar that controls the importance given to regularization loss.
(default: ``0.1``)
alpha (float, optional):
A scalar in [0,1] controlling the importance granted
to higher-order information in the loss function.
(default: ``0.5``)
hosc_ortho (bool, optional):
Specifies either to use the hosc_orthogonality_loss or the
orthogonality_loss.
(default: :obj:`False`)
remove_self_loops (bool, optional):
If :obj:`True`, the self-loops will be removed from the adjacency matrix.
(default: :obj:`True`)
degree_norm (bool, optional):
If :obj:`True`, the adjacency matrix will be symmetrically normalized.
(default: :obj:`True`)
edge_weight_norm (bool, optional):
Whether to normalize the edge weights by dividing by the maximum absolute value per graph.
(default: :obj:`False`)
adj_transpose (bool, optional):
If :obj:`True`, the preprocessing step in :class:`~tgp.src.DenseSRCPooling` 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`)
cache_preprocessing (bool, optional):
If :obj:`True`, caches the dense adjacency produced during preprocessing.
This should only be enabled when the same graph is reused across iterations.
(default: :obj:`False`)
lift (~tgp.utils.typing.LiftType, optional):
Defines how to compute the matrix :math:`\mathbf{S}_\text{inv}` to lift the pooled node features.
- ``"precomputed"`` (default): Use as :math:`\mathbf{S}_\text{inv}` what is
already stored in the ``"s_inv"`` attribute of the :class:`~tgp.select.SelectOutput`.
- ``"transpose"``: Recomputes :math:`\mathbf{S}_\text{inv}` as :math:`\mathbf{S}^\top`,
the transpose of :math:`\mathbf{S}`.
- ``"inverse"``: Recomputes :math:`\mathbf{S}_\text{inv}` as :math:`\mathbf{S}^+`,
the Moore-Penrose pseudoinverse of :math:`\mathbf{S}`.
s_inv_op (~tgp.utils.typing.SinvType, optional):
The operation used to compute :math:`\mathbf{S}_\text{inv}` from the select matrix
:math:`\mathbf{S}`. :math:`\mathbf{S}_\text{inv}` is stored in the ``"s_inv"`` attribute of
the :class:`~tgp.select.SelectOutput`. It can be one of:
- ``"transpose"`` (default): Computes :math:`\mathbf{S}_\text{inv}` as :math:`\mathbf{S}^\top`,
the transpose of :math:`\mathbf{S}`.
- ``"inverse"``: Computes :math:`\mathbf{S}_\text{inv}` as :math:`\mathbf{S}^+`,
the Moore-Penrose pseudoinverse of :math:`\mathbf{S}`.
"""
def __init__(
self,
in_channels: Union[int, List[int]],
k: int,
act: str = None,
dropout: float = 0.0,
mu: float = 0.1,
alpha: float = 0.5,
hosc_ortho: bool = False,
remove_self_loops: bool = True,
degree_norm: bool = True,
edge_weight_norm: bool = False,
adj_transpose: bool = True,
lift: LiftType = "precomputed",
s_inv_op: SinvType = "transpose",
batched: bool = True,
sparse_output: bool = False,
cache_preprocessing: bool = False,
):
super().__init__(
selector=MLPSelect(
in_channels=in_channels,
k=k,
batched_representation=batched,
act=act,
dropout=dropout,
s_inv_op=s_inv_op,
),
reducer=BaseReduce(),
lifter=BaseLift(matrix_op=lift),
connector=DenseConnect(
remove_self_loops=remove_self_loops,
degree_norm=degree_norm,
adj_transpose=adj_transpose,
edge_weight_norm=edge_weight_norm,
sparse_output=sparse_output,
),
adj_transpose=adj_transpose,
cache_preprocessing=cache_preprocessing,
batched=batched,
sparse_output=sparse_output,
)
self.k = k
self.mu = mu
self.alpha = alpha
self.hosc_ortho = hosc_ortho
[docs]
def forward(
self,
x: Tensor,
adj: Optional[Adj] = None,
edge_weight: Optional[Tensor] = None,
so: Optional[SelectOutput] = None,
mask: Optional[Tensor] = None,
batch: Optional[Tensor] = None,
batch_pooled: Optional[Tensor] = None,
lifting: bool = False,
**kwargs,
) -> PoolingOutput:
r"""Forward pass.
Args:
x (~torch.Tensor): Node feature tensor
:math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`, with
batch-size :math:`B`, (maximum) number of nodes :math:`N` for
each graph, and feature dimension :math:`F`.
adj (~torch_geometric.typing.Adj, optional): The connectivity matrix.
In batched mode, this accepts sparse connectivity
(``edge_index``, ``torch_sparse.SparseTensor``, or torch COO),
which is internally converted to a dense padded tensor
:math:`\mathbf{A} \in \mathbb{R}^{B \times N \times N}`, or an
already dense adjacency tensor with the same shape.
(default: :obj:`None`)
edge_weight (~torch.Tensor, optional): Edge weights associated with
``adj`` when sparse connectivity is provided.
(default: :obj:`None`)
so (~tgp.select.SelectOutput, optional): The output of the :math:`\texttt{select}` operator.
(default: :obj:`None`)
mask (~torch.Tensor, optional): Input-node validity mask
:math:`\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}` with
:obj:`True` on real (non-padded) nodes in each graph. Only used
when inputs are already dense/padded. (default: :obj:`None`)
batch (~torch.Tensor, optional): Batch assignment vector for input
nodes. Required in sparse mode and optional in dense mode.
(default: :obj:`None`)
batch_pooled (~torch.Tensor, optional): Optional precomputed batch
assignment for pooled nodes, used when ``lifting=True``.
(default: :obj:`None`)
lifting (bool, optional): If set to :obj:`True`, the :math:`\texttt{lift}` operation is performed.
(default: :obj:`False`)
Returns:
~tgp.src.PoolingOutput: The output of the pooling operator.
"""
if lifting:
# Lift
x_lifted = self.lift(
x_pool=x, so=so, batch=batch, batch_pooled=batch_pooled
)
return x_lifted
# === Batched path ===
if self.batched:
x, adj, mask = self._ensure_batched_inputs(
x=x,
edge_index=adj,
edge_weight=edge_weight,
batch=batch,
mask=mask,
)
so = self.select(x=x, mask=mask)
x_pooled, batch_pooled = self.reduce(x=x, so=so, batch=batch)
adj_pool = self.connector.dense_connect(adj=adj, s=so.s)
adj_pool = postprocess_adj_pool_dense(
adj_pool,
remove_self_loops=self.connector.remove_self_loops,
degree_norm=self.connector.degree_norm,
adj_transpose=self.connector.adj_transpose,
edge_weight_norm=self.connector.edge_weight_norm,
)
loss = self.compute_loss(adj, so.s, adj_pool, mask)
if self.sparse_output:
x_pooled, edge_index_pooled, edge_weight_pooled, batch_pooled = (
self._finalize_sparse_output(
x_pool=x_pooled,
adj_pool=adj_pool,
batch=batch,
batch_pooled=batch_pooled,
so=so,
)
)
return PoolingOutput(
x=x_pooled,
edge_index=edge_index_pooled,
edge_weight=edge_weight_pooled,
batch=batch_pooled,
so=so,
loss=loss,
)
return PoolingOutput(x=x_pooled, edge_index=adj_pool, so=so, loss=loss)
# === Unbatched (sparse-loss) path ===
so = self.select(x=x, batch=batch)
return_batched = not self.sparse_output
x_pooled, batch_pooled = self.reduce(
x=x, so=so, batch=batch, return_batched=return_batched
)
edge_index_pooled, edge_weight_pooled = self.connect(
edge_index=adj,
so=so,
edge_weight=edge_weight,
batch=batch,
batch_pooled=batch_pooled,
)
loss = self.compute_sparse_loss(
adj,
edge_weight,
so.s,
batch,
adj_pool=edge_index_pooled,
edge_weight_pool=edge_weight_pooled,
batch_pooled=batch_pooled,
)
return PoolingOutput(
x=x_pooled,
edge_index=edge_index_pooled,
edge_weight=edge_weight_pooled,
batch=batch_pooled,
so=so,
loss=loss,
)
def compute_loss(
self, adj: Tensor, S: Tensor, adj_pool: Tensor, mask: Optional[Tensor] = None
) -> Optional[dict]:
r"""Computes the auxiliary loss terms.
Args:
adj (~torch.Tensor): The dense adjacency matrix.
S (~torch.Tensor): The dense assignment matrix.
adj_pool (~torch.Tensor): The pooled adjacency matrix.
mask (~torch.Tensor, optional): Input-node validity mask
:math:`\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}` with
:obj:`True` on real (non-padded) nodes. (default: :obj:`None`)
Returns:
dict: A dictionary with the different terms of
the auxiliary loss.
"""
# Motif adj matrix - not sym. normalised
motif_adj = torch.matmul(torch.matmul(adj, adj), adj)
motif_adj_pool = torch.matmul(torch.matmul(S.transpose(1, 2), motif_adj), S)
cut_loss = ho_cut_loss = 0
# 1st order MinCUT loss
if self.alpha < 1:
cut_loss = mincut_loss(adj, S, adj_pool)
cut_loss = 1 / self.k * cut_loss
# Higher order cut
if self.alpha > 0:
ho_cut_loss = mincut_loss(motif_adj, S, motif_adj_pool)
ho_cut_loss = 1 / self.k * ho_cut_loss
# Combine ho and fo mincut loss.
hosc_loss = (1 - self.alpha) * cut_loss + self.alpha * ho_cut_loss
# Orthogonality loss
if self.mu == 0:
ortho_loss = torch.tensor(0)
elif self.hosc_ortho:
# Hosc orthogonality regularization
ortho_loss = hosc_orthogonality_loss(S, mask, batch_reduction="mean")
else:
# Standard orthogonality regularization of MinCutPool
ortho_loss = orthogonality_loss(S, batch_reduction="mean")
return {"hosc_loss": hosc_loss, "ortho_loss": self.mu * ortho_loss}
[docs]
def compute_sparse_loss(
self,
edge_index: Adj,
edge_weight: Optional[Tensor],
S: Tensor,
batch: Optional[Tensor],
adj_pool: Optional[Adj] = None,
edge_weight_pool: Optional[Tensor] = None,
batch_pooled: Optional[Tensor] = None,
) -> dict:
r"""Computes the auxiliary loss terms for unbatched (sparse) mode.
This method is used when ``batched=False`` and operates on sparse
adjacency matrices.
Args:
edge_index (~torch_geometric.typing.Adj): Graph connectivity in sparse format.
edge_weight (~torch.Tensor, optional): Edge weights of shape :math:`(E,)`.
S (~torch.Tensor): The dense assignment matrix of shape :math:`(N, K)`.
batch (~torch.Tensor, optional): Batch vector of shape :math:`(N,)`.
adj_pool (~torch_geometric.typing.Adj, optional): The postprocessed pooled
adjacency. When ``self.sparse_output=True``, an ``edge_index`` of
shape :math:`(2, E_\\text{pool})` over the block-diagonal supernode
graph. When ``self.sparse_output=False``, a dense tensor of shape
:math:`(B, K, K)`. Required when ``alpha < 1``.
edge_weight_pool (~torch.Tensor, optional): Edge weights of the
postprocessed pooled adjacency, of shape :math:`(E_\\text{pool},)`.
Required when ``self.sparse_output=True`` and ``alpha < 1``.
batch_pooled (~torch.Tensor, optional): Batch vector for the pooled
supernodes. Required when ``self.sparse_output=True``,
``alpha < 1`` and the input contains multiple graphs.
Returns:
dict: A dictionary with ``'hosc_loss'`` and ``'ortho_loss'``.
"""
edge_index_conv, edge_weight_conv = connectivity_to_edge_index(
edge_index, edge_weight
)
num_nodes = S.size(0)
device = S.device
if batch is None:
batch = torch.zeros(num_nodes, dtype=torch.long, device=device)
batch_size = int(batch.max().item()) + 1
cut_loss = torch.tensor(0.0, device=device, dtype=S.dtype)
ho_cut_loss = torch.tensor(0.0, device=device, dtype=S.dtype)
if self.alpha < 1:
if self.sparse_output:
# Sparse block-diagonal pooled adj: Tr per graph =
# sum of self-loop weights, scattered over `batch_pooled`.
src, dst = adj_pool[0], adj_pool[1]
loop_mask = src == dst
contrib = edge_weight_pool.view(-1) * loop_mask.to(
edge_weight_pool.dtype
)
if batch_pooled is None:
edge_batch = torch.zeros(
src.numel(), dtype=torch.long, device=device
)
else:
edge_batch = batch_pooled[src]
num_per_graph = scatter(
contrib, edge_batch, dim=0, dim_size=batch_size, reduce="sum"
)
else:
# Dense [B, K, K] pooled adj: Tr per graph = sum of diagonals.
num_per_graph = torch.diagonal(adj_pool, dim1=-2, dim2=-1).sum(dim=-1)
cut_loss = sparse_mincut_loss(
edge_index_conv,
S,
edge_weight_conv,
batch,
batch_reduction="mean",
num_per_graph=num_per_graph,
)
cut_loss = cut_loss / self.k
if self.alpha > 0:
ho_cut_loss = sparse_ho_mincut_loss(
edge_index_conv,
S,
edge_weight_conv,
batch,
batch_reduction="mean",
)
ho_cut_loss = ho_cut_loss / self.k
hosc_loss = (1 - self.alpha) * cut_loss + self.alpha * ho_cut_loss
if self.mu == 0:
ortho_loss = torch.tensor(0.0, device=device, dtype=S.dtype)
elif self.hosc_ortho:
ortho_loss = unbatched_hosc_orthogonality_loss(
S, batch, batch_reduction="mean"
)
else:
ortho_loss = unbatched_orthogonality_loss(S, batch, batch_reduction="mean")
return {"hosc_loss": hosc_loss, "ortho_loss": self.mu * ortho_loss}
def extra_repr_args(self) -> dict:
return {
"batched": self.batched,
"mu": self.mu,
"alpha": self.alpha,
"hosc_ortho": self.hosc_ortho,
}
