Skip to content

Commit 6aa3402

Browse files
KiddoZhuKiddoZhu
committed
init commit
0 parents  commit 6aa3402

24 files changed

+16708
-0
lines changed

‎.gitignore‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
# PyCharm
2+
/.idea
3+
4+
# VS Code
5+
/.vscode
6+
7+
# Python
8+
__pycache__
9+
*.pyc
10+
11+
# macOS
12+
.DS_Store

‎LICENSE‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
MIT License
2+
3+
Copyright (c) 2021 MilaGraph
4+
5+
Permission is hereby granted, free of charge, to any person obtaining a copy
6+
of this software and associated documentation files (the "Software"), to deal
7+
in the Software without restriction, including without limitation the rights
8+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9+
copies of the Software, and to permit persons to whom the Software is
10+
furnished to do so, subject to the following conditions:
11+
12+
The above copyright notice and this permission notice shall be included in all
13+
copies or substantial portions of the Software.
14+
15+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21+
SOFTWARE.

‎README.md‎

Lines changed: 214 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,214 @@
1+
# NBFNet: Neural Bellman-Ford Networks #
2+
3+
This is the official codebase of the paper
4+
5+
[Neural Bellman-Ford Networks: A General Graph Neural Network Framework for Link Prediction][paper]
6+
7+
[Zhaocheng Zhu](https://kiddozhu.github.io),
8+
[Zuobai Zhang](https://oxer11.github.io),
9+
[Louis-Pascal Xhonneux](https://github.com/lpxhonneux),
10+
[Jian Tang](https://jian-tang.com)
11+
12+
[paper]: https://arxiv.org/pdf/2106.06935.pdf
13+
14+
NeurIPS 2021
15+
16+
## Overview ##
17+
18+
NBFNet is a graph neural network framework inspired by traditional path-based
19+
methods. It enjoys the advantages of both traditional path-based methods and modern
20+
graph neural networks, including **generalization in the inductive setting**,
21+
**interpretability**, **high model capacity** and **scalability**. NBFNet can be
22+
applied to solve link prediction on both homogeneous graphs and knowledge graphs.
23+
24+
![NBFNet](asset/nbfnet.svg)
25+
26+
This codebase is based on PyTorch and [TorchDrug]. It supports training and inference
27+
with multiple GPUs or multiple machines.
28+
29+
[TorchDrug]: https://github.com/DeepGraphLearning/torchdrug
30+
31+
## Installation ##
32+
33+
You may install the dependencies via either conda or pip. Generally, NBFNet works
34+
with Python 3.7/3.8 and PyTorch version >= 1.8.0.
35+
36+
### From Conda ###
37+
38+
```bash
39+
conda install torchdrug pytorch=1.8.2 cudatoolkit=11.1 -c milagraph -c pytorch-lts -c pyg -c conda-forge
40+
conda install ogb easydict pyyaml -c conda-forge
41+
```
42+
43+
### From Pip ###
44+
45+
```bash
46+
pip install torch==1.8.2+cu111 -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html
47+
pip install torchdrug
48+
pip install ogb easydict pyyaml
49+
```
50+
51+
## Reproduction ##
52+
53+
To reproduce the results of NBFNet, use the following command. Alternatively, you
54+
may use `--gpus null` to run NBFNet on a CPU. All the datasets will be automatically
55+
downloaded in the code.
56+
57+
```bash
58+
python script/run.py -c config/inductive/wn18rr.yaml --gpus [0] --version v1
59+
```
60+
61+
We provide the hyperparameters for each experiment in configuration files.
62+
All the configuration files can be found in `config/*/*.yaml`.
63+
64+
For experiments on inductive relation prediction, you need to additionally specify
65+
the split version with `--version v1`.
66+
67+
To run NBFNet with multiple GPUs or multiple machines, use the following commands
68+
69+
```bash
70+
python -m torch.distributed.launch --nproc_per_node=4 script/run.py -c config/inductive/wn18rr.yaml --gpus [0,1,2,3]
71+
```
72+
73+
```bash
74+
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]
75+
```
76+
77+
### Visualize Interpretations on FB15k-237 ###
78+
79+
Once you have models trained on FB15k237, you can visualize the path interpretations
80+
with the following line. Please replace the checkpoint with your own path.
81+
82+
```bash
83+
python script/visualize.py -c config/knowledge_graph/fb15k237_visualize.yaml --checkpoint /path/to/nbfnet/experiment/model_epoch_20.pth
84+
```
85+
86+
### Evaluate ogbl-biokg ###
87+
88+
Due to the large size of ogbl-biokg, we only evaluate on a small portion of the
89+
validation set during training. The following line evaluates a model on the full
90+
validation / test sets of ogbl-biokg. Please replace the checkpoint with your own
91+
path.
92+
93+
```bash
94+
python script/run.py -c config/knowledge_graph/ogbl-biokg_test.yaml --checkpoint /path/to/nbfnet/experiment/model_epoch_10.pth
95+
```
96+
97+
## Results ##
98+
99+
Here are the results of NBFNet on standard benchmark datasets. All the results are
100+
obtained with 4 V100 GPUs (32GB). Note results may be slightly different if the
101+
model is trained with 1 GPU and/or a smaller batch size.
102+
103+
### Knowledge Graph Completion ###
104+
105+
<table>
106+
<tr>
107+
<th>Dataset</th>
108+
<th>MR</th>
109+
<th>MRR</th>
110+
<th>HITS@1</th>
111+
<th>HITS@3</th>
112+
<th>HITS@10</th>
113+
</tr>
114+
<tr>
115+
<th>FB15k-237</th>
116+
<td>114</td>
117+
<td>0.415</td>
118+
<td>0.321</td>
119+
<td>0.454</td>
120+
<td>0.599</td>
121+
</tr>
122+
<tr>
123+
<th>WN18RR</th>
124+
<td>636</td>
125+
<td>0.551</td>
126+
<td>0.497</td>
127+
<td>0.573</td>
128+
<td>0.666</td>
129+
</tr>
130+
<tr>
131+
<th>ogbl-biokg</th>
132+
<td>-</td>
133+
<td>0.829</td>
134+
<td>0.768</td>
135+
<td>0.870</td>
136+
<td>0.946</td>
137+
</tr>
138+
</table>
139+
140+
### Homogeneous Graph Link Prediction ###
141+
142+
<table>
143+
<tr>
144+
<th>Dataset</th>
145+
<th>AUROC</th>
146+
<th>AP</th>
147+
</tr>
148+
<tr>
149+
<th>Cora</th>
150+
<td>0.956</td>
151+
<td>0.962</td>
152+
</tr>
153+
<tr>
154+
<th>CiteSeer</th>
155+
<td>0.923</td>
156+
<td>0.936</td>
157+
</tr>
158+
<tr>
159+
<th>PubMed</th>
160+
<td>0.983</td>
161+
<td>0.982</td>
162+
</tr>
163+
</table>
164+
165+
### Inductive Relation Prediction ###
166+
167+
<table>
168+
<tr>
169+
<th rowspan="2">Dataset</th>
170+
<th colspan="4">HITS@10 (50 sample)</th>
171+
</tr>
172+
<tr>
173+
<th>v1</th>
174+
<th>v2</th>
175+
<th>v3</th>
176+
<th>v4</th>
177+
</tr>
178+
<tr>
179+
<th>FB15k-237</th>
180+
<td>0.834</td>
181+
<td>0.949</td>
182+
<td>0.951</td>
183+
<td>0.960</td>
184+
</tr>
185+
<tr>
186+
<th>WN18RR</th>
187+
<td>0.948</td>
188+
<td>0.905</td>
189+
<td>0.893</td>
190+
<td>0.890</td>
191+
</tr>
192+
</table>
193+
194+
Frequently Asked Questions
195+
--------------------------
196+
197+
1. **The code is stuck at the beginning of epoch 0.**
198+
199+
This is probably because the JIT cache is broken.
200+
Try `rm -r ~/.cache/torch_extensions/*` and run the code again.
201+
202+
Citation
203+
--------
204+
205+
If you find this codebase useful in your research, please cite the following paper.
206+
207+
```bibtex
208+
@article{zhu2021neural,
209+
title={Neural Bellman-Ford Networks: A General Graph Neural Network Framework for Link Prediction},
210+
author={Zhu, Zhaocheng and Zhang, Zuobai and Xhonneux, Louis-Pascal and Tang, Jian},
211+
journal={arXiv preprint arXiv:2106.06935},
212+
year={2021}
213+
}
214+
```

‎asset/nbfnet.svg‎

Lines changed: 1 addition & 0 deletions
Loading

‎config/homogeneous_graph/citeseer.yaml‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
output_dir: ~/experiments/
2+
3+
dataset:
4+
class: CiteSeerLinkPrediction
5+
path: ~/datasets/homogeneous_graphs/
6+
7+
task:
8+
class: LinkPrediction
9+
model:
10+
class: NBFNet
11+
input_dim: 32
12+
hidden_dims: [32, 32, 32, 32, 32, 32]
13+
message_func: distmult
14+
aggregate_func: pna
15+
short_cut: yes
16+
layer_norm: yes
17+
dependent: no
18+
remove_one_hop: yes
19+
symmetric: yes
20+
criterion: bce
21+
num_negative: 1
22+
strict_negative: yes
23+
24+
optimizer:
25+
class: Adam
26+
lr: 5.0e-3
27+
28+
engine:
29+
gpus: {{ gpus }}
30+
batch_size: 64
31+
32+
train:
33+
num_epoch: 20
34+
35+
metric: auroc

‎config/homogeneous_graph/cora.yaml‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
output_dir: ~/experiments/
2+
3+
dataset:
4+
class: CoraLinkPrediction
5+
path: ~/datasets/homogeneous_graphs/
6+
7+
task:
8+
class: LinkPrediction
9+
model:
10+
class: NBFNet
11+
input_dim: 32
12+
hidden_dims: [32, 32, 32, 32, 32, 32]
13+
message_func: distmult
14+
aggregate_func: pna
15+
short_cut: yes
16+
layer_norm: yes
17+
dependent: no
18+
remove_one_hop: yes
19+
symmetric: yes
20+
criterion: bce
21+
num_negative: 1
22+
strict_negative: yes
23+
24+
optimizer:
25+
class: Adam
26+
lr: 5.0e-3
27+
28+
engine:
29+
gpus: {{ gpus }}
30+
batch_size: 64
31+
32+
train:
33+
num_epoch: 20
34+
35+
metric: auroc

‎config/homogeneous_graph/pubmed.yaml‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
output_dir: ~/experiments/
2+
3+
dataset:
4+
class: PubMedLinkPrediction
5+
path: ~/datasets/homogeneous_graphs/
6+
7+
task:
8+
class: LinkPrediction
9+
model:
10+
class: NBFNet
11+
input_dim: 32
12+
hidden_dims: [32, 32, 32, 32, 32, 32]
13+
message_func: distmult
14+
aggregate_func: pna
15+
short_cut: yes
16+
layer_norm: yes
17+
dependent: no
18+
remove_one_hop: yes
19+
symmetric: yes
20+
criterion: bce
21+
num_negative: 1
22+
strict_negative: yes
23+
24+
optimizer:
25+
class: Adam
26+
lr: 5.0e-3
27+
28+
engine:
29+
gpus: {{ gpus }}
30+
batch_size: 16
31+
32+
train:
33+
num_epoch: 10
34+
35+
metric: auroc

0 commit comments

Comments
 (0)