# EdgePredictor¶

Bases: torch.nn.modules.module.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]

Predict for pairs of node representations.

Parameters
• h_src (torch.Tensor) – 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.

• h_dst (torch.Tensor) – 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.

Returns

The output features.

Return type

torch.Tensor

reset_parameters()[source]

Reinitialize learnable parameters.