JumpingKnowledge

class dgl.nn.pytorch.utils.JumpingKnowledge(mode='cat', in_feats=None, num_layers=None)

基类: Module

来自 Representation Learning on Graphs with Jumping Knowledge Networks 的 Jumping Knowledge 聚合模块

它通过以下方式聚合多个 GNN 层的输出表示：

拼接

\[h_i^{(1)} \, \Vert \, \ldots \, \Vert \, h_i^{(T)}\]

或 **最大池化**

\[\max \left( h_i^{(1)}, \ldots, h_i^{(T)} \right)\]

或 **LSTM**

\[\sum_{t=1}^T \alpha_i^{(t)} h_i^{(t)}\]

其中注意力分数 \(\alpha_i^{(t)}\) 通过 BiLSTM 获得

参数:

• mode (str) – 要应用的聚合方式。可以是 'cat'、'max' 或 'lstm'，分别对应上面的公式。

• in_feats (int, optional) – 仅当 mode 为 'lstm' 时需要此参数。单个 GNN 层的输出表示大小。注意，所有 GNN 层需要具有相同的输出表示大小。

• num_layers (int, optional) – 仅当 mode 为 'lstm' 时需要此参数。用于输出聚合的 GNN 层数。

示例

>>> import dgl
>>> import torch as th
>>> from dgl.nn import JumpingKnowledge

>>> # Output representations of two GNN layers
>>> num_nodes = 3
>>> in_feats = 4
>>> feat_list = [th.zeros(num_nodes, in_feats), th.ones(num_nodes, in_feats)]

>>> # Case1
>>> model = JumpingKnowledge()
>>> model(feat_list).shape
torch.Size([3, 8])

>>> # Case2
>>> model = JumpingKnowledge(mode='max')
>>> model(feat_list).shape
torch.Size([3, 4])

>>> # Case3
>>> model = JumpingKnowledge(mode='max', in_feats=in_feats, num_layers=len(feat_list))
>>> model(feat_list).shape
torch.Size([3, 4])

forward(feat_list)

描述

聚合多个 GNN 层的输出表示。

参数 feat_list:

feat_list[i] 是一个 GNN 层的输出表示。

类型 feat_list:

list[Tensor]

返回:

聚合后的表示。

返回类型:

Tensor

reset_parameters()

描述

重新初始化可学习参数。此操作仅在 lstm 模式下生效。