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
  • num_rels (int) – Number of relation types.

  • feats (int) – Embedding size.

  • p (int, optional) – The p to use for Lp norm, which can be 1 or 2.

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

reset_parameters()[source]

Reinitialize learnable parameters.