EdgePredictor

class dgl.nn.pytorch.link.EdgePredictor(op, in_feats=None, out_feats=None, bias=False)

Bases: Module

用于节点表示对的预测器/评分函数

给定一对节点表示，\(h_i\) 和 \(h_j\)，它将它们与

点积

\[h_i^{T} h_j\]

或 余弦相似度

\[\frac{h_i^{T} h_j}{{\| h_i \|}_2 \cdot {\| h_j \|}_2}\]

或 元素乘积

\[h_i \odot h_j\]

或 连接

\[h_i \Vert h_j\]

可选地，它将组合结果传递给线性层以进行最终预测。

Parameters:

• op (str) – 要应用的操作。它可以是 'dot'、'cos'、'ele' 或 'cat'，分别对应上述公式。

• in_feats (int, optional) – 输入特征的大小，即 \(h_i\) 和 \(h_j\) 的大小。只有在需要应用线性层时才需要指定。

• out_feats (int, optional) – 输出特征大小。仅在需要应用线性层时才需要。

• bias (bool, optional) – 如果适用，是否为线性层使用偏置。

示例

>>> import dgl
>>> import torch as th
>>> from dgl.nn import EdgePredictor
>>> num_nodes = 2
>>> num_edges = 3
>>> in_feats = 4
>>> g = dgl.rand_graph(num_nodes=num_nodes, num_edges=num_edges)
>>> h = th.randn(num_nodes, in_feats)
>>> src, dst = g.edges()
>>> h_src = h[src]
>>> h_dst = h[dst]

案例1：点积

>>> predictor = EdgePredictor('dot')
>>> predictor(h_src, h_dst).shape
torch.Size([3, 1])
>>> predictor = EdgePredictor('dot', in_feats, out_feats=3)
>>> predictor.reset_parameters()
>>> predictor(h_src, h_dst).shape
torch.Size([3, 3])

案例2：余弦相似度

>>> predictor = EdgePredictor('cos')
>>> predictor(h_src, h_dst).shape
torch.Size([3, 1])
>>> predictor = EdgePredictor('cos', in_feats, out_feats=3)
>>> predictor.reset_parameters()
>>> predictor(h_src, h_dst).shape
torch.Size([3, 3])

案例3：元素乘积

>>> predictor = EdgePredictor('ele')
>>> predictor(h_src, h_dst).shape
torch.Size([3, 4])
>>> predictor = EdgePredictor('ele', in_feats, out_feats=3)
>>> predictor.reset_parameters()
>>> predictor(h_src, h_dst).shape
torch.Size([3, 3])

案例4：连接

>>> predictor = EdgePredictor('cat')
>>> predictor(h_src, h_dst).shape
torch.Size([3, 8])
>>> predictor = EdgePredictor('cat', in_feats, out_feats=3)
>>> predictor.reset_parameters()
>>> predictor(h_src, h_dst).shape
torch.Size([3, 3])

forward(h_src, h_dst)

Description

预测节点表示对。

param h_src:

源节点特征。张量的形状为 \((E, D_{in})\)， 其中 \(E\) 是边/节点对的数量，\(D_{in}\) 是输入特征的大小。

type h_src:

torch.Tensor

param h_dst:

目标节点特征。张量的形状为 \((E, D_{in})\)， 其中 \(E\) 是边/节点对的数量，\(D_{in}\) 是输入特征的大小。

type h_dst:

torch.Tensor

returns:

输出特征。

rtype:

torch.Tensor

reset_parameters()

Description

重新初始化可学习的参数。