NBFNet: Neural Bellman-Ford Networks

PyG re-implementation of NBFNet. Authored by Zhaocheng Zhu and Michael Galkin.

Overview

NBFNet is a graph neural network framework inspired by traditional path-based methods. It enjoys the advantages of both traditional path-based methods and modern graph neural networks, including generalization in the inductive setting, interpretability, high model capacity and scalability. This repo implements NBFNet for transductive and inductive knowledge graph reasoning.

This codebase is based on PyTorch and PyTorch-Geometric. It supports training and inference with multiple GPUs or multiple machines.

Installation

You may install the dependencies via either conda or pip. Generally, NBFNet works with Python >= 3.7 and PyTorch >= 1.8.0.

From Conda

conda install pytorch=1.8.0 cudatoolkit=11.1 pyg -c pytorch -c pyg -c conda-forge
conda install ninja easydict pyyaml -c conda-forge

From Pip

pip install torch==1.8.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
pip install torch-scatter==2.0.8 torch-sparse==0.6.12 torch-geometric -f https://data.pyg.org/whl/torch-1.8.0+cu111.html
pip install ninja easydict pyyaml

Reproduction

To reproduce the results of NBFNet, use the following command. Alternatively, you may use --gpus null to run NBFNet on a CPU. All the datasets will be automatically downloaded in the code.

python script/run.py -c config/inductive/wn18rr.yaml --gpus [0] --version v1

We provide the hyperparameters for each experiment in configuration files. All the configuration files can be found in config/*/*.yaml.

For experiments on inductive relation prediction, you need to additionally specify the split version with --version v1.

To run NBFNet with multiple GPUs or multiple machines, use the following commands

python -m torch.distributed.launch --nproc_per_node=4 script/run.py -c config/inductive/wn18rr.yaml --gpus [0,1,2,3]

python -m torch.distributed.launch --nnodes=4 --nproc_per_node=4 script/run.py -c config/inductive/wn18rr.yaml --gpus [0,1,2,3,0,1,2,3,0,1,2,3,0,1,2,3]

Visualize Interpretations on FB15k-237

Once you have models trained on FB15k237, you can visualize the path interpretations with the following line. Please replace the checkpoint with your own path.

python script/visualize.py -c config/transductive/fb15k237_visualize.yaml --checkpoint /path/to/nbfnet/experiment/model_epoch_20.pth

Results

Here are the benchmark results of this re-implementation. All the results are obtained with 4 V100 GPUs (32GB). Note results may be slightly different if the model is trained with 1 GPU and/or a smaller batch size.

Knowledge Graph Completion

Dataset	MR	MRR	HITS@1	HITS@3	HITS@10
FB15k-237	113	0.416	0.322	0.456	0.602
WN18RR	650	0.551	0.496	0.573	0.661

Dataset	Training Time	Test Time	GPU Memory
FB15k-237	22.8 min / epoch	64 s	22.9 GiB
WN18RR	12.0 min / epoch	16 s	16.1 GiB

Inductive Relation Prediction

Dataset	HITS@10 (50 sample)
	v1	v2	v3	v4
FB15k-237	0.821	0.948	0.957	0.959
WN18RR	0.954	0.903	0.909	0.888

Dataset	v1			v2			v3			v4
	Training Time	Test Time	GPU Memory	Training Time	Test Time	GPU Memory	Training Time	Test Time	GPU Memory	Training Time	Test Time	GPU Memory
FB15k-237	2 s / epoch	< 1 s	3.51 GiB	6 s / epoch	< 1 s	5.02 GiB	15 s / epoch	< 1 s	6.56 GiB	29 s / epoch	< 1 s	8.10 GiB
WN18RR	3 s / epoch	< 1 s	5.17 GiB	22 s / epoch	< 1 s	11.3 GiB	3 s / epoch	1 s	18.7 GiB	7 s / epoch	1 s	6.84 GiB

Citation

If you find this codebase useful in your research, please cite the original paper.

@article{zhu2021neural,
  title={Neural bellman-ford networks: A general graph neural network framework for link prediction},
  author={Zhu, Zhaocheng and Zhang, Zuobai and Xhonneux, Louis-Pascal and Tang, Jian},
  journal={Advances in Neural Information Processing Systems},
  volume={34},
  year={2021}
}