# coding:utf-8
#
# unike/module/model/SimplE.py
#
# git pull from OpenKE-PyTorch by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on May 7, 2023
# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 7, 2023
#
# 该头文件定义了 SimplE.
"""
SimplE - 简单的双线性模型，能够为头实体和尾实体学习不同的嵌入向量。
"""
import math
import torch
import typing
import numpy as np
import torch.nn as nn
from .Model import Model
from typing_extensions import override
[文档]class SimplE(Model):
"""
``SimplE`` :cite:`SimplE` 提出于 2018 年，简单的双线性模型，能够为头实体和尾实体学习不同的嵌入向量。
评分函数为:
.. math::
1/2(<\mathbf{h}_{i}, \mathbf{v}_r, \mathbf{t}_{j}> + <\mathbf{h}_{j}, \mathbf{v}_{r^{-1}}, \mathbf{t}_{i}>)
:math:`< \mathbf{a}, \mathbf{b}, \mathbf{c} >` 为逐元素多线性点积（element-wise multi-linear dot product）。
正三元组的评分函数的值越大越好，负三元组越小越好，如果想获得更详细的信息请访问 :ref:`SimplE <simple>`。
例子::
from unike.config import Trainer, Tester
from unike.module.model import SimplE
from unike.module.loss import SoftplusLoss
from unike.module.strategy import NegativeSampling
# define the model
simple = SimplE(
ent_tol = train_dataloader.get_ent_tol(),
rel_tol = train_dataloader.get_rel_tol(),
dim = config.dim
)
# define the loss function
model = NegativeSampling(
model = simple,
loss = SoftplusLoss(),
batch_size = train_dataloader.get_batch_size(),
regul_rate = config.regul_rate
)
# dataloader for test
test_dataloader = TestDataLoader(in_path = config.in_path)
# test the model
tester = Tester(model = simple, data_loader = test_dataloader, use_gpu = config.use_gpu, device = config.device)
# train the model
trainer = Trainer(model = model, data_loader = train_dataloader, epochs = config.epochs,
lr = config.lr, opt_method = config.opt_method, use_gpu = config.use_gpu, device = config.device,
tester = tester, test = config.test, valid_interval = config.valid_interval,
log_interval = config.log_interval, save_interval = config.save_interval,
save_path = config.save_path, wandb_logger = wandb_logger)
trainer.run()
"""
[文档] def __init__(
self,
ent_tol: int,
rel_tol: int,
dim: int = 100):
"""创建 SimplE 对象。
:param ent_tol: 实体的个数
:type ent_tol: int
:param rel_tol: 关系的个数
:type rel_tol: int
:param dim: 实体嵌入向量和关系嵌入向量的维度
:type dim: int
"""
super(SimplE, self).__init__(ent_tol, rel_tol)
#: 实体嵌入向量和关系嵌入向量的维度
self.dim: int = dim
#: 根据实体个数，创建的实体嵌入
self.ent_embeddings: torch.nn.Embedding = nn.Embedding(self.ent_tol, self.dim * 2)
#: 根据关系个数，创建的关系嵌入
self.rel_embeddings: torch.nn.Embedding = nn.Embedding(self.rel_tol, self.dim * 2)
nn.init.xavier_uniform_(self.ent_embeddings.weight.data)
nn.init.xavier_uniform_(self.rel_embeddings.weight.data)
[文档] @override
def forward(
self,
triples: torch.Tensor,
negs: torch.Tensor = None,
mode: str = 'single') -> torch.Tensor:
"""
定义每次调用时执行的计算。
:py:class:`torch.nn.Module` 子类必须重写 :py:meth:`torch.nn.Module.forward`。
:param triples: 正确的三元组
:type triples: torch.Tensor
:param negs: 负三元组类别
:type negs: torch.Tensor
:param mode: 模式
:type triples: str
:returns: 三元组的得分
:rtype: torch.Tensor
"""
head_emb, relation_emb, tail_emb = self.tri2emb(triples, negs, mode)
score = self._calc(head_emb, relation_emb, tail_emb)
return score
[文档] def _calc(
self,
h: torch.Tensor,
r: torch.Tensor,
t: torch.Tensor) -> torch.Tensor:
"""计算 SimplE 的评分函数。
:param h: 头实体的向量。
:type h: torch.Tensor
:param r: 关系的向量。
:type r: torch.Tensor
:param t: 尾实体的向量。
:type t: torch.Tensor
:returns: 三元组的得分
:rtype: torch.Tensor
"""
hh_embs, th_embs = torch.chunk(h, 2, dim=-1)
r_embs, r_inv_embs = torch.chunk(r, 2, dim=-1)
ht_embs, tt_embs = torch.chunk(t, 2, dim=-1)
scores1 = torch.sum(hh_embs * r_embs * tt_embs, -1)
scores2 = torch.sum(ht_embs * r_inv_embs * th_embs, -1)
# Without clipping, we run into NaN problems.
# 基于论文作者的实现。
return torch.clamp((scores1 + scores2) / 2, -20, 20)
[文档] @override
def predict(
self,
data: dict[str, typing.Union[torch.Tensor,str]],
mode) -> torch.Tensor:
"""SimplE 的推理方法。
:param data: 数据。
:type data: dict[str, typing.Union[torch.Tensor,str]]
:returns: 三元组的得分
:rtype: torch.Tensor
"""
triples = data["positive_sample"]
head_emb, relation_emb, tail_emb = self.tri2emb(triples, mode=mode)
score = self._calc(head_emb, relation_emb, tail_emb)
return score
[文档] def regularization(
self,
data: dict[str, typing.Union[torch.Tensor, str]]) -> torch.Tensor:
"""L2 正则化函数（又称权重衰减），在损失函数中用到。
:param data: 数据。
:type data: dict[str, typing.Union[torch.Tensor, str]]
:returns: 模型参数的正则损失
:rtype: torch.Tensor
"""
pos_sample = data["positive_sample"]
neg_sample = data["negative_sample"]
mode = data["mode"]
pos_head_emb, pos_relation_emb, pos_tail_emb = self.tri2emb(pos_sample)
if mode == "bern":
neg_head_emb, neg_relation_emb, neg_tail_emb = self.tri2emb(neg_sample)
else:
neg_head_emb, neg_relation_emb, neg_tail_emb = self.tri2emb(pos_sample, neg_sample, mode)
pos_regul = (torch.mean(pos_head_emb ** 2) +
torch.mean(pos_relation_emb ** 2) +
torch.mean(pos_tail_emb ** 2)) / 3
neg_regul = (torch.mean(neg_head_emb ** 2) +
torch.mean(neg_relation_emb ** 2) +
torch.mean(neg_tail_emb ** 2)) / 3
regul = (pos_regul + neg_regul) / 2
return regul
[文档]def get_simple_hpo_config() -> dict[str, dict[str, typing.Any]]:
"""返回 :py:class:`SimplE` 的默认超参数优化配置。
默认配置为::
parameters_dict = {
'model': {
'value': 'SimplE'
},
'dim': {
'values': [50, 100, 200]
}
}
:returns: :py:class:`SimplE` 的默认超参数优化配置
:rtype: dict[str, dict[str, typing.Any]]
"""
parameters_dict = {
'model': {
'value': 'SimplE'
},
'dim': {
'values': [50, 100, 200]
}
}
return parameters_dict