hell-to-heaven / mrotate

Knowledge Graph Embedding by Relational and Entity Rotation

MRotatE

Knowledge Graph Embedding by Relational and Entity Rotation

Implemented features

Models:

• MRotatE

Evaluation Metrics:

• MRR, MR, HITS@1, HITS@3, HITS@10 (filtered)

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 MRotatE 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 MRotatE \
 -n 256 -b 512 -d 1000 \
 -g 24.0 -a 1.0 -adv \
 -lr 0.0001 --max_steps 50000 \
 -save models/RotatE_FB15k_0 --test_batch_size 8 -te

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

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

bash run.sh train MRotatE FB15k 0 0 512 256 1000 24.0 1.0 0.0001 50000 8 -te

Speed

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

Dataset	FB15k	FB15k-237	wn18	wn18rr
MAX_STEPS	50000	30000	60000	30000
TIME	2 h	1 h	1.5h	2 h

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.