Equivariant Subgraph Aggregation Networks (ESAN)
This repository contains the official code of the paper Equivariant Subgraph Aggregation Networks (ICLR 2022)
Install
First create a conda environment
conda env create -f environment.yml
and activate it.
Prepare the data
Run
python data.py --dataset $DATASET
where $DATASET
is one of the following:
- MUTAG
- PTC
- PROTEINS
- NCI1
- NCI109
- IMDB-BINARY
- IMDB-MULTI
- ogbg-molhiv
- ogbg-moltox21
- ZINC
- CSL
- EXP
- CEXP
Run the models
To perform hyperparameter tuning, make use of wandb
:
-
In
configs/
folder, choose theyaml
file corresponding to the dataset and setting (deterministic vs sampling) of interest, say -
Run
wandb sweep configs/<config-name>
to obtain a sweep id
-
Run the hyperparameter tuning with
wandb agent <sweep-id>
You can run the above command multiple times on each machine you would like to contribute to the grid-search
-
Open your project in your wandb account on the browser to see the results:
-
For the TUDatasets, the CSL and the EXP/CEXP datasets, refer to
Metric/valid_mean
andMetric/valid_std
to obtain the results. -
For the ogbg datasets and the ZINC dataset, compute mean and std of
Metric/train_mean
,Metric/valid_mean
,Metric/test_mean
over the different seeds of the same configuration. Then, take the results corresponding to the configuration obtaining the best validation metric.
-
Credits
For attribution in academic contexts, please cite
@inproceedings{bevilacqua2022equivariant,
title={Equivariant Subgraph Aggregation Networks},
author={Beatrice Bevilacqua and Fabrizio Frasca and Derek Lim and Balasubramaniam Srinivasan and Chen Cai and Gopinath Balamurugan and Michael M. Bronstein and Haggai Maron},
booktitle={International Conference on Learning Representations},
year={2022},
}