from typing import Callable, Optional, Union
from torch import Tensor
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 TopkPooling(SRCPooling):
r"""The :math:`\mathrm{top}_k` pooling operator from the papers `"Graph U-Nets"
<https://arxiv.org/abs/1905.05178>`_ (Gao & Ji, ICML 2019),
`"Towards Sparse Hierarchical Graph Classifiers" <https://arxiv.org/abs/1811.01287>`_ (Cangea et al., 2018),
and `"Understanding Attention and Generalization in Graph Neural
Networks" <https://arxiv.org/abs/1905.02850>`_ (Knyazev et al., NeurIPS 2019).
Args:
in_channels (int):
Size of each input sample.
ratio (float or int):
The 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``)
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`)
"""
def __init__(
self,
in_channels: int,
ratio: Union[int, float] = 0.5,
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,
):
super().__init__(
selector=TopkSelect(
in_channels=in_channels,
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,
),
)
self.multiplier = multiplier
[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
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(
so=so,
edge_index=adj,
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,
}
