from typing import Callable, Optional
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 EdgeContractionSelect, SelectOutput
from tgp.src import PoolingOutput, SRCPooling
from tgp.utils.ops import connectivity_to_edge_index
from tgp.utils.typing import ConnectionType, LiftType, ReduceType, SinvType
[docs]
class EdgeContractionPooling(SRCPooling):
r"""The edge pooling operator from the papers `"Towards Graph Pooling by Edge
Contraction" <https://graphreason.github.io/papers/17.pdf>`_ (Diehl et al. 2019) and
`"Edge Contraction Pooling for Graph Neural Networks"
<https://arxiv.org/abs/1905.10990>`_ (Diehl, 2019).
This implementation is based on the paper `"Revisiting Edge Pooling in Graph Neural Networks"
<https://www.esann.org/sites/default/files/proceedings/2022/ES2022-92.pdf>`_ (Landolfi, 2022).
+ The :math:`\texttt{select}` operator is implemented with :class:`~tgp.select.EdgeContractionSelect`.
+ 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`.
To duplicate the configuration of the paper `"Towards Graph Pooling by Edge
Contraction" <https://graphreason.github.io/papers/17.pdf>`_ (Diehl et al. 2019), use
either :func:`~tgp.select.EdgeContractionSelect.compute_edge_score_softmax`
or :func:`~tgp.select.EdgeContractionSelect.compute_edge_score_tanh`, and set
``add_to_edge_score`` to ``0.0``. To duplicate the configuration of the paper
`"Edge Contraction Pooling for Graph Neural Networks"
<https://arxiv.org/abs/1905.10990>`_ (Diehl, 2019), set ``dropout`` to ``0.2``.
Args:
in_channels (int): Size of each input sample.
edge_score_method (callable, optional): The function to apply to compute the edge
score from raw edge scores. By default, this is the softmax over all incoming
edges for each node. This function takes in a ``raw_edge_score`` tensor of
shape ``[num_nodes]``, an ``edge_index`` tensor and the number of nodes
``num_nodes``, and produces a new tensor of the same size as
``raw_edge_score`` describing normalized edge scores. Included functions are
:func:`~tgp.select.EdgeContractionSelect.compute_edge_score_softmax`,
:func:`~tgp.select.EdgeContractionSelect.compute_edge_score_tanh`, and
:func:`~tgp.select.EdgeContractionSelect.compute_edge_score_sigmoid`.
(default: :func:`~tgp.select.EdgeContractionSelect.compute_edge_score_softmax`)
dropout (float, optional): The probability with which to drop edge scores during
training. (default: ``0.0``)
add_to_edge_score (float, optional): A value to be added to each computed edge
score. Adding this greatly helps with unpooling stability. (default: ``0.5``)
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,
edge_score_method: Optional[Callable] = None,
dropout: Optional[float] = 0.0,
add_to_edge_score: float = 0.5,
lift: LiftType = "precomputed",
s_inv_op: SinvType = "transpose",
connect_red_op: ConnectionType = "sum",
lift_red_op: ReduceType = "sum",
remove_self_loops: bool = True,
degree_norm: bool = False,
edge_weight_norm: bool = False,
):
super().__init__(
selector=EdgeContractionSelect(
in_channels=in_channels,
edge_score_method=edge_score_method,
dropout=dropout,
add_to_edge_score=add_to_edge_score,
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,
remove_self_loops=remove_self_loops,
degree_norm=degree_norm,
edge_weight_norm=edge_weight_norm,
),
)
[docs]
def forward(
self,
x: Tensor,
adj: Optional[Adj] = None,
edge_weight: Optional[Tensor] = None,
so: Optional[SelectOutput] = None,
batch: 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`)
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
edge_index, edge_weight = connectivity_to_edge_index(adj, edge_weight)
so = self.select(x=x, edge_index=edge_index, batch=batch)
# Reduce
x_pooled, batch_pooled = self.reduce(x=x, so=so, batch=batch)
# Connect
edge_index_pooled, edge_weight_pooled = self.connect(
edge_index=edge_index,
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
