wayne9qiu / pykg2vec

Python library for knowledge graph embedding and representation learning.

Pykg2vec: Python Library for KGE Methods

Pykg2vec is a library for learning the representation of entities and relations in Knowledge Graphs built on top of PyTorch 1.5 (TF2 version is available in tf-master branch as well). We have attempted to bring state-of-the-art Knowledge Graph Embedding (KGE) algorithms and the necessary building blocks in the pipeline of knowledge graph embedding task into a single library. We hope Pykg2vec is both practical and educational for people who want to explore the related fields.

Features:

• Support state-of-the-art KGE model implementations and benchmark datasets. (also support custom datasets)
• Support automatic discovery for hyperparameters.
• Tools for inspecting the learned embeddings.
  • Support exporting the learned embeddings in TSV or Pandas-supported format.
  • Interactive result inspector.
  • TSNE-based, KPI summary visualization (mean rank, hit ratio) in various format. (csvs, figures, latex table)

We welcome any form of contribution! Please refer to CONTRIBUTING.md for more details.

To Get Started

Before using pykg2vec, we recommend users to have the following libraries installed:

• python >=3.6 (recommended)
• pytorch>= 1.5

Quick Guide for Anaconda users:

• Setup a Virtual Environment: we encourage you to use anaconda to work with pykg2vec:

(base) $ conda create --name pykg2vec python=3.6
(base) $ conda activate pykg2vec

• Setup Pytorch: we encourage to use pytorch with GPU support for good training performance. However, a CPU version also runs. The following sample commands are for setting up pytorch:

# if you have a GPU with CUDA 10.1 installed
(pykg2vec) $ conda install pytorch torchvision cudatoolkit=10.1 -c pytorch
# or cpu-only
(pykg2vec) $ conda install pytorch torchvision cpuonly -c pytorch

• Setup Pykg2vec:

(pykg2vec) $ git clone https://github.com/Sujit-O/pykg2vec.git
(pykg2vec) $ cd pykg2vec
(pykg2vec) $ python setup.py install

For beginners, these papers, A Review of Relational Machine Learning for Knowledge Graphs, Knowledge Graph Embedding: A Survey of Approaches and Applications, and An overview of embedding models of entities and relationships for knowledge base completion can be good starting points!

User Documentation

The documentation is here.

Usage Examples

With pykg2vec, you can

1. Run a single algorithm with various models and datasets (customized dataset also supported).
2. Tune a single algorithm.
3. Perform Inference Tasks (more advanced).

We pasted one programming example (train.py) as below,

from pykg2vec.data.kgcontroller import KnowledgeGraph
from pykg2vec.common import Importer, KGEArgParser
from pykg2vec.utils.trainer import Trainer

def main():
    # getting the customized configurations from the command-line arguments.
    args = KGEArgParser().get_args(sys.argv[1:])
    
    # Preparing data and cache the data for later usage
    knowledge_graph = KnowledgeGraph(dataset=args.dataset_name, custom_dataset_path=args.dataset_path)
    knowledge_graph.prepare_data()

    # Extracting the corresponding model config and definition from Importer().
    config_def, model_def = Importer().import_model_config(args.model_name.lower())
    config = config_def(args)
    model = model_def(**config.__dict__)

    # Create, Compile and Train the model. While training, several evaluation will be performed.
    trainer = Trainer(model, config)
    trainer.build_model()
    trainer.train_model()

if __name__ == "__main__":
    main()

With train.py you can try KGE methods using the following commands:

# check all tunnable parameters.
$ python train.py -h 

# Train TransE on FB15k benchmark dataset.
$ python train.py -mn TransE

# Train using different KGE methods.
$ python train.py -mn [TransE|TransD|TransH|TransG|TransM|TransR|Complex|ComplexN3|CP|RotatE|Analogy|
                       DistMult|KG2E|KG2E_EL|NTN|Rescal|SLM|SME|SME_BL|HoLE|ConvE|ConvKB|Proje_pointwise]

# For KGE using projection-based loss function, use more processes for batch generation.
$ python train.py -mn [ConvE|ConvKB|Proje_pointwise] -npg [the number of processes, 4 or 6]

# Train TransE model using different benchmark datasets.
$ python train.py -mn TransE -ds [fb15k|wn18|wn18_rr|yago3_10|fb15k_237|ks|nations|umls|dl50a|nell_955]

# Train TransE model using different hyperparameters.
$ python train.py -exp True -mn TransE -ds fb15k -hpf ./examples/custom_hp.yaml

For more other pykg2vec usage, please check the programming examples.

Citation

Please kindly consider citing our paper if you find pykg2vec useful for your research.

  @article{yu2019pykg2vec,
  title={Pykg2vec: A Python Library for Knowledge Graph Embedding},
  author={Yu, Shih Yuan and Rokka Chhetri, Sujit and Canedo, Arquimedes and Goyal, Palash and Faruque, Mohammad Abdullah Al},
  journal={arXiv preprint arXiv:1906.04239},
  year={2019}
  }