使用 Deep Graph Library 训练知识图谱嵌入

Original URL: https://aws.amazon.com/cn/blogs/machine-learning/training-knowledge-graph-embeddings-at-scale-with-the-deep-graph-library/

今天，我们很高兴与大家分享知识图谱嵌入库（DGL-KE），这是一套以深度图谱库（Deep Graph Library，DGL）为基础构建而成的知识图谱（KG）嵌入库。深度图谱库（DGL）是一套易于使用、性能出色且可扩展的Python库，主要用于实现对图谱的深度学习。现在，您可以为包含数十亿个节点与边的大型知识图谱库创建嵌入，且执行速度是其他同类手段的2到5倍。

例如，DGL-KE在“药物再利用知识图谱（DRKG）”之上创建嵌入，旨在显示哪些已经批准上市的药物能够用于对抗COVID-19。这些嵌入可用于预测药物治疗疾病的可能性或药物与疾病相关蛋白质结合的可能性。

在本文中，我们将重点介绍如何使用Kensho派生维基媒体数据集（KDWD）创建知识图谱嵌入（KGE knowledge graph embeddings）。例如，在自然语言处理（NLP）与信息检索用例当中，大家可以解析新查询，并将其句法转换为三元组（subject, predicate, object 即主语、谓词、宾语）。在将新的三元组添加至KG之后，则可以对节点做分类并根据现有的KGE来对关系作出推理，从而进一步扩充KG中的节点和关系。您可以借此指导聊天机器人发现对话意图，并为客户提供正确的FAQ（常见问题的回答）或提示信息。

知识图谱（KG Knowledge Graph）

知识图谱的本质，是对事实的结构化表示，包括为特定领域或应用场景专门构建的实体、关系与语义描述。知识图谱也被称为异构图，其中存在多种实体类型与关系类型。KG中存储的信息通常使用三元组进行指定，其中包含三个元素，即头（head）、关系（relation）与尾（tail），表示为[h,r,t]。头与尾也被称为实体（entities）。三元组组合在一起也被称为陈述（statements）。（译者注：statements也叫语句，三元组也对应上面的SPO主谓宾三元组）

在KG的帮助下，您可以以直观且富有表现力的方式对信息进行建模，从而轻松实现数据集成。例如，您可以使用Amazon Neptune为客户构建360度全方位身份信息图，也可以借此实现金融服务业中常见的“了解您的客户”应用。在大量涉及稀疏数据的医疗保健与生命科学领域，KG则可使用分类法与词汇表对来自不同孤岛的数据进行集成与协调。在电子商务与通信行业中，KG通常用于问答、聊天机器人与推荐系统。关于在各种实际用例中使用Amazon Neptune的更多详细信息，请参阅Amazon Neptune项目主页。

知识图谱嵌入（KGE Knowledge graph embeddings）

知识图谱嵌入，是知识图谱中实体与关系的低维表示形式。这些嵌入概括了给定节点的语义与局部结构信息。

当前有很多流行的KGE模型，例如TransE, TransR, RESCAL, DistMult, ComplEx以及 RotatE等。

每个模型都有一个不同的评分函数，该函数通过两个关联实体之间的关系来衡量它们之间的距离。通常可以直观的描述为：通过关系相连的实体往往距离较近，而没有连接的实体在向量空间中通常相距很远。

DGL-KE目前支持的模型评分函数如下所示：

维基媒体数据集

在本用例中，我们将使用Kensho派生维基媒体数据集（KDWD）。大家可以在DGL-KE GitHub repo当中找到相应的notebook示例与代码。

作为维基百科（Wikipedia）与维基数据（Wikidata）的组合，这套维基媒体数据集由以下三个数据层组成：

• 基础层——包含英文维基百科语料库。
• 中间层——标识哪些文本块是超链接，并在语料库中做标注。
• 顶部层——将维基百科里的超链接与维基数据KG中的条目相连接。

下图所示，为这些数据层的基本情况。

KDWD当中包含：

• 2,315,761,359个令牌
• 121,835,453个页面链接
• 5,343,564个维基百科页面
• 51,450,317个维基数据条目
• 141,206,854 条维基数据陈述

以下代码，为entity.txt文件中的示例：

ID     Label         Description1      Universe      totality of space and all contents2      Earth         third planet from the Sun in the Solar System

3      life          matter capable of extracting energy from the environment for replication

在创建嵌入之前，我们需要对数据进行预处理。为了计算KGE，DGL-KE支持两种输入格式。

针对原始的用户定义知识图谱（raw user-defined knowledge graphs）格式，我们需要提供三元组；实体与关系可以是任意字符串，数据加载器则自动生成ID映射。

如下表格来自Train User-Defined Knowledge Graphs，给出了三元组的示例：

针对用户定义知识图谱（user-defined knowledge graphs）格式，大家还需要提供实体与关系之间的ID映射（三元组内应只包含这些ID）。ID从0开始，且以连续形式递增。下面的表格来自Train User-Defined Knowledge Graphs，给出了这种情况三元组的示例：

关于更多详细信息，请参阅DGL-KE命令行。

虽然KDWD数据集直接提供字典，但由于其索引并非以0开头，且各索引值之间并不连续，所以无法在本示例中直接使用。我们需要对数据进行预处理，并使用原始格式生成嵌入。在对数据进行合并与清洗之后，我们最终得到具有以下属性的KG：

