torch_geometric.nn.kge.TransE
- class TransE(num_nodes: int, num_relations: int, hidden_channels: int, margin: float = 1.0, p_norm: float = 1.0, sparse: bool = False)[source]
Bases:
KGEModel
The TransE model from the “Translating Embeddings for Modeling Multi-Relational Data” paper.
TransE
models relations as a translation from head to tail entities such that\[\mathbf{e}_h + \mathbf{e}_r \approx \mathbf{e}_t,\]resulting in the scoring function:
\[d(h, r, t) = - {\| \mathbf{e}_h + \mathbf{e}_r - \mathbf{e}_t \|}_p\]Note
For an example of using the
TransE
model, see examples/kge_fb15k_237.py.- Parameters:
num_nodes (int) – The number of nodes/entities in the graph.
num_relations (int) – The number of relations in the graph.
hidden_channels (int) – The hidden embedding size.
margin (int, optional) – The margin of the ranking loss. (default:
1.0
)p_norm (int, optional) – The order embedding and distance normalization. (default:
1.0
)sparse (bool, optional) – If set to
True
, gradients w.r.t. to the embedding matrices will be sparse. (default:False
)
- forward(head_index: Tensor, rel_type: Tensor, tail_index: Tensor) Tensor [source]
Returns the score for the given triplet.
- Parameters:
head_index (torch.Tensor) – The head indices.
rel_type (torch.Tensor) – The relation type.
tail_index (torch.Tensor) – The tail indices.
- Return type:
- loss(head_index: Tensor, rel_type: Tensor, tail_index: Tensor) Tensor [source]
Returns the loss value for the given triplet.
- Parameters:
head_index (torch.Tensor) – The head indices.
rel_type (torch.Tensor) – The relation type.
tail_index (torch.Tensor) – The tail indices.
- Return type: