NeuralKG是一个支持多种知识图谱表示学习/知识图谱嵌入(Knowledge Graph Embedding, KGE)模型的Python工具包,其中实现了多种传统知识图谱嵌入、基于图神经网络的知识图谱嵌入以及基于规则的知识图谱嵌入方法。同时为初学者提供了详细的文档以及一个开放共享的知识图谱表示学习社区网站。
- 我们添加了SEGNN模型
- 我们添加了DualE模型
- 我们添加了PairRE模型
- 我们添加了HAKE模型
- 我们提供了一个Google Colab教程帮助用户使用我们的工具包
- 我们提供了一篇博客来介绍我们的工具包在自定义数据集上的使用方式
- 我们发布了关于该工具包的论文NeuralKG: An Open Source Library for Diverse Representation Learning of Knowledge Graphs
NeuralKG工具包整体基于PyTorch Lightning框架,并提供了一个用于多种知识图谱表示学习模型的通用工作流程且高度模块化。NeuralKG具有如下特性:
-
支持多种方法。 NeuralKG提供了对三类知识图谱嵌入方法的代码实现,包括传统知识图谱嵌入, 基于图神经网络的知识图谱嵌入, 以及基于规则的知识图谱嵌入。
-
方便快速的客制化。 NeuralKG对知识图谱表示模型进行细化的模块解耦以方便使用者快速定制自己的模型,其中包括知识图谱数据处理模块,负采样模块,超参数监控模块,训练模块以及模型验证模块。这些模块被广泛应用于不同的知识图谱嵌入模型中
-
长期技术支持。 NeuralKG的核心开发团队将提供长期的技术支持,同时我们也欢迎开发者们对本项目进行pull requests。
NeuralKG在自定义知识图谱demo_kg上运行的示例。
类别 | 模型 |
---|---|
传统知识图谱嵌入(KGEModel) | TransE, TransH, TransR, ComplEx, DistMult, RotatE, ConvE, BoxE, CrossE, SimplE, HAKE, PairRE, DualE |
基于图神经网络的知识图谱嵌入(GNNModel) | RGCN, KBAT, CompGCN, XTransE, SEGNN |
基于规则的知识图谱嵌入(RuleModel) | ComplEx-NNE+AER, RUGE, IterE |
Step1 使用 Anaconda
创建虚拟环境,并进入虚拟环境
conda create -n neuralkg python=3.8
conda activate neuralkg
Step2 下载适用您CUDA版本的的PyTorch的DGL,下面我们提供一个基于CUDA 11.1的下载样例
- 下载PyTorch
pip install torch==1.9.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html
- 下载DGL
pip install dgl-cu111 dglgo -f https://data.dgl.ai/wheels/repo.html
Step3 安装NeuralKG
- 基于Pypi
pip install neuralkg
- 或基于源码
git clone https://github.com/zjukg/NeuralKG.git
cd NeuralKG
python setup.py install
# Use bash script
sh ./scripts/your-sh
# Use config
python main.py --load_config --config_path <your-config>
python main.py --test_only --checkpoint_dir <your-model-path>
NeuralKG使用Weights&Biases进行超参数调节,支持多种超参优化例如网格搜索、随机搜索和贝叶斯优化。搜索类型和搜索空间可以通过配置(*.yaml)文件进行设置。
下面展示了在FB15k-237上训练TransE,并使用贝叶斯搜索(bayes search)进行超参数调节的配置文件:
command:
- ${env}
- ${interpreter}
- ${program}
- ${args}
program: main.py
method: bayes
metric:
goal: maximize
name: Eval|hits@10
parameters:
dataset_name:
value: FB15K237
model_name:
value: TransE
loss_name:
values: [Adv_Loss, Margin_Loss]
train_sampler_class:
values: [UniSampler, BernSampler]
emb_dim:
values: [400, 600]
lr:
values: [1e-4, 5e-5, 1e-6]
train_bs:
values: [1024, 512]
num_neg:
values: [128, 256]
下面展示了使用NeuralKG的不同模型在FB15k-237上的结果,更多结果请访问此处。
Method | MRR | Hit@1 | Hit@3 | Hit@10 |
---|---|---|---|---|
TransE | 0.32 | 0.23 | 0.36 | 0.51 |
TransR | 0.23 | 0.16 | 0.26 | 0.38 |
TransH | 0.31 | 0.2 | 0.34 | 0.50 |
DistMult | 0.30 | 0.22 | 0.33 | 0.48 |
ComplEx | 0.25 | 0.17 | 0.27 | 0.40 |
SimplE | 0.16 | 0.09 | 0.17 | 0.29 |
ConvE | 0.32 | 0.23 | 0.35 | 0.50 |
RotatE | 0.33 | 0.23 | 0.37 | 0.53 |
BoxE | 0.32 | 0.22 | 0.36 | 0.52 |
HAKE | 0.34 | 0.24 | 0.38 | 0.54 |
PairRE | 0.35 | 0.25 | 0.38 | 0.54 |
DualE | 0.33 | 0.24 | 0.36 | 0.52 |
XTransE | 0.29 | 0.19 | 0.31 | 0.45 |
RGCN | 0.25 | 0.16 | 0.27 | 0.43 |
KBAT* | 0.28 | 0.18 | 0.31 | 0.46 |
CompGCN | 0.34 | 0.25 | 0.38 | 0.52 |
SEGNN | 0.36 | 0.27 | 0.39 | 0.54 |
IterE | 0.26 | 0.19 | 0.29 | 0.41 |
*:在KBAT的原论文作者实现中存在标签泄漏的问题,所以正确的结果相对较低,具体可以查看deepakn97/relationPrediction#28
😃我们使用colab提供部分notebook供用户使用我们的工具包
https://zjukg.github.io/NeuralKG/neuralkg.html
如果您使用了NeuralKG,请引用我们的论文
@article{zhang2022neuralkg,
title={NeuralKG: An Open Source Library for Diverse Representation Learning of Knowledge Graphs},
author={Zhang, Wen and Chen, Xiangnan and Yao, Zhen and Chen, Mingyang and Zhu, Yushan and Yu, Hongtao and Huang, Yufeng and others},
journal={arXiv preprint arXiv:2202.12571},
year={2022},
}
浙江大学: 张文,陈湘楠,姚祯,陈名杨,朱渝珊,俞洪涛,黄雨峰,许泽众,徐雅静,叶鹏,张溢弛,张宁豫,郑国轴,陈华钧