Repository files navigation

This is the official codebase of the paper

Neural Bellman-Ford Networks: A General Graph Neural Network Framework for Link Prediction

Zhaocheng Zhu,Zuobai Zhang,Louis-Pascal Xhonneux,Jian Tang

A PyG re-implementation of NBFNet can be foundhere.

NBFNet is a graph neural network framework inspired by traditional path-basedmethods. It enjoys the advantages of both traditional path-based methods and moderngraph neural networks, includinggeneralization in the inductive setting,interpretability,high model capacity andscalability. NBFNet can beapplied to solve link prediction on both homogeneous graphs and knowledge graphs.

This codebase is based on PyTorch andTorchDrug. It supports training and inferencewith multiple GPUs or multiple machines.

You may install the dependencies via either conda or pip. Generally, NBFNet workswith Python 3.7/3.8 and PyTorch version >= 1.8.0.

conda install torchdrug pytorch=1.8.2 cudatoolkit=11.1 -c milagraph -c pytorch-lts -c pyg -c conda-forgeconda install ogb easydict pyyaml -c conda-forge

pip install torch==1.8.2+cu111 -f https://download.pytorch.org/whl/lts/1.8/torch_lts.htmlpip install torchdrugpip install ogb easydict pyyaml

To reproduce the results of NBFNet, use the following command. Alternatively, youmay use--gpus null to run NBFNet on a CPU. All the datasets will be automaticallydownloaded 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 inconfig/*/*.yaml.

For experiments on inductive relation prediction, you need to additionally specifythe 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]

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

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

Due to the large size of ogbl-biokg, we only evaluate on a small portion of thevalidation set during training. The following line evaluates a model on the fullvalidation / test sets of ogbl-biokg. Please replace the checkpoint with your ownpath.

python script/run.py -c config/knowledge_graph/ogbl-biokg_test.yaml --checkpoint /path/to/nbfnet/experiment/model_epoch_10.pth

Here are the results of NBFNet on standard benchmark datasets. All the results areobtained with 4 V100 GPUs (32GB). Note results may be slightly different if themodel is trained with 1 GPU and/or a smaller batch size.

1. The code is stuck at the beginning of epoch 0.

   This is probably because the JIT cache is broken.Tryrm -r ~/.cache/torch_extensions/* and run the code again.

If you find this codebase useful in your research, please cite the following 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}}