import inspect
from typing import Callable, Optional, Union
import torch
from torch import Tensor
from torch_geometric.nn import GraphConv
from torch_geometric.typing import Adj
from tgp.connect import SparseConnect
from tgp.lift import BaseLift
from tgp.reduce import BaseReduce
from tgp.select import SelectOutput, TopkSelect
from tgp.src import PoolingOutput, SRCPooling
from tgp.utils.typing import LiftType, ReduceType, SinvType
[docs]
class SAGPooling(SRCPooling):
r"""The self-attention pooling operator from the paper `"Self-Attention Graph Pooling" <https://arxiv.org/abs/1904.08082>`_ (Lee et al., ICML 2019).
It computes the attention scores :math:`\mathbf{a}` top-:math:`k` selector as:
.. math::
\mathbf{a} = \textrm{GNN}(\mathbf{X}, \mathbf{A})
+ The :math:`\texttt{select}` operator is implemented with :class:`~tgp.select.TopkSelect`.
+ The :math:`\texttt{reduce}` operator is implemented with :class:`~tgp.reduce.BaseReduce`.
+ The :math:`\texttt{connect}` operator is implemented with :class:`~tgp.connect.SparseConnect`.
+ The :math:`\texttt{lift}` operator is implemented with :class:`~tgp.lift.BaseLift`.
Args:
in_channels (int):
Size of each input sample.
ratio (float or int): Graph pooling ratio, which is used to compute
:math:`k = \lceil \mathrm{ratio} \cdot N \rceil`, or the value
of :math:`k` itself, depending on whether the type of ``ratio``
is :obj:`float` or :obj:`int`.
This value is ignored if ``min_score`` is not :obj:`None`.
(default: ``0.5``)
GNN (~torch.nn.Module, optional): A graph neural network layer for
calculating projection scores (one of
:class:`~torch_geometric.nn.conv.GraphConv`,
:class:`~torch_geometric.nn.conv.GCNConv`,
:class:`~torch_geometric.nn.conv.GATConv` or
:class:`~torch_geometric.nn.conv.SAGEConv`). (default:
:class:`~torch_geometric.nn.conv.GraphConv`)
min_score (float, optional):
Minimal node score :math:`\tilde{\alpha}`
which is used to compute indices of pooled nodes
:math:`\mathbf{i} = \mathbf{s}_i > \tilde{\alpha}`.
When this value is not :obj:`None`, the ``ratio`` argument is
ignored. (default: :obj:`None`)
multiplier (float, optional):
Coefficient by which features gets
multiplied after pooling. This can be useful for large graphs and
when ``min_score`` is used. (default: ``1``)
nonlinearity (str or callable, optional):
The non-linearity to use when computing the score.
(default: ``"tanh"``)
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}`.
connect_red_op (~tgp.utils.typing.ConnectionType, optional):
The aggregation function to be applied to all edges connecting nodes assigned
to supernodes :math:`i` and :math:`j`.
Can be any string of class :class:`~tgp.utils.typing.ConnectionType` admitted by
:obj:`~torch_geometric.utils.coalesce`,
e.g., ``'sum'``, ``'mean'``, ``'max'``)
(default: ``"sum"``)
lift_red_op (~tgp.utils.typing.ReduceType, optional):
The aggregation function to be applied to the lifted node features.
Can be any string of class :class:`~tgp.utils.typing.ReduceType` admitted by
:obj:`~torch_geometric.utils.scatter`,
e.g., ``'sum'``, ``'mean'``, ``'max'``)
(default: ``"sum"``)
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:`False`)
edge_weight_norm (bool, optional):
Whether to normalize the edge weights by dividing by the maximum absolute value per graph.
(default: :obj:`False`)
**kwargs (any, optional):
Additional parameters for initializing the graph
neural network layer.
"""
def __init__(
self,
in_channels: int,
ratio: Union[float, int] = 0.5,
GNN: Optional["torch.nn.Module"] = None,
min_score: Optional[float] = None,
multiplier: float = 1.0,
nonlinearity: Union[str, "Callable"] = "tanh",
lift: "LiftType" = "precomputed",
s_inv_op: "SinvType" = "transpose",
connect_red_op: "ReduceType" = "sum",
lift_red_op: "ReduceType" = "sum",
remove_self_loops: bool = True,
degree_norm: bool = False,
edge_weight_norm: bool = False,
**kwargs,
):
super().__init__(
selector=TopkSelect(
ratio=ratio, min_score=min_score, act=nonlinearity, s_inv_op=s_inv_op
),
reducer=BaseReduce(),
lifter=BaseLift(matrix_op=lift, reduce_op=lift_red_op),
connector=SparseConnect(
reduce_op=connect_red_op,
degree_norm=degree_norm,
edge_weight_norm=edge_weight_norm,
remove_self_loops=remove_self_loops,
),
)
# keep only the kwargs that are used in the GNN (signature works when __code__ is not available)
_gnn_cls = GNN or GraphConv
try:
_params = set(inspect.signature(_gnn_cls).parameters.keys())
except (ValueError, TypeError):
_params = set()
kwargs = {k: v for k, v in kwargs.items() if k in _params}
self.gnn = (GNN or GraphConv)(in_channels, 1, **kwargs)
self.multiplier = multiplier
def reset_parameters(self):
r"""Resets all learnable parameters of the module."""
self.gnn.reset_parameters()
super().reset_parameters()
[docs]
def forward(
self,
x: Tensor,
adj: Optional[Adj] = None,
edge_weight: Optional[Tensor] = None,
so: Optional[SelectOutput] = None,
batch: Optional[Tensor] = None,
attn: Optional[Tensor] = None,
lifting: bool = False,
**kwargs,
) -> PoolingOutput:
r"""Forward pass.
Args:
x (~torch.Tensor): The node feature matrix of shape :math:`[N, F]`,
where :math:`N` is the number of nodes in the batch and
:math:`F` is the number of node features.
adj (~torch_geometric.typing.Adj, optional): The connectivity matrix.
It can either be a :class:`~torch_sparse.SparseTensor` of (sparse) shape :math:`[N, N]`,
where :math:`N` is the number of nodes in the batch or a :obj:`~torch.Tensor` of shape
:math:`[2, E]`, where :math:`E` is the number of edges in the batch.
If ``lifting`` is :obj:`False`, it cannot be :obj:`None`.
(default: :obj:`None`)
edge_weight (~torch.Tensor, optional): A vector of shape
:math:`[E]` containing the weights of the edges.
(default: :obj:`None`)
so (~tgp.select.SelectOutput, optional): The output of the :math:`\texttt{select}` operator.
(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`)
attn (~torch.Tensor, optional):
Optional node-level matrix to use
for computing attention scores instead of using the node
feature matrix ``x``. (default: :obj:`None`)
lifting (bool, optional): If set to :obj:`True`, the :math:`\texttt{lift}` operation is performed.
(default: :obj:`False`)
Returns:
PoolingOutput: The output of the pooling operator.
"""
if lifting:
# Lift
x_lifted = self.lift(x_pool=x, so=so)
return x_lifted
else:
# Select
attn = x if attn is None else attn
attn = attn.view(-1, 1) if attn.dim() == 1 else attn
attn = self.gnn(attn, adj)
so = self.select(x=attn, batch=batch)
# Reduce
x_pooled, batch_pooled = self.reduce(x=x, so=so, batch=batch)
x_pooled = self.multiplier * x_pooled if self.multiplier != 1 else x_pooled
# Connect
edge_index_pooled, edge_weight_pooled = self.connect(
edge_index=adj,
so=so,
edge_weight=edge_weight,
batch_pooled=batch_pooled,
)
out = PoolingOutput(
x=x_pooled,
edge_index=edge_index_pooled,
edge_weight=edge_weight_pooled,
batch=batch_pooled,
so=so,
)
return out
def extra_repr_args(self) -> dict:
return {
"multiplier": self.multiplier,
}
