TransRο
- class dgl.nn.pytorch.link.TransR(num_rels, rfeats, nfeats, p=1)[source]ο
Bases:
Module
Similarity measure from Learning entity and relation embeddings for knowledge graph completion
Mathematically, it is defined as follows:
\[- {\| M_r h + r - M_r t \|}_p\]where \(M_r\) is a relation-specific projection matrix, \(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
- rel_projectο
The learnable relation-type-specific projection.
- Type:
torch.nn.Embedding
Examples
>>> import dgl >>> import torch as th >>> from dgl.nn import TransR
>>> # input features >>> num_nodes = 10 >>> num_edges = 30 >>> num_rels = 3 >>> feats = 4
>>> scorer = TransR(num_rels=num_rels, rfeats=2, nfeats=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