ogb>=1.3.0
torch>=1.7.0
pytorch-lightning>=1.2.0
torch-geometric==master (pip install git+https://github.com/rusty1s/pytorch_geometric.git)
jupyterlab (for post-processing)
The MAG240M-LSC dataset will be automatically downloaded to the path denoted in root.py.
Please change its content accordingly if you want to download the dataset to a custom hard-drive or folder.
For each experiment, the test submission will be automatically saved in ./results/ after training is done.
Due to the file size of the MAG240M-LSC node feature matrix, training requires at least 256GB RAM.
For training the 2-layer model on 4 gpus, run:
python rgnn.py --exp_name rgat_2layers --device=k --accelerator='ddp' --model=rgat --hidden_channels=2048 --precision=16 --scheduler=cosine --optimizer=radam --extra_mlp --train_set=train --author_labels
For evaluating the 2-layer model on the best validation checkpoint with a neighborhood of 5*(sizes) and save prediction logits, run:
python rgnn.py --exp_name rgat_2layers --device=4 --accelerator='ddp' --evaluate --eval_size=5 --eval_size_dynamic --save_eval_probsFor training the 3-layer model on 4 gpus, run:
python rgnn.py --exp_name rgat_3layers --device=4 --accelerator='ddp' --model=rgat --hidden_channels=1800 --precision=16 --scheduler=cosine --optimizer=radam --extra_mlp --train_set=train --author_labels --num_layers=3 --sizes='25-20-15' --batch_size=512For evaluating the 3-layer model on the best validation checkpoint with a neighborhood of 5*(sizes) and save prediction logits, run:
python rgnn.py --exp_name rgat_3layers --device=4 --accelerator='ddp' --evaluate --eval_size=5 --eval_size_dynamic --save_eval_probsTo train the models on the training + validation sets (only for the model used on the hidden test set) replace --train_set=train with --train_set=train_val
| Model | Valid Accuracy (%) | Test Accuracy (%)* | #Parameters | Hardware |
|---|---|---|---|---|
| R-GAT [2] | 70.48 | 69.49 | 12.3M | GeForce RTX 2080 Ti (11GB GPU) |
| Ours (2-layer) | 71.08 | - | 81.2M | NVIDIA V100 (32GB GPU) |
| Ours (3-layer) | 71.87 | - | 99.5M | NVIDIA RTX 6000 (48GB GPU) |
| Ours (Ensemble) | 72.72 | 73.53 | 180.7M | NVIDIA V100 (32GB GPU), NVIDIA RTX 6000 (48GB GPU) |
The ensemble model is generated by aggregating 10 inference runs of the 2-layer model, and 10 total inference runs (with dynamic evaluation with sizes 1-10). In order to obtain different inference results with multiple runs, a change must be done to the pytorch sparse sampling code. Specifically, a random seed must be added to the sample_adj_cpu function in sample_cpu.cpp. This is worked around by doing evaluations with different neighborhood sizes, but was used to generate the inference results of the 2-layer model.
* Test Accuracy is evaluated on the hidden test set.
This code is heavily based on [1].
[1] Hu et al.: Open Graph Benchmark: Datasets for Machine Learning on Graphs
[2] Schlichtkrull et al.: Modeling Relational Data with Graph Convolutional Networks