This is the source code for paper Neural Link Prediction with Walk Pooling.

python>=3.3.7

torch>=1.9.0

torch-cluster>=1.5.9

torch-geometric>=2.0.0

torch-scatter>=2.0.8

torch-sparse>=0.6.11

tqdm

This code was tested on macOS and Linux.

python ./src/main.py --data-name USAir

--data-name: supported data:

1. Without node attributes: USAir NS Power Celegans Router PB Ecoli Yeast

2. With node attributes: cora citeseer pubmed

--use-splitted: when it is True, we use the splitted data from SEAL. When it is False, we will randomly split train, validation and test data.

--data-split-num: the index of splitted data when --use-splitted is True. From 1 to 10.

--test-ratio and --val-ratio: Test ratio and validation ratio of the data set when --use-splitted is False. Defaults are 0.1 and 0.05 respectively.

--observe-val-and-injection: whether to contain the validation set in the observed graph and apply injection trick.

--practical-neg-sample: whether only see the train positive edges when sampling negative.

--num-hops: number of hops in sampling subgraph. Default is 2.

--max-nodes-per-hop: When the graph is too large or too dense, we need max node per hop threshold to avoid OOM. Default is None.

--init-attribute: the initial attribute for graphs without node attributes. options: n2v, one_hot, spc, ones, zeros, None. Default is ones.

--init-representation: node feature representation . options: gic, vgae, argva, None. Default is None.

--drnl: whether to use drnl labeling. Default is False.

--seed: random seed. Default is 1.

--lr: learning rate. Default is 0.00005.

-heads: using multi-heads in the attention link weight encoder. Default is 2.

--hidden-channels: Default is 32.

--batch-size: Default is 32.

--epoch-num: Default is 50.

Reproduce Table 1, 2, 3, 4 in the paper.

./bash/run.sh

If you find our work useful in your research, please cite our paper:

@article{pan2021neural,
  title={Neural Link Prediction with Walk Pooling},
  author={Pan, Liming and Shi, Cheng and Dokmani{\'c}, Ivan},
  journal={arXiv preprint arXiv:2110.04375},
  year={2021}
}