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