torch_geometric.nn.dense.DMoNPooling

class DMoNPooling(channels: Union[int, List[int]], k: int, dropout: float = 0.0)[source]

Bases: Module

来自“使用图神经网络进行图聚类”论文的光谱模块化池化操作符。

\[ \begin{align}\begin{aligned}\mathbf{X}^{\prime} &= {\mathrm{softmax}(\mathbf{S})}^{\top} \cdot \mathbf{X}\\\mathbf{A}^{\prime} &= {\mathrm{softmax}(\mathbf{S})}^{\top} \cdot \mathbf{A} \cdot \mathrm{softmax}(\mathbf{S})\end{aligned}\end{align} \]

基于密集学习分配 \(\mathbf{S} \in \mathbb{R}^{B \times N \times C}\). 返回学习到的聚类分配矩阵、池化的节点特征 矩阵、粗化的对称归一化邻接矩阵,以及三个 辅助目标:(1) 谱损失

\[\mathcal{L}_s = - \frac{1}{2m} \cdot{\mathrm{Tr}(\mathbf{S}^{\top} \mathbf{B} \mathbf{S})}\]

其中 \(\mathbf{B}\) 是模块化矩阵,(2) 正交性损失

\[\mathcal{L}_o = {\left\| \frac{\mathbf{S}^{\top} \mathbf{S}} {{\|\mathbf{S}^{\top} \mathbf{S}\|}_F} -\frac{\mathbf{I}_C}{\sqrt{C}} \right\|}_F\]

其中 \(C\) 是簇的数量,以及 (3) 簇损失

\[\mathcal{L}_c = \frac{\sqrt{C}}{n} {\left\|\sum_i\mathbf{C_i}^{\top}\right\|}_F - 1.\]

注意

有关使用 DMoNPooling 的示例,请参见 examples/proteins_dmon_pool.py

Parameters:
  • channels (intList[int]) – 每个输入样本的大小。如果以列表形式给出,将根据给定的特征大小构建一个MLP。

  • k (int) – 聚类的数量。

  • dropout (float, optional) – Dropout probability. (default: 0.0)

reset_parameters()[source]

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

forward(x: Tensor, adj: Tensor, mask: Optional[Tensor] = None) Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor][source]

前向传播。

Parameters:
  • x (torch.Tensor) – 节点特征张量 \(\mathbf{X} \in \mathbb{R}^{B \times N \times F}\), 其中 批次大小 \(B\), (最大) 节点数 \(N\) 对于 每个图,以及特征维度 \(F\). 请注意,聚类分配矩阵 \(\mathbf{S} \in \mathbb{R}^{B \times N \times C}\) 是 在此方法中创建的。

  • adj (torch.Tensor) – 邻接张量 \(\mathbf{A} \in \mathbb{R}^{B \times N \times N}\).

  • mask (torch.Tensor, optional) – Mask matrix \(\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}\) indicating the valid nodes for each graph. (default: None)

Return type:

(torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor)