This repository contains the code and resources of the following paper:
KPGT: Knowledge-Guided Pre-training of Graph Transformer for Molecular Property Prediction (accepted to KDD 2022). DOI:https://doi.org/10.1145/3534678.3539426.
KPGT is a novel self-supervised learning framework for the representation learning of molecular graphs, consisting of a novel graph transformer architecture, LiGhT, and a knowledge-guided pre-training strategy.
Setup the required environment using environment.yml with Anaconda. While in the project directory run:
conda env create
Activate the environment
conda activate KPGT
We upload the datasets and splits used in our computational tests to figshare.
To download the datasets: https://figshare.com/s/8bbb8cad9ac644bf9caa.
Then unzip the file and put it in the KPGT/ directory.
Download the pre-trained model at: https://figshare.com/s/d488f30c23946cf6898f.
Then unzip it and put it in the KPGT/models/ directory.
You also can follow the steps below to pre-train a new model.
Extract the molecular descriptors and fingerprints of the SMILES in the ChEMBL dataset:
python preprocess_pretrain_dataset.py --data_path ../datasets/chembl29/
Use DistributedDataParallel to pre-train KPGT:
CUDA_VISIBLE_DEVICES=0,1,2,3 python -u -m torch.distributed.run --nproc_per_node=4 --nnodes=1 --master_port 12312 train_kpgt.py --save_path ../models/pretrained/base/ --n_threads 8 --n_devices 4 --config base --n_steps 100000 --data_path ../datasets/chembl29/
You can configure the model by modifing the KPGT/src/model_config.py.
Construct molecular line graphs and extract the molecular descriptors and the fingerprints from SMILES in a downstream dataset (e.g., bace):
python preprocess_downstream_dataset.py --data_path ../datasets/ --dataset bace
Fine-tune pre-trained model on a specific downstream task:
For classification tasks:
python finetune.py --config base --model_path ../models/pretrained/base/base.pth --dataset bace --data_path ../datasets/ --dataset_type classification --metric rocauc --split scaffold-0 --weight_decay 0 --dropout 0 --lr 3e-5
For regression tasks:
python finetune.py --config base --model_path ../models/pretrained/base/base.pth --dataset freesolv --data_path ../datasets/ --dataset_type regression --metric rmse --split scaffold-0 --weight_decay 0 --dropout 0 --lr 3e-5
Weight decay, dropout and lr are tunable hyper-parameters.
Due to the non-deterministic function in PyTorch, it is hard to exactly reproduce the fine-tuning results. Therefore, we provide the fine-tuned model for eleven datasets under the transfer learning setting, to guarantee the reproducibility of the test results reported in our paper.
To download the fine-tuned models: https://figshare.com/s/fd55d94d6bb21b8d7c39
Then unzip it and put the files in the KPGT/models/downstream/ directory.
Then the results can be reproduced by:
python evaluation.py --config base --model_path ../models/downstream/bace/scaffold_0.pth --dataset bace --data_path ../datasets/ --dataset_type classification --metric rocauc --split scaffold-0
The dataset, split and model can be specified using parameters --dataset, --split and --model_path, respectively.
To generate latent features for molecules from arbitrary datasets using the pre-trained KPGT:
python extract_features.py --config base --model_path ../models/pretrained/base/base.pth --data_path ../datasets/ --dataset bace
Modify --dataset to specify the target dataset.
KPGT is licensed under the Apache License, Version 2.0: http://www.apache.org/licenses/LICENSE-2.0.