[ICML24] Improving Equivariant Graph Neural Networks on Large Geometric Graphs via Virtual Nodes Learning
[25/05/13] 🚧 This project is currently under reconstruction. Stay tuned for updates!
pip install -r requirements.txtBefore executing following shell, please make sure that you specify the right data_directory and log_directory based on your machine. After the Training and Evaluating Ends, a log file will be generated in the log_directory, containing args, and losses.
1. N-body System Dataset
Data Generation:
cd ./datasets/nbody/datagen
bash run.sh
cd ../../..Train and Evaluate Model:
bash run_nbody.sh2. Protein Molecular Dynamics Dataset
Train and Evaluate Model:
bash run_protein.shThe Dataset will automatically download in the directory you specified.
3. Water-3D Dataset
Follow the instruction introduced here to download the .tfrecord format Water-3D data and transform them to .h5. Place the .h5 format data in your data directory.
Train and Evaluate Model:
bash run_simulation.shpython equivariant.pyIt will random generate a graph G, rotation matrix R and translation vector t, and check FastEGNN(G @ R + t) equals to FastEGNN(G) @ R + t or not.
If you find our work helpful, please cite as:
@inproceedings{
zhang2024improving,
title={Improving Equivariant Graph Neural Networks on Large Geometric Graphs via Virtual Nodes Learning},
author={Yuelin Zhang and Jiacheng Cen and Jiaqi Han and Zhiqiang Zhang and JUN ZHOU and Wenbing Huang},
booktitle={Forty-first International Conference on Machine Learning},
year={2024},
url={https://openreview.net/forum?id=wWdkNkUY8k}
}