No description, website, or topics provided.
Branch: master
Clone or download
Latest commit 3b9e7e7 Jan 28, 2019
Permalink
Type Name Latest commit message Commit time
Failed to load latest commit information.
codes delete __pycache__ Jan 28, 2019
data KGE first commit Jan 28, 2019
.gitignore KGE first commit Jan 28, 2019
LICENSE KGE first commit Jan 28, 2019
README.md Update README.md Jan 28, 2019
best_config.sh KGE first commit Jan 28, 2019
requirements.txt KGE first commit Jan 28, 2019
run.sh KGE first commit Jan 28, 2019

README.md

RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space

Introduction

This is the PyTorch implementation of the RotatE model for knowledge graph embedding (KGE). We provide a toolkit that gives state-of-the-art performance of several popular KGE models. The toolkit is quite efficient, which is able to train a large KGE model within a few hours on a single GPU.

Implemented features

Models:

  • RotatE
  • pRotatE
  • TransE
  • ComplEx
  • DistMult

Evaluation Metrics:

  • MRR, MR, HITS@1, HITS@3, HITS@10 (filtered)
  • AUC-PR (for Countries data sets)

Loss Function:

  • Uniform Negative Sampling
  • Self-Adversarial Negative Sampling

Usage

Knowledge Graph Data:

  • entities.dict: a dictionary map entities to unique ids
  • relations.dict: a dictionary map relations to unique ids
  • train.txt: the KGE model is trained to fit this data set
  • valid.txt: create a blank file if no validation data is available
  • test.txt: the KGE model is evaluated on this data set

Train

For example, this command train a RotatE model on FB15k dataset with GPU 0.

CUDA_VISIBLE_DEVICES=0 python -u codes/run.py --do_train \
 --cuda \
 --do_valid \
 --do_test \
 --data_path data/FB15k \
 --model RotatE \
 -n 256 -b 1024 -d 1000 \
 -g 24.0 -a 1.0 -adv \
 -lr 0.0001 --max_steps 150000 \
 -save models/RotatE_FB15k_0 --test_batch_size 16 -de

Check argparse configuration at codes/run.py for more arguments and more details.

Test

CUDA_VISIBLE_DEVICES=$GPU_DEVICE python -u $CODE_PATH/run.py --do_test --cuda -init $SAVE

Reproducing the best results

To reprocude the results in the ICLR 2019 paper RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space, you can run the bash commands in best_config.sh to get the best performance of RotatE, TransE, and ComplEx on five widely used datasets (FB15k, FB15k-237, wn18, wn18rr, Countries).

The run.sh script provides an easy way to search hyper-parameters:

bash run.sh train RotatE FB15k 0 0 1024 256 1000 24.0 1.0 0.0001 200000 16 -de

Speed

The KGE models usually take about half an hour to run 10000 steps on a single GeForce GTX 1080 Ti GPU with default configuration. And these models need different max_steps to converge on different data sets:

Dataset FB15k FB15k-237 wn18 wn18rr Countries S*
MAX_STEPS 150000 100000 80000 80000 40000
TIME 9 h 6 h 4 h 4 h 2 h

Results of the RotatE model

Dataset FB15k FB15k-237 wn18 wn18rr
MRR .797 ± .001 .949 ± .000 .337 ± .001 .477 ± .001
MR 40 177 309 3340
HITS@1 .746 .241 .944 .428
HITS@3 .830 .375 .952 .492
HITS@10 .884 .533 .959 .571

Using the library

The python libarary is organized around 3 objects:

  • TrainDataset (dataloader.py): prepare data stream for training
  • TestDataSet (dataloader.py): prepare data stream for evluation
  • KGEModel (model.py): calculate triple score and provide train/test API

The run.py file contains the main function, which parses arguments, reads data, initilize the model and provides the training loop.

Add your own model to model.py like:

def TransE(self, head, relation, tail, mode):
    if mode == 'head-batch':
        score = head + (relation - tail)
    else:
        score = (head + relation) - tail

    score = self.gamma.item() - torch.norm(score, p=1, dim=2)
    return score

Citation

If you use the codes, please cite the following paper:

@inproceedings{
 sun2018rotate,
 title={RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space},
 author={Zhiqing Sun and Zhi-Hong Deng and Jian-Yun Nie and Jian Tang},
 booktitle={International Conference on Learning Representations},
 year={2019},
 url={https://openreview.net/forum?id=HkgEQnRqYQ},
}