torch_geometric.nn.pool.SAGPooling

class SAGPooling(in_channels: int, ratio: ~typing.Union[float, int] = 0.5, GNN: ~torch.nn.modules.module.Module = <class 'torch_geometric.nn.conv.graph_conv.GraphConv'>, min_score: ~typing.Optional[float] = None, multiplier: float = 1.0, nonlinearity: ~typing.Union[str, ~typing.Callable] = 'tanh', **kwargs)[source]

Bases: Module

自注意力池化操作符来自“自注意力图池化”和“理解图神经网络中的注意力和泛化”论文。

如果 min_score \(\tilde{\alpha}\) 是 None，则计算：

\[ \begin{align}\begin{aligned}\mathbf{y} &= \textrm{GNN}(\mathbf{X}, \mathbf{A})\\\mathbf{i} &= \mathrm{top}_k(\mathbf{y})\\\mathbf{X}^{\prime} &= (\mathbf{X} \odot \mathrm{tanh}(\mathbf{y}))_{\mathbf{i}}\\\mathbf{A}^{\prime} &= \mathbf{A}_{\mathbf{i},\mathbf{i}}\end{aligned}\end{align} \]

如果 min_score \(\tilde{\alpha}\) 是 [0, 1] 中的一个值， 计算：

\[ \begin{align}\begin{aligned}\mathbf{y} &= \mathrm{softmax}(\textrm{GNN}(\mathbf{X},\mathbf{A}))\\\mathbf{i} &= \mathbf{y}_i > \tilde{\alpha}\\\mathbf{X}^{\prime} &= (\mathbf{X} \odot \mathbf{y})_{\mathbf{i}}\\\mathbf{A}^{\prime} &= \mathbf{A}_{\mathbf{i},\mathbf{i}}.\end{aligned}\end{align} \]

投影分数是基于图神经网络层学习的。

Parameters:

  • in_channels (int) – Size of each input sample.

  • ratio (float 或 int) – 图池化比例，用于计算 \(k = \lceil \mathrm{ratio} \cdot N \rceil\)，或者直接作为 \(k\) 的值，具体取决于 ratio 的类型是 float 还是 int。 如果 min_score 不是 None，则忽略此值。 (默认值: 0.5)

  • GNN (torch.nn.Module, optional) – 用于计算投影分数的图神经网络层（可以是 torch_geometric.nn.conv.GraphConv, torch_geometric.nn.conv.GCNConv, torch_geometric.nn.conv.GATConv 或 torch_geometric.nn.conv.SAGEConv）。 (默认: torch_geometric.nn.conv.GraphConv)

  • min_score (float, optional) – 最小节点分数 \(\tilde{\alpha}\) 用于计算池化节点的索引 \(\mathbf{i} = \mathbf{y}_i > \tilde{\alpha}\). 当此值不为 None 时，ratio 参数将被忽略。（默认值：None）

  • multiplier (float, optional) – 池化后特征乘以的系数。这对于大型图和使用 min_score 时非常有用。（默认值：1）

  • 非线性 (str 或 可调用, 可选) – 使用的非线性函数。 (默认: "tanh")

  • **kwargs (可选) – 用于初始化图神经网络层的额外参数。

reset_parameters()[source]

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

forward(x: Tensor, edge_index: Tensor, edge_attr: Optional[Tensor] = None, batch: Optional[Tensor] = None, attn: Optional[Tensor] = None) → Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor], Tensor, Tensor][source]

前向传播。

Parameters:

  • x (torch.Tensor) – 节点特征矩阵。

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

  • edge_attr (torch.Tensor, optional) – The edge features. (default: None)

  • batch (torch.Tensor, optional) – 批次向量 \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), 它将 每个节点分配到一个特定的示例中。(默认值: None)

  • attn (torch.Tensor, optional) – 可选的节点级矩阵，用于计算注意力分数，而不是使用节点特征矩阵 x。（默认值：None）

Return type:

Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor], Tensor, Tensor]