KGTuner
KGTuner copied to clipboard
LARS-research
KGTuner: Efficient Hyper-parameter Search for Knowledge Graph Learning" (ACL 2022 long paper)
KGTuner
Official code for the paper "KGTuner: Efficient Hyper-parameter Search for Knowledge Graph Learning" (ACL 2022 long paper).
Overview
We propose an efficient two-stage search algorithm, KGTuner, which efficiently explores HP configurations on small subgraph at the first stage and transfers the top-performed configurations for fine-tuning on the large full graph at the second stage.
Instructions
0. Setup
pip install -r requirements.txt
1. Generate sampled KG
Supported datasets: DATASET = {wn18rr/FB15k_237/ogbl-biokg/ogbl-wikikg2}.
usage:
python3 generator_random_walks.py -dataset {DATASET} -sample_ratio {FLOAT}
example:
python3 generator_random_walks.py -dataset wn18rr -sample_ratio 0.2
The generated new dataset will be store in dataset/sampled_wn18rr_0.2_starts_10.
2. Run hyper-parameter searching
Supported models: MODEL = {ComplEx/DistMult/TransE/RotatE/TuckER/RESCAL/ConvE/AutoSF}.
usage:
python3 run.py -search -earlyStop -space {SPACE} -dataset {DATASET} -cpu 2 -valid_steps {e.g., 2500} -max_steps {e.g., 100000} -model {MODEL} -eval_test -test_batch_size {e.g., 16} -gpu 0
Two-stage searching
example (firstly search on sampled_wn18rr_0.2_starts_10 and then on wn18rr):
stage1:
python3 run.py -search -earlyStop -space reduced -dataset sampled_wn18rr_0.2_starts_10 -cpu 2 -valid_steps 2500 -max_steps 20000 -model ComplEx -eval_test -test_batch_size 16 -gpu 0
stage2:
python3 run.py -search -earlyStop -space full -dataset wn18rr -pretrain_dataset sampled_wn18rr_0.2_starts_10 -cpu 2 -valid_steps 2500 -max_steps 100000 -model ComplEx -eval_test -test_batch_size 16 -gpu 0
One-stage searching
Alternatively, you can proceed searching directly on the original KG.
example:
python3 run.py -search -earlyStop -space full -dataset wn18rr -cpu 2 -valid_steps 2500 -max_steps 100000 -model ComplEx -eval_test -test_batch_size 16 -gpu 0
3. Show searched results
You can use the showResults.py to check your local HPO trials.
usage:
python3 showResults.py -dataset {DATASET} -model {MODEL}
example:
python3 showResults.py -dataset wn18rr -model ComplEx
Results
WN18RR
| Model | Test MRR | Test Hit@1 | Test Hit@3 | Test Hit@10 |
|---|---|---|---|---|
| ComplEx | 0.484 | 0.440 | 0.506 | 0.562 |
| DistMult | 0.453 | 0.407 | 0.468 | 0.548 |
| RESCAL | 0.479 | 0.436 | 0.496 | 0.557 |
| ConvE | 0.437 | 0.399 | 0.449 | 0.515 |
| TransE | 0.233 | 0.032 | 0.399 | 0.542 |
| RotatE | 0.480 | 0.427 | 0.501 | 0.582 |
| TuckER | 0.480 | 0.437 | 0.500 | 0.557 |
FB15k-237
| Model | Test MRR | Test Hit@1 | Test Hit@3 | Test Hit@10 |
|---|---|---|---|---|
| ComplEx | 0.352 | 0.263 | 0.387 | 0.530 |
| DistMult | 0.345 | 0.254 | 0.377 | 0.527 |
| RESCAL | 0.357 | 0.268 | 0.390 | 0.535 |
| ConvE | 0.335 | 0.242 | 0.368 | 0.523 |
| TransE | 0.327 | 0.228 | 0.369 | 0.522 |
| RotatE | 0.338 | 0.243 | 0.373 | 0.527 |
| TuckER | 0.347 | 0.255 | 0.382 | 0.534 |
ogbl-biokg
| Model | Test MRR | Val MRR | #parameters |
|---|---|---|---|
| ComplEx | 0.8385±0.0009 | 0.8394±0.0007 | 187,648,000 |
| DistMult | 0.8241±0.0008 | 0.8245±0.0009 | 93,824,000 |
| RotatE | 0.8013±0.0015 | 0.8024±0.0012 | 187,597,000 |
| TransE | 0.7781±0.0009 | 0.7787±0.0008 | 187,648,000 |
| AutoSF | 0.8354±0.0013 | 0.8361±0.0012 | 187,648,000 |
ogbl-wikikg2
| Model | Test MRR | Val MRR | #parameters |
|---|---|---|---|
| ComplEx | 0.4942±0.0017 | 0.5099±0.0023 | 250,113,900 |
| DistMult | 0.4837±0.0078 | 0.5004±0.0075 | 250,113,900 |
| RotatE | 0.2948±0.0026 | 0.2650±0.0034 | 250,087,150 |
| TransE | 0.4739±0.0021 | 0.4932±0.0013 | 250,113,900 |
| AutoSF | 0.5222±0.0021 | 0.5397±0.0023 | 250,113,900 |
Reproduction
Detailed commands can be found in code/scripts/reproduce.sh.
Citation
If you find this repository useful in your research, please kindly cite our paper.
@inproceedings{zhang2022kgtuner,
title={Efficient Hyper-parameter Search for Knowledge Graph Embedding},
author={Zhang, Yongqi and Zhou, Zhanke and Yao, Quanming and Li, Yong},
booktitle={ACL (long paper)},
pages={2715--2735},
year={2022}
}
Metadata
Owner
LARS-research
Metadata
KGTuner: Efficient Hyper-parameter Search for Knowledge Graph Learning" (ACL 2022 long paper)