Implementation of Hirerachical Relational Learning for Few Shot Knowledge Graph Completion
- python==3.6.7
- torch==1.7.0
- torchvision==0.8.0
# other python/pytorch version might also work
Download Nell dataset from github and Wiki dataset from github.
Download pre-trained embedding Nell and Wiki and put them into the emb folder.
To train HiRe on Nell-One under 1-shot setting:
python main.py --dataset NELL-One --few 1 --prefix example-train-nell --learning_rate 0.001 --checkpoint_epoch 1000 --eval_epoch 1000 --batch_size 1024 --device 0 --step train
To test HiRe on Wiki-One under 3-shot setting:
python main.py --dataset Wiki-One --data_path ./Wiki-Hire/ --few 3 --prefix example-test-wiki --learning_rate 0.001 --checkpoint_epoch 1000 --eval_epoch 1000 --batch_size 1024 --device 0 --step test
To test HiRe on Nell-One under 5-shot setting using checkpoints:
python main.py --dataset NELL-One --few 5 --prefix example-test-ckpt-nell --learning_rate 0.001 --checkpoint_epoch 1000 --eval_epoch 1000 --batch_size 1024 --device 0 --eval_ckpt ./best_ckpt/nell_5shot_best.ckpt --step test
This repo is based on MetaR and InfoNCE.
If you find this project useful in your research, please consider cite:
@inproceedings{wu2023hierarchical,
title={Hierarchical Relational Learning for Few-Shot Knowledge Graph Completion},
author={Han Wu and Jie Yin and Bala Rajaratnam and Jianyuan Guo},
booktitle={The Eleventh International Conference on Learning Representations },
year={2023},
}