• 39,569,815个实体
• 1,213个关系
• 约1.2亿条陈述

以下代码，为三元组示例：

Head                Relation                   TailEiksteinen           located in the administrative…     RogalandTrivellona marlowi     instance of              taxonActa Numerica        main subject               mathematical analysisGünther Neukirchner    given name               GüntherRuth Pointer         given name                 Ruth

DGL-KE拥有多种不同训练模式。您可以根据数据集及训练要求，从中选择CPU、GPU、CPU-GPU混合模式以及分布式训练。在本文的用例中，我们将使用混合模式生成嵌入。如果能够将所有数据包含在GPU内存中，则GPU的训练性能最好。由于我们是基于大型KG来训练，因此选择了混合模式以获得更大的CPU与GPU内存池，且同时可以从GPU强大的训练加速能力中获益。

我们使用dgl-ke命令行工具创建嵌入，具体参见以下代码（译者注：要使用dgl 0.4.3 版本，不能用0.5以上的版本）：

!DGLBACKEND=pytorch dglke_train \
--model_name TransE_l2 \
--batch_size 1000 \
--neg_sample_size 200 \
 --hidden_dim 400 \
--gamma 19.9 \
--lr 0.25 \
--max_step 24000 \
--log_interval 100 \
--batch_size_eval 16 \
-adv \
--regularization_coef 1.00E-09 \
--test \
--gpu 0 1 2 3 \
--mix_cpu_gpu \
--save_path ./wikimedia \
--data_path ./data/wikimedia/ \
--format raw_udd_hrt \
--data_files train.txt valid.txt test.txt \
--neg_sample_size_eval 10000

关于DGL-KE参数的更多详细信息，请参阅 DGL-KE官网。

我们在大约7分钟时间内，使用p3.8xl实例训练了包含约4000万个实体、1200个关系以及约1.2亿条陈述的KG。

输入以下代码来做模型验证：

!DGLBACKEND=pytorch dglke_eval --dataset wikimedia --model_name TransE_l2 \
--neg_sample_size 200 --hidden_dim 400 --gamma 19.9 \
--batch_size_eval 16 --gpu 0 1 2 3  --model_path ./wikimedia/TransE_l2_wikimedia_0/ \
--data_path ./data/wikimedia/ --format raw_udd_hrt --data_files train.txt valid.txt test.txt --neg_sample_size_eval 10000 --no_eval_filter

以下代码为输出结果：

-------------- Test result --------------
Test average MRR: 0.4159753346227368
Test average MR: 1001.1689418833716
Test average HITS@1: 0.3540242971873324
Test average HITS@3: 0.45541123141672746
Test average HITS@10: 0.5213350742247359
-----------------------------------------

DGL-KE允许您使用[h,r,t]的任意组合执行KG下游任务。

在下面的示例中，我们试图为两个人找到相似的实体（也就是图谱中的节点），先创建文件head.list，该文件包含的条目如下：

head.list
Jeff Bezos
Barack Obama

DGL-KE提供对于实体及其关系执行离线推理的功能。读取head.list文件中的条目查找相似的节点实体，请输入以下代码：

!DGLBACKEND=pytorch dglke_emb_sim \
--format 'l_*' --data_files /home/ec2-user/SageMaker/DGL-Neptune/notebooks/data/wikimedia/head.list \
--mfile ./data/wikimedia/entities.tsv \
--emb_file ./wikimedia/TransE_l2_wikimedia_1/wikimedia_TransE_l2_entity.npy \
--raw_data --gpu 0 \
--exec_mode 'batch_left' \
--sim_func 'cosine' --topK 5
The following code is the output:
resulst.tsv
head tail score

以下为输出结果：

resulst.tsv
head tail score
Jeff Bezos Jeff Bezos 1.0
Jeff Bezos Aga Khan IV 0.8602205514907837
Jeff Bezos Alisher Usmanov 0.8584005236625671
Jeff Bezos Klaus Tschira 0.8512368202209473
Jeff Bezos Bill Gates 0.8441287875175476
Barack Obama Barack Obama 1.0
Barack Obama Donald Trump 0.9529082179069519
Barack Obama George W. Bush 0.9426612854003906
Barack Obama Harry S. Truman 0.9414601922035217
Barack Obama Ronald Reagan 0.9393566250801086

有趣的是，所有与杰夫-贝索斯（Jeff Bezos）相似的节点都显示为科技大亨；而所有与巴拉克-奥巴马（Barack Obama）相关的节点都显示为美国前任与现任总统。

总结

目前，化学、生物学、金融服务以及社交网络等领域存在大量图谱，我们可以使用这些图谱以直观方式了解各实体及其关系之间的复杂概念。这些图谱既可以同构、也可以异构，能够涵盖多种类型的实体与关系。

知识图谱嵌入KGE则为大家提供一种强大的方法，可以对特定节点上的语义与局部结构信息进行编码，您也可以将它们作为机器学习与深度学习模型的输入。DGL-KE支持多种流行嵌入模型，并能够让您通过CPU或GPU的大规模训练得到嵌入结果，训练速度是其他同类技术的2到5倍。

我们期待有更多朋友把KGE应用到现有KG和新的机器学习问题上。关于该库的更多详细信息，请参阅DGL-KE GitHub repo。关于在KG中使用维基媒体（Wikimedia KG嵌入的具体操作说明，请参阅DGL-KE notebook示例。

本篇作者