An effective self-supervised framework for learning expressive molecular global representations to drug discovery
This repository is the official implementation of Our Paper Published on BIB.
📋 We provide the code of MPG implementation for pre-training, and fine-tuning on downstream tasks including molecular properties, DDI, and DTI predictions.
To install requirements:
pip install -r requirements.txt
To pre-train the model from scrach, please run the following command to preprocess the data by PSD and AttrMasking:
python pretraining/loader.py
Then, run:
bash pretraining/run_pretraining.sh
📋 Here, we provide an input example of unlabelled molecular data at
data/pretraining/raw/little_set_smi. To train the model on 11 M molecules thoroughly, please download the complete pre-training dataset first and put it in thedata/pretraining/raw/folder.
📋 To finetune the model, please download the pretrained models first and put the model in the
pretrained_model/folder.
| Methods | Classification (AUC-ROC) | Regression (RMSE) | |||||||
|---|---|---|---|---|---|---|---|---|---|
| Dataset | Tox21 | ToxCast | SIDER | ClinTox | BACE | BBBP | FreeSolv | ESOL | Lipo |
| # Molecules | 7831 | 8575 | 1427 | 1478 | 1513 | 2039 | 642 | 1128 | 4200 |
| ECFP | |||||||||
| TF_Robust | |||||||||
| GraphConv | |||||||||
| Weave | |||||||||
| SchNet | |||||||||
| MPNN | |||||||||
| DMPNN | |||||||||
| MGCN | |||||||||
| AttentiveFP | |||||||||
| TrimNet | |||||||||
| Mol2Vec | |||||||||
| N-GRAM | - | ||||||||
| SMILES-BERT | |||||||||
| HU. et.al. | - | - | - | ||||||
| GROVER | |||||||||
| MPG(ours) |
Classification tasks
python property/finetune.py --dataset toxcast --lr 0.0001 --input_model_file pretrained_model/MolGNet.pt
Note: If you want to run the fine-tune on CPU, please add an argument --cpu to the command above.
There are more hyper-parameters which can be tuned during finetuning. Please refer to the arguments in finetune.py.
Regression tasks
python property/finetune_regress.py --dataset esol --lr 0.0001 --input_model_file pretrained_model/MolGNet.pt
BIOSNAP
python DDI/finetune_snap.py --input_model_file pretrained_model/MolGNet.pt
TWOSIDES
python DDI/finetune_twosides.py --input_model_file pretrained_model/MolGNet.pt
python CPI/cross_validate.py --dataset human --input_model_file pretrained_model/MolGNet.pt
Due to the non-deterministic behavior of the function index_select_nd(See link) and the randomless of dataset split, it is hard to exactly reproduce the training process of finetuning. Therefore, we provide the finetuned model and the splitted datasets for thirteen datasets to guarantee the reproducibility of the experiments. Note: these results are fine-tuned in different hardware environments, resulting in slightly difference from reported statistics in the paper.
Molecular property prediction
DDI prediction
DTI prediction
We provide the eval.py function in property, DDI and DTI folders to reproduce the experiments.
For example, to evaluate the performance on BBBP dataset, suppose the finetuned model is placed in finetuned_model/, please run:
python property/eval.py --dataset bbbp --model_dir finetuned_model/property
Please cite our paper if you find this code useful for your research:
@article{li2021effective,
title={An effective self-supervised framework for learning expressive molecular global representations to drug discovery},
author={Li, Pengyong and
Wang, Jun and
Qiao, Yixuan and
Chen, Hao and
Yu, Yihuan and
Yao, Xiaojun and
Gao, Peng and
Xie, Guotong and
Song, Sen},
journal={Briefings in Bioinformatics},
volume={22},
number={6},
pages={bbab109},
year={2021},
publisher={Oxford University Press}
}