torch_geometric.nn.pool.TopKPooling

class TopKPooling(in_channels: int, ratio: Union[int, float] = 0.5, min_score: Optional[float] = None, multiplier: float = 1.0, nonlinearity: Union[str, Callable] = 'tanh')

Bases: Module

\(\mathrm{top}_k\) 池化操作符来自 “Graph U-Nets”, “Towards Sparse Hierarchical Graph Classifiers” 和 “Understanding Attention and Generalization in Graph Neural Networks” 论文。

If min_score \(\tilde{\alpha}\) is None, computes:

\[ \begin{align}\begin{aligned}\mathbf{y} &= \sigma \left( \frac{\mathbf{X}\mathbf{p}}{\| \mathbf{p} \|} \right)\\\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} \]

If min_score \(\tilde{\alpha}\) is a value in [0, 1], computes:

\[ \begin{align}\begin{aligned}\mathbf{y} &= \mathrm{softmax}(\mathbf{X}\mathbf{p})\\\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} \]

其中节点基于可学习的投影分数 \(\mathbf{p}\)被丢弃。

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)

• min_score (float, optional) – Minimal node score \(\tilde{\alpha}\) which is used to compute indices of pooled nodes \(\mathbf{i} = \mathbf{y}_i > \tilde{\alpha}\). When this value is not None, the ratio argument is ignored. (default: None)

• multiplier (float, optional) – Coefficient by which features gets multiplied after pooling. This can be useful for large graphs and when min_score is used. (default: 1)

• 非线性 (str 或 可调用, 可选) – 非线性函数 \(\sigma\). (默认: "tanh")

reset_parameters()

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

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]

前向传播。

Parameters:

• x (torch.Tensor) – The node feature matrix.

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

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

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

• attn (torch.Tensor, optional) – Optional node-level matrix to use for computing attention scores instead of using the node feature matrix x. (default: None)

Return type:

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