torch_geometric.nn.kge.RotatE

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

Bases: KGEModel

来自“RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space”论文的RotatE模型。

RotatE 将关系建模为从头部到尾部的复数空间中的旋转，使得

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

导致评分函数

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

注意

有关使用RotatE模型的示例，请参见 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) – 排名损失的边距。

  • sparse (bool, 可选) – 如果设置为 True，关于嵌入矩阵的梯度将是稀疏的。（默认值：False）

reset_parameters()[source]

重置模块的所有可学习参数。

forward(head_index: Tensor, rel_type: Tensor, tail_index: Tensor) → Tensor[source]

返回给定三元组的分数。

Parameters:
Return type:

Tensor

loss(head_index: Tensor, rel_type: Tensor, tail_index: Tensor) → Tensor[source]

返回给定三元组的损失值。

Parameters:
Return type:

Tensor