TransEο
- class dgl.nn.pytorch.link.TransE(num_rels, feats, p=1)[source]ο
Bases:
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]ο
Descriptionο
Score triples.
- param h_head:
Head entity features. The tensor is of shape \((E, D)\), where \(E\) is the number of triples, and \(D\) is the feature size.
- type h_head:
torch.Tensor
- param h_tail:
Tail entity features. The tensor is of shape \((E, D)\), where \(E\) is the number of triples, and \(D\) is the feature size.
- type h_tail:
torch.Tensor
- param rels:
Relation types. It is a LongTensor of shape \((E)\), where \(E\) is the number of triples.
- type rels:
torch.Tensor
- returns:
The triple scores. The tensor is of shape \((E)\).
- rtype:
torch.Tensor