Awesome

T-GAP

This is an official implementation of T-GAP: Learning to Walk across Time for Temporal Knowledge Graph Completion (submitted to KDD2021).

Requirements

python >= 3.7
torch >= 1.4.0
numpy == 1.17.2
dgl-cu101 == 0.4.2
tqdm == 4.36.1
ipdb == 0.12.3

You can install the requirements using pip install -r requirements.txt.

Datasets, Pre-processing, and Pretrained Model

We provide a shell script to download the datasets and pre-process them. The datasets used in the paper are ICEWS14, ICEWS05-15, and Wikidata11k. Go to the data folder, and type sh preprocess.sh. The script downloads, and pre-processes the official distribution of the 3 public datasets provided by Garcia-Duran et al.

We also provide the version of T-GAP trained on ICEWS14 dataset. The checkpoint file can be found at results/checkpoint/trained_icews14.tar.

Simple Demo

You can run a demo script to evaluate the trained model on ICEWS14, by running sh scripts/demo.sh in the project root.
The experimental results of the trained model are as follows:

MRR	Hits@1	Hits@3	Hits@10
0.610	50.9	67.7	79.0

Note that we have used 1 Tesla V100 GPU for all experiments.

Training

We provide additional scripts for each dataset that trains T-GAP with the same hyperparameter configurations as used in the paper. The scripts can be run as follows:

ICEWS14: sh scripts/icews14.sh
ICEWS05-15: sh scripts/icews05-15.sh
Wikidata11k: sh sh scripts/wikidata11k.sh

Running these files will automatically create checkpoint files and tensorboard log in results/checkpoint/${DATASET}, and results/log/${DATASET} respectively. For a tutorial on how to use tensorboard, refer to Pytorch Tensorboard Tutorial.

To train with custom hyperparameters, run main.py in the project root folder as follows:

python main.py \
  --run=${DIR_NAME} \
  --dataset=${DATASET_NAME} \
  --device=${cuda/cpu} \
  --seed=... \
  --epoch=... \
  --lr=... \
  --grad_clip=... \
  --node_dim=... \
  --num_in_heads=... \
  --num_out_heads=... \
  --num_step=... \
  --num_sample_from=... \
  --max_num_neighbor=... \

Option	Description	Default
`run`	Output dir name (automatically created under `results/checkpoint`, `results/log`)	Current Time
`dataset`	Dataset to train (among `data/icews14_aug`, `data/icews05-15_aug`, `data/wikidata11k_aug`)	`data/icews14_aug`
`device`	Torch device to use (`cuda` or `cpu`)	`cuda`
`seed`	Random seed	999
`epoch`	Number of training epochs	10
`lr`	Learning Rate	5e-4
`grad_clip`	Gradient clipping norm	3
`node_dim`	Embedding dimension	100
`num_in_heads`	Number of PGNN, SGNN attention heads	5
`num_out_heads`	Number of Attention Flow attention heads	5
`num_step`	Number of propagation steps	3
`num_sample_from`	Number of core nodes	10
`max_num_neighbor`	Number of sampled / added edges	100

Evaluation

You can evaluate a trained model, using Mean Reciprocal Rank (MRR), Hits@1/3/10 on the test set.
For evaluation, you need to locate the checkpoint file generated during training. To the same command used in the training, add --test and --ckpt=${CKPT_DIR} to evaluate.

python main.py
  --test \
  --ckpt=${CKPT_DIR} \
  --other_options (Same configuration with training)

References

Garcia-Duran. A, Dumancic. S, and Niepert. M, Learning Sequence Encoders for Temporal Knowledge Graph Completion. EMNLP 2018