torch_geometric.nn.aggr.SoftmaxAggregation

class SoftmaxAggregation(t: float = 1.0, learn: bool = False, semi_grad: bool = False, channels: int = 1)[source]

Bases: Aggregation

The softmax aggregation operator based on a temperature term, as described in the “DeeperGCN: All You Need to Train Deeper GCNs” paper.

\[\mathrm{softmax}(\mathcal{X}|t) = \sum_{\mathbf{x}_i\in\mathcal{X}} \frac{\exp(t\cdot\mathbf{x}_i)}{\sum_{\mathbf{x}_j\in\mathcal{X}} \exp(t\cdot\mathbf{x}_j)}\cdot\mathbf{x}_{i},\]

where \(t\) controls the softness of the softmax when aggregating over a set of features \(\mathcal{X}\).

Parameters:

  • t (float, optional) – Initial inverse temperature for softmax aggregation. (default: 1.0)

  • learn (bool, optional) – If set to True, will learn the value t for softmax aggregation dynamically. (default: False)

  • semi_grad (bool, optional) – If set to True, will turn off gradient calculation during softmax computation. Therefore, only semi-gradients are used during backpropagation. Useful for saving memory and accelerating backward computation when t is not learnable. (default: False)

  • channels (int, optional) – Number of channels to learn from \(t\). If set to a value greater than 1, \(t\) will be learned per input feature channel. This requires compatible shapes for the input to the forward calculation. (default: 1)

reset_parameters()[source]

Resets all learnable parameters of the module.

forward(x: Tensor, index: Optional[Tensor] = None, ptr: Optional[Tensor] = None, dim_size: Optional[int] = None, dim: int = -2) → Tensor[source]

Forward pass.

Parameters:

  • x (torch.Tensor) – The source tensor.

  • index (torch.Tensor, optional) – The indices of elements for applying the aggregation. One of index or ptr must be defined. (default: None)

  • ptr (torch.Tensor, optional) – If given, computes the aggregation based on sorted inputs in CSR representation. One of index or ptr must be defined. (default: None)

  • dim_size (int, optional) – The size of the output tensor at dimension dim after aggregation. (default: None)

  • dim (int, optional) – The dimension in which to aggregate. (default: -2)

  • max_num_elements – (int, optional): The maximum number of elements within a single aggregation group. (default: None)