EdgePredictor

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

Bases: Module

Predictor/score function for pairs of node representations

Given a pair of node representations, \(h_i\) and \(h_j\), it combines them with

dot product

\[h_i^{T} h_j\]

or cosine similarity

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

or elementwise product

\[h_i \odot h_j\]

or concatenation

\[h_i \Vert h_j\]

Optionally, it passes the combined results to a linear layer for the final prediction.

Parameters:

  • op (str) – The operation to apply. It can be β€˜dot’, β€˜cos’, β€˜ele’, or β€˜cat’, corresponding to the equations above in order.

  • in_feats (int, optional) – The input feature size of \(h_i\) and \(h_j\). It is required only if a linear layer is to be applied.

  • out_feats (int, optional) – The output feature size. It is reuiqred only if a linear layer is to be applied.

  • bias (bool, optional) – Whether to use bias for the linear layer if it applies.

Examples

>>> 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]

Case1: dot product

>>> 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])

Case2: cosine similarity

>>> 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])

Case3: elementwise product

>>> 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])

Case4: concatenation

>>> 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)[source]

Description

Predict for pairs of node representations.

param h_src:

Source node features. The tensor is of shape \((E, D_{in})\), where \(E\) is the number of edges/node pairs, and \(D_{in}\) is the input feature size.

type h_src:

torch.Tensor

param h_dst:

Destination node features. The tensor is of shape \((E, D_{in})\), where \(E\) is the number of edges/node pairs, and \(D_{in}\) is the input feature size.

type h_dst:

torch.Tensor

returns:

The output features.

rtype:

torch.Tensor

reset_parameters()[source]

Description

Reinitialize learnable parameters.