Awesome
Mole-BERT: Rethinking Pre-training Graph Neural Networks for Molecules (ICLR 2023)
This is a Pytorch implementation of the Mole-BERT paper:
Installation
We used the following Python packages for core development. We tested on Python 3.7
.
pytorch 1.0.1
torch-cluster 1.2.4
torch-geometric 1.0.3
torch-scatter 1.1.2
torch-sparse 0.2.4
torch-spline-conv 1.0.6
rdkit 2019.03.1.0
tqdm 4.31.1
tensorboardx 1.6
Dataset download
All the necessary data files can be downloaded from the following links.
For the chemistry dataset, download from chem data (2.5GB), unzip it, and put it under dataset/
.
Tokenizer Training
python vqvae.py --output_model_file OUTPUT_MODEL_PATH
This will save the resulting tokenizer to OUTPUT_MODEL_PATH
.
Pre-training and fine-tuning
1. Self-supervised pre-training
python pretrain.py --output_model_file OUTPUT_MODEL_PATH
This will save the resulting pre-trained model to OUTPUT_MODEL_PATH
.
2. Fine-tuning
python finetune.py --model_file INPUT_MODEL_PATH --dataset DOWNSTREAM_DATASET --filename OUTPUT_FILE_PATH
This will finetune pre-trained model specified in INPUT_MODEL_PATH
using dataset DOWNSTREAM_DATASET.
The result of fine-tuning will be saved to OUTPUT_FILE_PATH.
Reproducing results in the paper
Our results in the paper can be reproduced using a random seed ranging from 0 to 9 with scaffold splitting.
Useful resources for Chemical Pre-trained Models
- The first comprehensive survey on this topic: A Systematic Survey of Chemical Pre-trained Models (IJCAI 2023)
- A curated list of must-read papers, open-source pretrained models and pretraining datasets.
Acknowledgement
[1] Strategies for Pre-training Graph Neural Networks (Hu et al., ICLR 2020)
Citation
@inproceedings{
xia2023molebert,
title={Mole-{BERT}: Rethinking Pre-training Graph Neural Networks for Molecules},
author={Jun Xia and Chengshuai Zhao and Bozhen Hu and Zhangyang Gao and Cheng Tan and Yue Liu and Siyuan Li and Stan Z. Li},
booktitle={The Eleventh International Conference on Learning Representations },
year={2023},
url={https://openreview.net/forum?id=jevY-DtiZTR}
}