TorchDrug教程--知识图推理

TorchDrug教程–知识图推理

教程来源TorchDrug开源

目录

• TorchDrug安装
• 分子数据结构
• 属性预测
• 预训练的分子表示
• 分子生成
• 逆合成
• 知识图推理

在知识图中，一个重要的任务是知识图推理，它的目的是在已知知识图中存在(h,r,t)-链接的情况下预测缺失的(h,r,t)-链接。有两种众所周知的知识图推理方法。一种是知识图嵌入，另一种是神经归纳逻辑规划。

在本教程中，我们提供了两个示例来说明如何使用TorchDrug进行知识图推理。

知识图谱嵌入

对于知识图推理，第一种比较流行的方法是知识图嵌入法。其基本思想是基于现有的(h,r,t)-链接学习知识图中每个实体和关系的嵌入向量。然后，这些嵌入进一步用于预测缺失的环节。

接下来，我们将介绍如何使用知识图嵌入模型进行知识图推理。

准备分子

我们使用FB15k-237数据集进行说明。FB15k-237是基于Freebase构建的，数据集有14,541个实体和237个关系。对于数据集，有一个标准的raining/validation/test分割。我们可以使用以下代码加载数据集

import torch
from torchdrug import core, datasets, tasks, models

dataset = datasets.FB15k237("~/kg-datasets/")
train_set, valid_set, test_set = dataset.split()

定义我们的模型

一旦我们加载了数据集，我们就可以构建模型了。让我们以RotatE模型为例，我们可以使用下面的代码来构建模型。

model = models.RotatE(num_entity=dataset.num_entity,
                      num_relation=dataset.num_relation,
                      embedding_dim=2048, max_score=9)

在这里，embedding_dim指定实体和关系嵌入的维度。max_score指定了推断(h,r,t)三组的合理性的偏差。

您可以考虑使用较小的嵌入维度以获得更好的效率。

之后，我们需要进一步定义我们的任务。对于知识图嵌入任务，我们可以简单地使用以下代码。

task = tasks.KnowledgeGraphCompletion(model, num_negative=256,
                                      adversarial_temperature=1)

这里，num_negative是用于训练的反例数，adversarial_temperature是采样反例的温度。

Train and Test

之后，我们可以训练和测试我们的模型。对于模型训练，我们需要设置一个优化器，并使用以下代码将所有内容放在一个Engine实例中。

optimizer = torch.optim.Adam(task.parameters(), lr=2e-5)
solver = core.Engine(task, train_set, valid_set, test_set, optimizer, gpus=[0], batch_size=1024)
solver.train(num_epoch=200)

在这里，我们可以减少num_epoch以获得更好的效率。
然后，我们可以使用下面的代码在验证集上进一步评估模型。

solver.evaluate("valid")

神经归纳逻辑规划

另一种流行的方法是神经归纳逻辑编程。神经归纳逻辑编程的思想是从训练数据中学习逻辑规则。一旦学习了逻辑规则，就可以进一步使用它们来预测缺失的环节。

神经归纳逻辑编程的一种流行方法是NeuralLP。NeuralLP考虑到最大长度的所有链式规则(例如，nationality = born_in + city_of)。此外，还使用注意机制为每个逻辑规则分配标量权重。在训练过程中，对注意力模块进行训练，这样我们就可以学习到每个规则的合适权重。在测试期间，逻辑规则及其权重一起使用来预测缺失的链接。

接下来，我们将介绍如何部署一个用于知识图推理的NeuralLP模型。

准备数据集

我们从加载数据集开始。与知识图嵌入教程类似，使用FB15k-237数据集进行说明。我们可以通过运行以下命令来加载数据集

import torch
from torchdrug import core, datasets, tasks, models

dataset = datasets.FB15k237("~/kg-datasets/")
train_set, valid_set, test_set = dataset.split()

定义我们的模型

然后，我们现在可以用以下代码定义NeuralLP模型

model = models.NeuralLP(num_relation=dataset.num_relation,
                        hidden_dim=128,
                        num_step=3,
                        num_lstm_layer=2)

这里hidden dim是NeuralLP中使用的实体和关系嵌入的维数。num_step是链式规则的最大长度(即链式规则主体中关系的最大数量)，通常设置为3。num_lstm_layer是NeuralLP中使用的lstm层数。

一旦我们定义了模型，我们就可以开始定义任务了。由于训练NeuralLP与训练知识图嵌入思想相似，我们也使用了下面的知识图嵌入任务

task = tasks.KnowledgeGraphCompletion(model, fact_ratio=0.75,
                                      num_negative=256,
                                      sample_weight=False)

不同之处在于，我们需要指定fact_ratio，它告诉代码有多少事实被用于构建我们执行推理的背景知识图，这个超参数通常设置为0.75。

Train and Test

对于我们已经定义的模型和任务，我们不能执行模型训练和测试。模型训练类似于知识图嵌入模型，其中我们需要创建一个优化器，并通过运行以下代码将每个组件提供给Engine实例

optimizer = torch.optim.Adam(task.parameters(), lr=1.0e-3)
solver = core.Engine(task, train_set, valid_set, test_set, optimizer,
                     gpus=[0, 1, 2, 3], batch_size=64)
solver.train(num_epoch=10)

这里，gpu指定了我们想要在其上训练模型的gpu。我们可以使用上面的表格指定多个gpu。为了提高效率，我们可以降低num_epoch这个值。

在模型训练之后，我们可以进一步使用以下代码来评估验证集上的模型

solver.evaluate("valid")