TransE¶
-
class
dgl.nn.pytorch.link.
TransE
(num_rels, feats, p=1)[source]¶ Bases:
torch.nn.modules.module.Module
Similarity measure from Translating Embeddings for Modeling Multi-relational Data
Mathematically, it is defined as follows:
\[- {\| h + r - t \|}_p\]where \(h\) is the head embedding, \(r\) is the relation embedding, and \(t\) is the tail embedding.
- Parameters
-
rel_emb
¶ The learnable relation type embedding.
- Type
torch.nn.Embedding
Examples
>>> import dgl >>> import torch as th >>> from dgl.nn import TransE
>>> # input features >>> num_nodes = 10 >>> num_edges = 30 >>> num_rels = 3 >>> feats = 4
>>> scorer = TransE(num_rels=num_rels, feats=feats) >>> g = dgl.rand_graph(num_nodes=num_nodes, num_edges=num_edges) >>> src, dst = g.edges() >>> h = th.randn(num_nodes, feats) >>> h_head = h[src] >>> h_tail = h[dst] >>> # Randomly initialize edge relation types for demonstration >>> rels = th.randint(low=0, high=num_rels, size=(num_edges,)) >>> scorer(h_head, h_tail, rels).shape torch.Size([30])
-
forward
(h_head, h_tail, rels)[source]¶ Score triples.
- Parameters
h_head (torch.Tensor) – Head entity features. The tensor is of shape \((E, D)\), where \(E\) is the number of triples, and \(D\) is the feature size.
h_tail (torch.Tensor) – Tail entity features. The tensor is of shape \((E, D)\), where \(E\) is the number of triples, and \(D\) is the feature size.
rels (torch.Tensor) – Relation types. It is a LongTensor of shape \((E)\), where \(E\) is the number of triples.
- Returns
The triple scores. The tensor is of shape \((E)\).
- Return type
torch.Tensor