torch_geometric.nn.aggr.AttentionalAggregation

class AttentionalAggregation(gate_nn: Module, nn: Optional[Module] = None)[source]

Bases: Aggregation

来自“Graph Matching Networks for Learning the Similarity of Graph Structured Objects”论文的软注意力聚合层。

\[\mathbf{r}_i = \sum_{n=1}^{N_i} \mathrm{softmax} \left( h_{\mathrm{gate}} ( \mathbf{x}_n ) \right) \cdot h_{\mathbf{\Theta}} ( \mathbf{x}_n ),\]

其中 \(h_{\mathrm{gate}} \colon \mathbb{R}^F \to \mathbb{R}\) 和 \(h_{\mathbf{\Theta}}\) 表示神经网络，即 多层感知机（MLPs）。

Parameters:

  • gate_nn (torch.nn.Module) – 一个神经网络 \(h_{\mathrm{gate}}\) 通过将节点特征 x 从形状 [-1, in_channels] 映射到形状 [-1, 1]（用于节点级门控）或 [1, out_channels]（用于特征级门控）来计算注意力分数，例如，由 torch.nn.Sequential 定义。

  • nn (torch.nn.Module, optional) – 一个神经网络 \(h_{\mathbf{\Theta}}\)，它将节点特征 x 从形状 [-1, in_channels] 映射到形状 [-1, out_channels]， 然后在与注意力分数结合之前，例如，由 torch.nn.Sequential 定义。（默认值：None）

reset_parameters()[source]

重置模块的所有可学习参数。

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

前向传播。

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)

Return type:

Tensor