teslacool / unifiedmolpretrain

Unified 2D and 3D Pre-Training of Molecular Representations

Unified 2D and 3D Pre-Training of Molecular Representations

This repository contains the code for Unified 2D and 3D Pre-Training of Molecular Representations, which is introduced in KDD2022.

Dataset

We use the PCQM4Mv2 dataset for pre-training, which has 3.38M data.

Requirements and Installation

• PyTorch
• Torch-Geometric

You can build a Docker image with the Dockerfile. To install code and develop it locally

pip install -e . 

Pre-training

bash run_training.sh --num-layers 12 --batch-size 128 --enable-tb \
    --node-attn --use-bn --pred-pos-residual --mask-prob 0.25 \
    -c 0,1,2,3 --dist

Finetuning

bash run_finetune.sh --num-layers 12 --batch-size 128 \
        --dropout 0.3 --dataset ogbg-molpcba \
        --pooler-dropout 0.1 --epochs 50 --seed 42 \
        -m /yourpretrainedmodel \
        --lr 0.0005 --weight-decay 0.01 --grad-norm 1 --prefix molpcba

unifiedmolpretrain's Issues

Could you provide a saved pre-trained model?

Hi,

Could you provide your saved pre-trained model so that I can apply the model to other datasets quickly?

Thank you

Bad Input File

It shows no input files, but in the "qcqm4m-v2" dataset, it doesn't consist of the "qcqm4m-v2_xyz", could you tell me how to get these missing files?

Generate 3D conformation molecules

In the paper, you describe the process to take 2D information for 3D conformation generation, however, the codebase does not have any examples on how to generate this. Could you add this?

num_graphs doesn't exit

Hi!
thanks for sharing your work!
I noticed that "num_graphs"doesn't exit (batch.num_graphs)
there is no num_graphs in dataset.py
could you please provide complete code or explain why?

thanks:)

Error regarding loss.backward() when enabling ddp

We face the following problem when enabling the DDP. The environment is set up by following Dockerfile.

Thank you for your reply.