MTL-KGC

This is PyTorch implementation of the Multi-Task Learning for Knowledge Graph Completion with Pre-trained Language Models.

Train

Train multitask learning with link prediction (LP), relation prediction (RP), and relevance ranking (RR). if you get AssertionError: Default process group is not initialized, please try with python -m torch.distributed.launch

python run_bert_multitask.py \
    --do_train \
    --task_list lp,rp,rr \
    --data_dir ./data/wn18rr \
    --bert_model bert-base-cased \
    --max_seq_length 128 \
    --output_dir ./output_dir \
    --num_train_epochs 5.0 \
    --learning_rate=2e-5 \
    --tb_log_dir=runs/log_dir

Train link prediction only

python run_bert_multitask.py \
    --do_train \
    --task_list lp \
    --data_dir ./data/wn18rr \
    --bert_model bert-base-cased \
    --max_seq_length 128 \
    --output_dir ./output_dir \
    --num_train_epochs 5.0 \
    --learning_rate=2e-5 \
    --tb_log_dir=runs/log_dir

Evaulation

python run_bert_multitask.py \
    --do_predict \
    --data_dir ./data/wn18rr \
    --bert_model {test_model} \
    --max_seq_length 128 \
    --output_dir ./output_dir \
    --task_list lp

Performance

WN18RR

FB15k-237