torch_geometric.nn.pool.EdgePooling

class EdgePooling(in_channels: int, edge_score_method: Optional[Callable] = None, dropout: float = 0.0, add_to_edge_score: float = 0.5)[source]

Bases: Module

The edge pooling operator from the “Towards Graph Pooling by Edge Contraction” and “Edge Contraction Pooling for Graph Neural Networks” papers.

In short, a score is computed for each edge. Edges are contracted iteratively according to that score unless one of their nodes has already been part of a contracted edge.

To duplicate the configuration from the “Towards Graph Pooling by Edge Contraction” paper, use either EdgePooling.compute_edge_score_softmax() or EdgePooling.compute_edge_score_tanh(), and set add_to_edge_score to 0.0.

To duplicate the configuration from the “Edge Contraction Pooling for Graph Neural Networks” paper, set dropout to 0.2.

Parameters:

  • 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 EdgePooling.compute_edge_score_softmax(), EdgePooling.compute_edge_score_tanh(), and EdgePooling.compute_edge_score_sigmoid(). (default: EdgePooling.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)

reset_parameters()[source]

Resets all learnable parameters of the module.

static compute_edge_score_softmax(raw_edge_score: Tensor, edge_index: Tensor, num_nodes: int) → Tensor[source]

Normalizes edge scores via softmax application.

Return type:

Tensor

static compute_edge_score_tanh(raw_edge_score: Tensor, edge_index: Optional[Tensor] = None, num_nodes: Optional[int] = None) → Tensor[source]

Normalizes edge scores via hyperbolic tangent application.

Return type:

Tensor

static compute_edge_score_sigmoid(raw_edge_score: Tensor, edge_index: Optional[Tensor] = None, num_nodes: Optional[int] = None) → Tensor[source]

Normalizes edge scores via sigmoid application.

Return type:

Tensor

forward(x: Tensor, edge_index: Tensor, batch: Tensor) → Tuple[Tensor, Tensor, Tensor, UnpoolInfo][source]

Forward pass.

Parameters:

  • x (torch.Tensor) – The node features.

  • edge_index (torch.Tensor) – The edge indices.

  • batch (torch.Tensor) – The batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which assigns each node to a specific example.

Return types:

  • x (torch.Tensor) - The pooled node features.

  • edge_index (torch.Tensor) - The coarsened edge indices.

  • batch (torch.Tensor) - The coarsened batch vector.

  • unpool_info (UnpoolInfo) - Information that is consumed by EdgePooling.unpool() for unpooling.

Return type:

Tuple[Tensor, Tensor, Tensor, UnpoolInfo]

unpool(x: Tensor, unpool_info: UnpoolInfo) → Tuple[Tensor, Tensor, Tensor][source]

Unpools a previous edge pooling step.

For unpooling, x should be of same shape as those produced by this layer’s forward() function. Then, it will produce an unpooled x in addition to edge_index and batch.

Parameters:
Return types:

  • x (torch.Tensor) - The unpooled node features.

  • edge_index (torch.Tensor) - The new edge indices.

  • batch (torch.Tensor) - The new batch vector.

Return type:

Tuple[Tensor, Tensor, Tensor]