RotatE¶

class unike.module.model.RotatE(*args: Any, **kwargs: Any)[源代码]¶

RotatE [SDNT19] 提出于 2019 年，将实体表示成复数向量，关系建模为复数向量空间的旋转。

评分函数为:

\[\gamma - \parallel \mathbf{h} \circ \mathbf{r} - \mathbf{t} \parallel_{L_2}\]

\(\circ\) 表示哈达玛积（Hadamard product），正三元组的评分函数的值越大越好，如果想获得更详细的信息请访问 RotatE。

例子:

from unike.data import KGEDataLoader, UniSampler, TradTestSampler
from unike.module.model import RotatE
from unike.module.loss import SigmoidLoss
from unike.module.strategy import NegativeSampling
from unike.config import Trainer, Tester

# dataloader for training
dataloader = KGEDataLoader(
        in_path = '../../benchmarks/WN18RR/', 
        batch_size = 2000,
        neg_ent = 64,
        test = True,
        test_batch_size = 10,
        num_workers = 16,
        train_sampler = UniSampler,
        test_sampler = TradTestSampler
)

# define the model
rotate = RotatE(
        ent_tol = dataloader.get_ent_tol(),
        rel_tol = dataloader.get_rel_tol(),
        dim = 1024,
        margin = 6.0,
        epsilon = 2.0,
)

# define the loss function
model = NegativeSampling(
        model = rotate, 
        loss = SigmoidLoss(adv_temperature = 2), 
        regul_rate = 0.0,
)

# test the model
tester = Tester(model = rotate, data_loader = dataloader, use_gpu = True, device = 'cuda:1')

# train the model
trainer = Trainer(model = model, data_loader = dataloader.train_dataloader(), epochs = 6000,
        lr = 2e-5, opt_method = 'adam', use_gpu = True, device = 'cuda:1',
        tester = tester, test = True, valid_interval = 100,
        log_interval = 100, save_interval = 100,
        save_path = '../../checkpoint/rotate.pth', use_wandb = False)
trainer.run()

__call__(*args: Any, **kwargs: Any) → Any¶: Call self as a function.

__init__(ent_tol: int, rel_tol: int, dim: int = 100, margin: float = 6.0, epsilon: float = 2.0)[源代码]¶

创建 RotatE 对象。

参数:

ent_tol (int) – 实体的个数
rel_tol (int) – 关系的个数
dim (int) – 实体和关系嵌入向量的维度
margin (float) – 原论文中损失函数的 gamma。
epsilon (float) – RotatE 原论文对应的源代码固定为 2.0。

static __new__(cls, *args: Any, **kwargs: Any) → Any¶

__repr__() → str¶: Return repr(self).

__weakref__¶: list of weak references to the object (if defined)

_calc(h: torch.Tensor, r: torch.Tensor, t: torch.Tensor) → torch.Tensor[源代码]¶

计算 RotatE 的评分函数。

利用 torch.chunk() 拆分实体嵌入向量获得复数的实部和虚部。原论文使用 L1-norm 作为距离函数，而这里使用的 L2-norm 作为距离函数。

参数:

h (torch.Tensor) – 头实体的向量。
r (torch.Tensor) – 关系的向量。
t (torch.Tensor) – 尾实体的向量。

返回:

三元组的得分

返回类型:

torch.Tensor

dim_e: int¶: RotatE 原论文的实现中将实体嵌入向量的维度指定为 dim 的 2 倍。因为实体嵌入向量需要划分为实部和虚部。

dim_r: int¶: 关系嵌入向量的维度，为 dim。

ent_embeddings: torch.nn.Embedding¶: 根据实体个数，创建的实体嵌入。

ent_tol: int¶: 实体的种类

epsilon: int¶: RotatE 原论文对应的源代码固定为 2.0。

forward(triples: torch.Tensor, negs: torch.Tensor | None = None, mode: str = 'single') → torch.Tensor[源代码]¶

定义每次调用时执行的计算。 torch.nn.Module 子类必须重写 torch.nn.Module.forward()。

参数:

triples (str) – 正确的三元组
negs (torch.Tensor) – 负三元组类别
mode – 模式

返回:

三元组的得分

返回类型:

torch.Tensor

get_parameters(mode: str = 'numpy', param_dict: dict[str, Any] | None = None) → dict[str, numpy.ndarray] | dict[str, list] | dict[str, torch.Tensor]¶

获得模型权重。

参数:

mode – 模型保存的格式，可以选择 numpy 、 list 和 Tensor 。
param_dict (dict[str, Any] | None) – 可以选择从哪里获得模型权重。

返回:

模型权重字典。

返回类型:

dict[str, numpy.ndarray] | dict[str, list] | dict[str, torch.Tensor]

load_checkpoint(path: str)¶

加载模型权重。

参数:: path (str) – 模型保存的路径

load_parameters(path: str)¶

加载模型权重。

参数:: path (str) – 模型保存的路径

margin: torch.nn.parameter.Parameter¶: 原论文中损失函数的 gamma。

pi_const: torch.nn.parameter.Parameter¶: 常数 pi

predict(data: dict[str, Union[torch.Tensor, str]], mode) → torch.Tensor[源代码]¶

RotatE 的推理方法。

参数:: data (dict[str, Union[torch.Tensor,str]]) – 数据。
返回:: 三元组的得分
返回类型:: torch.Tensor

regularization(data: dict[str, Union[torch.Tensor, str]]) → torch.Tensor[源代码]¶

L2 正则化函数（又称权重衰减），在损失函数中用到。

参数:: data (dict[str, Union[torch.Tensor, str]]) – 数据。
返回:: 模型参数的正则损失
返回类型:: torch.Tensor

rel_embeddings: torch.nn.Embedding¶: 根据关系个数，创建的关系嵌入。

rel_tol: int¶: 关系的种类

save_checkpoint(path: str)¶

保存模型权重。

参数:: path (str) – 模型保存的路径

save_parameters(path: str)¶

用 json 格式保存模型权重。

参数:: path (str) – 模型保存的路径

set_parameters(parameters: dict[str, Any])¶

加载模型权重。

参数:: parameters (dict[str, Any]) – 模型权重字典。

tri2emb(triples: torch.Tensor, negs: torch.Tensor | None = None, mode: str = 'single') → tuple[torch.Tensor, torch.Tensor, torch.Tensor]¶

返回三元组对应的嵌入向量。

参数:

triples (str) – 正确的三元组
negs (torch.Tensor) – 负三元组类别
mode – 模式

返回:

头实体、关系和尾实体的嵌入向量

返回类型:

tuple[torch.Tensor, torch.Tensor, torch.Tensor]

zero_const: torch.nn.parameter.Parameter¶: 常数 0

RotatE¶

Docs