A Python Library for Inductive Knowledge Graph Representation Learning

English | 中文

NeuralKG-ind is a python-based library for inductive knowledge graph representation learning, which includes standardized processes, rich existing methods, decoupled modules, and comprehensive evaluation metrics. We provide comprehensive documents for beginners.

Table of Contents

• Table of Contents
• 😃What's New
  • Feb, 2023
• Overview
• Demo
• Implemented Methods
• Quick Start
  • Installation
  • Training
  • Evaluation
  • Hyperparameter Tuning
• Reproduced Results
• Notebook Guide
• Detailed Documentation
• NeuralKG-ind Core Team

😃What's New

Feb, 2023

• We have released a paper NeuralKG-ind: A Python Library for Inductive Knowledge Graph Representation Learning

Overview

NeuralKG-ind is built on PyTorch Lightning and based on NeuralKG. It provides a general workflow for developing models handling inductive tasks on KGs. It has the following features:

• Standardized process. According to existing methods, we standardized the overall process of constructing an inductive knowledge graph representation learning model, including data processing, sampler and trainer construction, and evaluation of link prediction and triple classification tasks. We also provide auxiliary functions, including log management and hyper-parameter tuning, for model training and analysis.

• Rich existing methods. We re-implemented 5 inductive knowledge graph representation learning methods proposed in recent 3 years, including GraIL, CoMPILE, SNRI, RMPI and MorsE, enabling users to apply these models off the shelf.

• Decoupled modules. We provide a lot of decoupled modules, such as the subgraph extraction function, the node labeling function, neighbor aggregation functions, compound graph neural network layers, and KGE score functions, enabling users to construct a new inductive knowledge graph representation learning model faster.

• Long-term supports. We provide long-term support on NeuralKG-ind, including maintaining detailed documentation, creating straightforward quick-start, adding new models, solving issues, and dealing with pull requests.

Demo

There is a demonstration of NeuralKG-ind.

Implemented Methods

Components	Models
KGEModel	TransE, TransH, TransR, ComplEx, DistMult, RotatE, ConvE, BoxE, CrossE, SimplE, HAKE, PairRE, DualE
GNNModel	RGCN, KBAT, CompGCN, XTransE, GraIL, CoMPILE, SNRI, RMPI, MorsE
RuleModel	ComplEx-NNE+AER, RUGE, IterE

Quick Start

Installation

Step1 Create a virtual environment using Anaconda and enter it

conda create -n neuralkg-ind python=3.8
conda activate neuralkg-ind

Step2 Install the appropriate PyTorch and DGL according to your cuda version

Here we give a sample installation based on cuda == 11.1

• Install PyTorch

pip install torch==1.9.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html

• Install DGL

pip install dgl-cu111 dglgo -f https://data.dgl.ai/wheels/repo.html

Step3 Install package

git clone https://github.com/zjukg/NeuralKG-ind.git
cd NeuralKG-ind
python setup.py install

Training

# Use bash script
sh ./scripts/your-sh

# Use config
python main.py --load_config --config_path <your-config>

Evaluation

# Testing AUC and AUC-PR 
python main.py --test_only --checkpoint_dir <your-model-path> --eval_task triple_classification 

# Testing MRR and hit@1,5,10
python main.py --test_only --checkpoint_dir <your-model-path> --eval_task link_prediction --test_db_path <your-db-path> 

Hyperparameter Tuning

NeuralKG-ind utilizes Weights&Biases supporting various forms of hyperparameter optimization such as grid search, Random search, and Bayesian optimization. The search type and search space are specified in the configuration file in the format "*.yaml" to perform hyperparameter optimization.

The following config file displays hyperparameter optimization of the Grail on the FB15K-237 dataset using bayes search:

command:
  - ${env}
  - ${interpreter}
  - ${program}
  - ${args}
program: main.py
method: bayes
metric:
  goal: maximize
  name: Eval|auc
parameters:
  dataset_name:
    value: FB15K237
  model_name:
    value: Grail
  loss_name:
    values: [Margin_Loss]
  train_sampler_class:
    values: [SubSampler]
  emb_dim:
    values: [32, 64]
  lr:
    values: [1e-2, 5e-3, 1e-3]
  train_bs:
    values: [64, 128]
  num_neg:
    values: [16, 32]

Reproduced Results

There are some reproduced model results on FB15K-237 dataset and partial results on NELL-995 using NeuralKG as below. See more results in here

Method	FB15K-237_v1	FB15K-237_v2
AUC	AUC-PR	MRR	Hits@1	Hit@5	Hit@10	AUC	AUC-PR	MRR	Hits@1	Hit@5	Hit@10
GraIL	0.802	0.821	0.452	0.359	0.561	0.624	0.873	0.900	0.642	0.539	0.767	0.831
CoMPILE	0.800	0.835	0.516	0.437	0.600	0.668	0.876	0.905	0.617	0.509	0.741	0.813
SNRI	0.792	0.883	0.495	0.390	0.600	0.720	0.884	0.906	0.646	0.536	0.781	0.857
RMPI	0.803	0.823	0.532	0.451	0.620	0.689	0.851	0.882	0.632	0.523	0.763	0.830
MorsE	0.844	0.847	0.591	0.470	0.723	0.833	0.963	0.960	0.754	0.643	0.897	0.950

Method	FB15K-237_v3	FB15K-237_v4
AUC	AUC-PR	MRR	Hits@1	Hit@5	Hit@10	AUC	AUC-PR	MRR	Hits@1	Hit@5	Hit@10
GraIL	0.871	0.899	0.637	0.530	0.765	0.828	0.911	0.921	0.639	0.521	0.797	0.880
CoMPILE	0.906	0.925	0.670	0.568	0.796	0.859	0.927	0.932	0.704	0.604	0.831	0.894
SNRI	0.870	0.884	0.642	0.525	0.775	0.871	0.899	0.916	0.681	0.573	0.821	0.894
RMPI	0.876	0.866	0.662	0.569	0.767	0.827	0.905	0.916	0.647	0.535	0.787	0.866
MorsE	0.959	0.952	0.745	0.637	0.878	0.954	0.963	0.952	0.742	0.662	0.888	0.958

Method	NELL-995_v1	NELL-995_v2
AUC	AUC-PR	MRR	Hits@1	Hit@5	Hit@10	AUC	AUC-PR	MRR	Hits@1	Hit@5	Hit@10
GraIL	0.814	0.750	0.467	0.395	0.515	0.575	0.929	0.947	0.735	0.624	0.884	0.933
SNRI	0.737	0.720	0.523	0.475	0.545	0.595	0.864	0.884	0.630	0.507	0.774	0.863

Notebook Guide

😃We use colab to provide some notebooks to help users use our library.

Detailed Documentation

https://zjukg.github.io/NeuralKG-ind/neuralkg_ind.model.html

NeuralKG-ind Core Team

Zhejiang University: Wen Zhang, Zhen Yao, Mingyang Chen, Zhiwei Huang, Huajun Chen