Enhancing geometric representations for molecules with equivariant vector-scalar interactive message passing

ViSNet (shorted for “Vector-Scalar interactive graph neural Network”) is an equivariant geometry-enhanced graph neural for molecules that significantly alleviate the dilemma between computational costs and sufficient utilization of geometric information.

• ViSNet is published on Nature Communications.

• ViSNet-Drug Team won the 1st place in the The First Global AI Drug Development Competition!

• ViSNet Team won the 2nd place in the OGB-LSC @ NeurIPS 2022 PCQM4Mv2 Track! Please check out the branch OGB-LSC@NIPS2022 and give it a star if you find it useful!

• Clone this repository

• Install the dependencies

conda create -y -n visnet python=3.9
conda activate visnet
conda install pytorch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0 cudatoolkit=11.3 -c pytorch
conda install pyg==2.1.0 -c pyg
pip install pytorch-lightning==1.8.0
pip install ase ase[test] ogb

To train ViSNet on MD17, just run:

CUDA_VISIBLE_DEVICES=0 python train.py --conf examples/ViSNet-MD17.yml --dataset-arg aspirin --dataset-root /path/to/data --log-dir /path/to/log

One can modify the dataset-arg to train another molecule like ethanol.

To train ViSNet on QM9, just run:

CUDA_VISIBLE_DEVICES=0 python train.py --conf examples/ViSNet-QM9.yml --dataset-arg energy_U0 --dataset-root /path/to/data --log-dir /path/to/log

One can modify the dataset-arg to train another property like energy_U.

We have also provided example training configuration files for other datasets within the examples, which can be used in a similar way.

Once ViSNet is trained, to use a pretrained checkpoint for inference, simply run:

CUDA_VISIBLE_DEVICES=0 python train.py --conf examples/ViSNet-MD17.yml --dataset-arg aspirin --dataset-root /path/to/data --log-dir /path/to/log --task inference --load-model /path/to/ckpt

CUDA_VISIBLE_DEVICES=0 python train.py --conf examples/ViSNet-QM9.yml --dataset-arg energy_U0 --dataset-root /path/to/data --log-dir /path/to/log --task inference --load-model /path/to/ckpt

Please contact Tong Wang (Project Lead) for technical support.

This project is licensed under the terms of the MIT license.