Shortcuts

TransR

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

TransR [LLS+15] 提出于 2015 年,是一个为实体和关系嵌入向量分别构建了独立的向量空间,将实体向量投影到特定的关系向量空间进行平移操作的模型。

评分函数为:

\[\Vert hM_r+r-tM_r \Vert_{L_1/L_2}\]

正三元组的评分函数的值越小越好,如果想获得更详细的信息请访问 TransR

例子:

from unike.data import KGEDataLoader, BernSampler, TradTestSampler
from unike.module.model import TransE, TransR
from unike.module.loss import MarginLoss
from unike.module.strategy import NegativeSampling
from unike.config import Trainer, Tester

# dataloader for training
dataloader = KGEDataLoader(
        in_path = "../../benchmarks/FB15K237/", 
        batch_size = 2048,
        neg_ent = 25, 
        test = True,
        test_batch_size = 10,
        num_workers = 16,
        train_sampler = BernSampler,
        test_sampler = TradTestSampler
)

# define the transe
transe = TransE(
        ent_tol = dataloader.get_ent_tol(),
        rel_tol = dataloader.get_rel_tol(),
        dim = 100, 
        p_norm = 1, 
        norm_flag = True)

transr = TransR(
        ent_tol = dataloader.get_ent_tol(),
        rel_tol = dataloader.get_rel_tol(),
        dim_e = 100,
        dim_r = 100,
        p_norm = 1, 
        norm_flag = True,
        rand_init = False)

model_e = NegativeSampling(
        model = transe, 
        loss = MarginLoss(margin = 5.0)
)

model_r = NegativeSampling(
        model = transr,
        loss = MarginLoss(margin = 4.0)
)

# pretrain transe
trainer = Trainer(model = model_e, data_loader = dataloader.train_dataloader(),
        epochs = 1, lr = 0.5, opt_method = "sgd", use_gpu = True, device = 'cuda:0')
trainer.run()
parameters = transe.get_parameters()
transe.save_parameters("../../checkpoint/transr_transe.json")

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

# train transr
transr.set_parameters(parameters)
trainer = Trainer(model = model_r, data_loader = dataloader.train_dataloader(),
        epochs = 1000, lr = 1.0, opt_method = "sgd", use_gpu = True, device = 'cuda:0',
        tester = tester, test = True, valid_interval = 100,
        log_interval = 100, save_interval = 100, save_path = '../../checkpoint/transr.pth')
trainer.run()

# test the model
transr.load_checkpoint('../../checkpoint/transr.pth')
tester.set_sampling_mode("link_test")
tester.run_link_prediction()
__call__(*args: Any, **kwargs: Any) Any

Call self as a function.

__init__(ent_tol: int, rel_tol: int, dim_e: int = 100, dim_r: int = 100, p_norm: int = 1, norm_flag: bool = True, rand_init: bool = False, margin: float | None = None)[源代码]

创建 TransR 对象。

参数:
  • ent_tol (int) – 实体的个数

  • rel_tol (int) – 关系的个数

  • dim_e (int) – 实体嵌入向量的维度

  • dim_r (int) – 关系嵌入向量的维度

  • p_norm (int) – 评分函数的距离函数, 按照原论文,这里可以取 1 或 2。

  • norm_flag (bool) – 是否利用 torch.nn.functional.normalize() 对实体和关系嵌入的最后一维执行 L2-norm。

  • rand_init (bool) – 关系矩阵是否采用随机初始化。

  • margin (float) – 当使用 RotatE [SDNT19] 的损失函数 unike.module.loss.SigmoidLoss,需要提供此参数,将 TransE [BUGD+13] 的正三元组的评分由越小越好转化为越大越好,如果想获得更详细的信息请访问 RotatE

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[源代码]

计算 TransR 的评分函数。

参数:
  • h (torch.Tensor) – 头实体的向量。

  • r (torch.Tensor) – 关系的向量。

  • t (torch.Tensor) – 尾实体的向量。

返回:

三元组的得分

返回类型:

torch.Tensor

_transfer(e: torch.Tensor, r_transfer: torch.Tensor) torch.Tensor[源代码]

将头实体或尾实体的向量投影到特定的关系向量空间。

参数:
  • e (torch.Tensor) – 头实体或尾实体向量。

  • r_transfer (torch.Tensor) – 特定关系矩阵

返回:

投影后的实体向量

返回类型:

torch.Tensor

dim_e: int

实体嵌入向量的维度

dim_r: int

关系嵌入向量的维度

ent_embeddings: torch.nn.Embedding

根据实体个数,创建的实体嵌入

ent_tol: int

实体的种类

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 – 模型保存的格式,可以选择 numpylistTensor

  • 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

当使用 RotatE [SDNT19] 的损失函数 unike.module.loss.SigmoidLoss,需要提供此参数,将 TransE [BUGD+13] 的正三元组的评分由越小越好转化为越大越好,如果想获得更详细的信息请访问 RotatE

norm_flag: bool

是否利用 torch.nn.functional.normalize() 对实体和关系嵌入向量的最后一维执行 L2-norm。

p_norm: int

评分函数的距离函数, 按照原论文,这里可以取 1 或 2。

pi_const: torch.nn.parameter.Parameter

常数 pi

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

TransR 的推理方法。

参数:
  • data (dict[str, Union[torch.Tensor,str]]) – 数据。

  • mode (str) – ‘head_predict’ 或 ‘tail_predict’

返回:

三元组的得分

返回类型:

torch.Tensor

rand_init: bool

关系矩阵是否采用随机初始化

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]) – 模型权重字典。

transfer_matrix: torch.nn.Embedding

关系矩阵

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

Docs

Access comprehensive developer documentation for UniKE

View Docs