文档

知识图谱（KGs）是一种用于存储不同实体（节点）及其关系（边）信息的数据结构。在各种机器学习任务中，使用知识图谱的一种常见方法是计算知识图谱嵌入。DGL-KE 是一个高性能、易用且可扩展的软件包，用于学习大规模知识图谱嵌入。该软件包基于 Deep Graph Library (DGL) 构建，开发者可以在 CPU 机器、GPU 机器以及集群上运行 DGL-KE，并支持一系列流行的模型，包括 TransE、TransR、RESCAL、DistMult、ComplEx 和 RotatE。

图: DGL-KE 总体架构

目前 DGL-KE 支持三项任务：

• 训练：使用 dglke_train（单机）或 dglke_dist_train（分布式环境）训练知识图谱嵌入。
• 评估：读取预训练的嵌入，并使用 dglke_eval 在测试集上通过链接预测任务对嵌入进行评估。
• 推理：读取预训练的嵌入，使用 dglke_predict 执行实体/关系链接预测推理任务，或使用 dglke_emb_sim 执行嵌入相似性推理任务。

注意

如果您只想使用 TransE、DistMult 或 RotatE 训练 KGE 模型，请访问 https://github.com/awslabs/graphstorm。

快速入门

要安装最新版本的 DGL-KE，请运行以下命令：

sudo pip3 install dgl
sudo pip3 install dglke

通过运行以下命令，在 FB15k 数据集上训练一个 transE 模型：

DGLBACKEND=pytorch dglke_train --model_name TransE_l2 --dataset FB15k --batch_size 1000 \
--neg_sample_size 200 --hidden_dim 400 --gamma 19.9 --lr 0.25 --max_step 500 --log_interval 100 \
--batch_size_eval 16 -adv --regularization_coef 1.00E-09 --test --num_thread 1 --num_proc 8

此命令将下载 FB15k 数据集，训练 transE 模型，并将训练好的嵌入保存到文件中。

性能与可扩展性

DGL-KE 专为大规模学习而设计。它引入了多种新颖的优化技术，能够加速包含数百万个节点和数十亿条边的知识图谱的训练。我们在由超过 86M 个节点和 338M 条边组成的知识图谱上的基准测试表明，DGL-KE 可以在配备 8 块 GPU 的 EC2 实例上用时 100 分钟完成嵌入计算，而在由 4 台机器（每台 48 核）组成的 EC2 集群上则只需 30 分钟。这些结果相比最佳竞争方案实现了 2×∼5× 的速度提升。

图: DGL-KE 与 GraphVite 在 FB15k 上的对比

图: DGL-KE 与 Pytorch-BigGraph 在 Freebase 上的对比

更多详细信息请参阅我们的 文档! 如果您对 DGL-KE 中的优化技术感兴趣，请查看 我们的论文 以获取更多信息。

引用

如果您在科学出版物中使用 DGL-KE，我们非常感谢您引用以下论文：

@inproceedings{DGL-KE,
author = {Zheng, Da and Song, Xiang and Ma, Chao and Tan, Zeyuan and Ye, Zihao and Dong, Jin and Xiong, Hao and Zhang, Zheng and Karypis, George},
title = {DGL-KE: Training Knowledge Graph Embeddings at Scale},
year = {2020},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
booktitle = {Proceedings of the 43rd International ACM SIGIR Conference on Research and Development in Information Retrieval},
pages = {739–748},
numpages = {10},
series = {SIGIR '20}
}

许可证

本项目采用 Apache-2.0 许可证授权。

DGL-KE 快速上手指南

DGL-KE 是一个高性能、易用且可扩展的知识图谱嵌入（KGE）学习包，基于深度图图书馆（DGL）构建。它支持在单机 CPU/GPU 或分布式集群上运行，内置 TransE、TransR、RotatE 等主流模型，适用于大规模知识图谱训练。

环境准备

• 操作系统：Linux (推荐 Ubuntu/CentOS)
• Python 版本：Python 3.6+
• 后端依赖：
  • PyTorch (需预先安装与系统 CUDA 版本匹配的 PyTorch)
  • DGL (Deep Graph Library)
• 硬件要求：
  • 单机模式：支持多核 CPU 或单/多 GPU 机器。
  • 分布式模式：需要多台机器组成的集群。

注意：如果您仅需使用 TransE、DistMult 或 RotatE 模型进行常规训练，也可以考虑使用功能更集成的 GraphStorm。

安装步骤

建议使用 pip3 进行安装。为确保国内下载速度，可指定清华或阿里镜像源。

1. 安装 DGL 核心库：

   sudo pip3 install dgl -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   (注：若需特定 CUDA 版本的 DGL，请参考 DGL 官方文档安装对应版本，例如 pip3 install dgl-cu11x)

2. 安装 DGL-KE：

   sudo pip3 install dglke -i https://pypi.tuna.tsinghua.edu.cn/simple
   

基本使用

以下示例演示如何在单机环境下，使用 TransE 模型在 FB15k 数据集上进行训练。该命令会自动下载数据集、启动训练并保存嵌入向量。

执行命令：

DGLBACKEND=pytorch dglke_train --model_name TransE_l2 --dataset FB15k --batch_size 1000 \
--neg_sample_size 200 --hidden_dim 400 --gamma 19.9 --lr 0.25 --max_step 500 --log_interval 100 \
--batch_size_eval 16 -adv --regularization_coef 1.00E-09 --test --num_thread 1 --num_proc 8

参数说明：

• --model_name: 指定模型类型（如 TransE_l2, RotatE, DistMult 等）。
• --dataset: 指定数据集名称（支持 FB15k, WN18, Freebase 等内置数据集）。
• --num_proc: 指定使用的 CPU 进程数，可根据机器核心数调整以提升速度。
• --test: 训练完成后自动在测试集上进行评估。

训练完成后，生成的嵌入文件将保存在当前目录下的输出文件夹中，可用于后续的链接预测 (dglke_predict) 或相似度计算 (dglke_emb_sim) 任务。