Skip to content

Latest commit

 

History

History

pytorch_geometric

README.md

Training Neural Bellman-Ford Networks (NBFnet) for Inductive Knowledge Graph Link Prediction on IPUs

This directory contains a PyTorch Geometric implementation of NBFNet used for link prediction in homogeneous and heterogeneous graphs. The model has been optimised for Graphcore's IPU.

Run inductive training on the FB15k-237 knowledge graph on Paperspace.
Gradient

Framework Domain Model Datasets Tasks Training Inference Reference
PyTorch GNNs NBFNet FB15k-237 Link Prediction
✅
Min. 4 IPUs (POD4) required

❌

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

Instructions summary

  1. Install and enable the Poplar SDK (see Poplar SDK setup)

  2. Install the system and Python requirements (see Environment setup)

  3. Running the model (see Running the model)

  4. Running the Jupyter notebook (see Running the notebook)

Poplar SDK setup

To check if your Poplar SDK has already been enabled, run:

 echo $POPLAR_SDK_ENABLED

If no path is provided, then follow these steps:

If no path is provided, then follow these steps:

  1. Download and unpack the Poplar SDK
  2. Navigate inside the unpacked Poplar SDK:
cd poplar_sdk*
  1. Enable the Poplar SDK with:
. enable

More detailed instructions on setting up your Poplar environment are available in the Poplar quick start guide.

Environment setup

To prepare your environment, follow these steps:

  1. Create and activate a Python3 virtual environment with your chosen name:
python3 -m venv <venv name>
source <venv name>/bin/activate

  1. Navigate to the Poplar SDK root directory

  2. Install the PopTorch (PyTorch) wheel:

cd <poplar sdk root dir>
pip3 install poptorch*x86_64.whl
  1. Install the PopTorch Geometric wheel:
cd <poplar sdk root dir>
pip3 install poptorch_geometric*.whl

  1. Navigate to this example's root directory, the directory of this readme.

  2. Install the Python requirements:

pip3 install -r requirements.txt

More detailed instructions on setting up your PyTorch environment are available in the PyTorch quick start guide.

Running

To train an inductive model on the FB15k-237 dataset from your terminal run

python run_nfbnet.py -c configs/IndFB15k-237_v1.yaml

The dataset.version parameter in the config file determines the data split (allowed values: ["v1", "v2", "v3", "v4"]).

Model specific hyperparameters can be set in the .yaml config file. Additionally, the following flags can be passed to run_nfbnet.py to adjust its behaviour:

-c, --config: Path to a config file
--device: The device to execute on ("ipu" or "cpu")
--profile: Generate a popvision profile
--profile_dir: Directory to save PopVision profile. See the PopVision documentation for more details.

Running the notebook

The model can be trained in an interactive way in a Jupyter notebook using the NBFNet_training.ipynb notebook.

  1. In your virtual environment, install the Jupyter notebook server:
pip3 install jupyter
  1. Launch a Jupyter Server on a specific port:
jupyter-notebook --no-browser --port <port number>
  1. Connect via SSH to your remote machine, forwarding your chosen port:
ssh -NL <port number>:localhost:<port number> <your username>@<remote machine>
  1. Use a browser to navigate to your Jupyter server and open the NBFNet_training.ipynb notebook.

For more details about this process, or if you need troubleshooting, see our guide on using IPUs from Jupyter notebooks.

License

This application is licensed under the MIT license, see the LICENSE file at the top-level of this repository.

This directory includes derived work from https://github.com/KiddoZhu/NBFNet-PyG
Copyright (c) 2021 MilaGraph
Licensed under the MIT License