torch_geometric.nn.kge.RotatE

class RotatE(num_nodes: int, num_relations: int, hidden_channels: int, margin: float = 1.0, sparse: bool = False)

Bases: KGEModel

The RotatE model from the “RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space” paper.

RotatE models relations as a rotation in complex space from head to tail such that

\[\mathbf{e}_t = \mathbf{e}_h \circ \mathbf{e}_r,\]

resulting in the scoring function

\[d(h, r, t) = - {\| \mathbf{e}_h \circ \mathbf{e}_r - \mathbf{e}_t \|}_p\]

Note

For an example of using the RotatE 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 (float, optional) – The margin of the ranking loss.

• sparse (bool, optional) – If set to True, gradients w.r.t. to the embedding matrices will be sparse. (default: False)

reset_parameters()

Resets all learnable parameters of the module.

forward(head_index: Tensor, rel_type: Tensor, tail_index: Tensor) → Tensor

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.

loss(head_index: Tensor, rel_type: Tensor, tail_index: Tensor) → Tensor

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.