Shortcuts

RGCN

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

R-GCN [SKB+18] 提出于 2017 年,是第一个图神经网络模型。

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

例子:

from pybind11_ke.data import GraphDataLoader
from pybind11_ke.module.model import RGCN
from pybind11_ke.module.loss import RGCNLoss
from pybind11_ke.module.strategy import RGCNSampling
from pybind11_ke.config import Trainer, GraphTester

dataloader = GraphDataLoader(
        in_path = "../../benchmarks/FB15K237/",
        batch_size = 60000,
        neg_ent = 10,
        test = True,
        test_batch_size = 100,
        num_workers = 16
)

# define the model
rgcn = RGCN(
        ent_tol = dataloader.get_ent_tol(),
        rel_tol = dataloader.get_rel_tol(),
        dim = 500,
        num_layers = 2
)

# define the loss function
model = RGCNSampling(
        model = rgcn,
        loss = RGCNLoss(model = rgcn, regularization = 1e-5)
)

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

# train the model
trainer = Trainer(model = model, data_loader = dataloader.train_dataloader(),
        epochs = 10000, lr = 0.0001, use_gpu = True, device = 'cuda:0',
        tester = tester, test = True, valid_interval = 500, log_interval = 500,
        save_interval = 500, save_path = '../../checkpoint/rgcn.pth'
)
trainer.run()
Loss_emb: torch.nn.Embedding

图神经网络层的输出

RGCN: torch.nn.ModuleList

R-GCN 的图神经网络层

__call__(*args: Any, **kwargs: Any) Any

Call self as a function.

__init__(ent_tol: int, rel_tol: int, dim: int, num_layers: int)[源代码]

创建 RGCN 对象。

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

  • rel_tol (int) – 关系的个数

  • dim (int) – 实体和关系嵌入向量的维度

  • num_layers (int) – 图神经网络的层数

static __new__(cls, *args: Any, **kwargs: Any) Any
__repr__() str

Return repr(self).

__weakref__

list of weak references to the object (if defined)

build_hidden_layer(idx: int) dgl.nn.pytorch.conv.RelGraphConv[源代码]

返回第 idx 的图神经网络层。

参数:

idx (int) – 数据。

返回:

图神经网络层

返回类型:

dgl.nn.pytorch.conv.RelGraphConv

build_model()[源代码]

构建模型

dim: int

实体和关系嵌入向量的维度

distmult_score_func(head_emb: torch.Tensor, relation_emb: torch.Tensor, tail_emb: torch.Tensor, mode: str) torch.Tensor[源代码]

计算 DistMult 的评分函数。

参数:
返回:

三元组的得分

返回类型:

torch.Tensor

ent_emb: torch.nn.Embedding

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

ent_tol: int

实体的种类

forward(graph: dgl.DGLGraph, ent: torch.Tensor, rel: torch.Tensor, norm: torch.Tensor, triples: torch.Tensor, mode: str = 'single') torch.Tensor[源代码]

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

参数:
返回:

三元组的得分

返回类型:

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) – 模型保存的路径

num_layers: int

图神经网络的层数

pi_const: torch.nn.parameter.Parameter

常数 pi

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

R-GCN 的推理方法。

参数:
返回:

三元组的得分

返回类型:

torch.Tensor

rel_emb: torch.nn.parameter.Parameter

根据关系个数,创建的关系嵌入

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(embedding: torch.Tensor, triples: torch.Tensor, mode: str = 'single') tuple[torch.Tensor, torch.Tensor, torch.Tensor][源代码]

获得三元组对应头实体、关系和尾实体的嵌入向量。

参数:
  • embedding (torch.Tensor) – 经过图神经网络更新的实体嵌入向量

  • triples (torch.Tensor) – 训练的三元组

  • mode (str) – 模式

返回:

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

返回类型:

torch.Tensor

zero_const: torch.nn.parameter.Parameter

常数 0

Docs

Access comprehensive developer documentation for Pybind11-OpenKE

View Docs