跳跃知识

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

Bases: Module

跳跃知识聚合模块来自使用跳跃知识网络的图表示学习

它聚合了多个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)}\]

使用从BiLSTM获得的注意力分数 \(\alpha_i^{(t)}\)

Parameters:

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

• in_feats (int, 可选) – 仅当 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)

Description

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

param feat_list:

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

type feat_list:

列表[张量]

returns:

聚合表示。

rtype:

张量

reset_parameters()

Description

重新初始化可学习的参数。这仅在lstm模式下生效。