EdgePredictor

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

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}{{\Vert h_i\Vert }_2\cdot {\Vert h_j\Vert }_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)

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

Description

Reinitialize learnable parameters.