train_autoencoder.py

from xbert import BertConfig, BertForMaskedLM
from transformers import BertModel
import torch
from torch import nn
import torch.distributed
from scheduler import create_scheduler
import copy
from pysmilesutils.augment import MolAugmenter
from torch.utils.data import DataLoader
from dataset import SMILESDataset_pretrain
import pytorch_lightning as pl
from pytorch_lightning.strategies import DDPStrategy
import torch.distributed
import argparse
from pathlib import Path
from utils import regexTokenizer


class AttrDict(dict):
    def __init__(self, *args, **kwargs):
        super(AttrDict, self).__init__(*args, **kwargs)
        self.__dict__ = self


class ldmol_autoencoder(pl.LightningModule):
    def __init__(self, cp=None, config=None, loader_len=0, no_train=False, tokenizer=None, use_linear=True, use_pr=False):
        super().__init__()
        self.save_hyperparameters()
        self.automatic_optimization = False
        self.config = config
        self.tokenizer = tokenizer
        self.training_step_outputs = []

        self.text_encoder = BertForMaskedLM(config=BertConfig.from_json_file(config['bert_config_decoder']))
        bert_config2 = BertConfig.from_json_file(config['bert_config_encoder'])
        self.text_encoder2 = BertModel(config=bert_config2)

        if cp:
            checkpoint = torch.load(cp, map_location='cpu')
            try:
                state_dict = copy.deepcopy(checkpoint['model'])
            except:
                state_dict = copy.deepcopy(checkpoint['state_dict'])
            for key in list(state_dict.keys()):
                if 'text_encoder.' in key:
                    new_key = key.replace('text_encoder.', '')
                    state_dict[new_key] = state_dict[key]
                del state_dict[key]
            msg = self.text_encoder2.load_state_dict(state_dict, strict=False)
            print('inside', msg)
            del state_dict
        for param in self.text_encoder2.parameters():
            param.requires_grad = False

        self.aug = MolAugmenter()
        self.use_linear = use_linear
        if use_linear:
            self.output_dim = 64
            final_dim = bert_config2.hidden_size
            self.encode_prefix = nn.Linear(final_dim, self.output_dim)
            self.decode_prefix = nn.Linear(self.output_dim, final_dim)

        if not no_train:
            self.loader_len = loader_len
            self.warmup_steps = config['schedular']['warmup_epochs']

    def forward(self, text_input_ids, text_attention_mask, text_input_ids2, text_attention_mask2):
        # ================= MLM ================= #
        input_ids = text_input_ids.clone()
        labels = input_ids.clone()[:, 1:]
        with torch.no_grad():
            text_embeds = self.text_encoder2(text_input_ids2, attention_mask=text_attention_mask2, return_dict=True).last_hidden_state
        if self.use_linear:
            text_embeds = self.decode_prefix(self.encode_prefix(text_embeds))
        mlm_output = self.text_encoder(input_ids,
                                       attention_mask=text_attention_mask,
                                       encoder_hidden_states=text_embeds,
                                       encoder_attention_mask=None,
                                       return_dict=True,
                                       is_decoder=True,
                                       return_logits=True,
                                       )[:, :-1, :]

        loss_fct = nn.CrossEntropyLoss(ignore_index=0, reduction='sum')
        loss_mlm = loss_fct(mlm_output.permute((0, 2, 1)), labels)
        return loss_mlm

    def configure_optimizers(self):
        arg_opt = self.config['optimizer']
        optimizer = torch.optim.AdamW(self.parameters(), lr=arg_opt['lr'], weight_decay=arg_opt['weight_decay'])
        arg_sche = AttrDict(self.config['schedular'])
        scheduler, _ = create_scheduler(arg_sche, optimizer)
        return [optimizer], [scheduler]

    def lr_scheduler_step(self, scheduler, optimizer_idx, metric):
        print('qqq', metric)

    def training_step(self, train_batch, batch_idx):
        optimizer = self.optimizers()
        scheduler = self.lr_schedulers()
        optimizer.zero_grad()
        text = train_batch
        # text_input = self.tokenizer(text, padding='longest', truncation=True, max_length=128, return_tensors="pt").to(self.device)
        text_input_ids = self.tokenizer(text, truncation='longest').to(self.device)
        text_attention_mask = torch.where(text_input_ids == 0, 0, 1).to(self.device)
        
        # text2 = ['[CLS]' + Chem.MolToSmiles(self.aug([Chem.MolFromSmiles(t[5:])])[0], canonical=False, isomericSmiles=True) if random.random()<0.1 else t for t in text]
        text2 = text
        # text_input2 = self.tokenizer(text2, padding='max_length', truncation=True, max_length=128, return_tensors="pt").to(self.device)
        text2_input_ids = self.tokenizer(text2, truncation='longest').to(self.device)
        text2_attention_mask = torch.where(text2_input_ids == 0, 0, 1).to(self.device)

        loss_mlm = self(text_input_ids, text_attention_mask, text2_input_ids, text2_attention_mask)
        loss = loss_mlm
        if loss != torch.tensor(0.):
            self.manual_backward(loss)
            torch.nn.utils.clip_grad_norm_(self.parameters(), 5.)
            optimizer.step()
        else:
            print('aaaaaaaaaaaa')
        if self.global_rank == 0:
            self.log('lr', optimizer.param_groups[0]["lr"], prog_bar=True)
            self.log('loss_mlm', loss_mlm, prog_bar=True)

        step_size = 100
        warmup_iterations = self.warmup_steps * step_size
        if self.current_epoch > 0 and batch_idx == 0:
            scheduler.step(self.current_epoch + self.warmup_steps)
        else:
            if self.current_epoch == 0 and batch_idx % step_size == 0 and batch_idx <= warmup_iterations:
                scheduler.step(batch_idx // step_size)
        self.training_step_outputs.append(torch.tensor([loss_mlm, ]))
        return torch.tensor([loss_mlm, ])

    def on_train_epoch_end(self):    # outputs: collection of returns from 'training_step'
        tmp = torch.stack(self.training_step_outputs[-1000:]).mean(dim=0).tolist()
        if self.global_rank == 0:
            print(f'\n mean loss: {tmp[0]:.4f}')
        self.training_step_outputs.clear()


def main(args, config):
    # data
    print("Creating dataset")
    dataset = SMILESDataset_pretrain(args.data_path, data_length=[0, 10000000], is_train=False, shuffle=True)
    print('#data:', len(dataset))
    data_loader = DataLoader(dataset, batch_size=config['batch_size'], num_workers=8, shuffle=False, pin_memory=True, drop_last=True)
    tokenizer = regexTokenizer(vocab_path=args.vocab_filename, max_len=127)#newtkn

    # model
    print("Creating model")
    model = ldmol_autoencoder(config=config, cp=args.enc_checkpoint, tokenizer=tokenizer, use_linear=True)
    print('#parameters:', sum(p.numel() for p in model.parameters() if p.requires_grad))

    if args.checkpoint:
        checkpoint = torch.load(args.checkpoint, map_location='cpu')
        try:
            state_dict = checkpoint['model']
        except:
            state_dict = checkpoint['state_dict']

        for key in list(state_dict.keys()):
            if 'queue' in key or 'property' in key or '_m' in key:
                del state_dict[key]
            if '_unk' in key:
                new_key = key.replace('_unk', '_mask')
                state_dict[new_key] = state_dict[key]
                del state_dict[key]

        msg = model.load_state_dict(state_dict, strict=False)
        print('load checkpoint from %s' % args.checkpoint)

    # training
    checkpoint_callback = pl.callbacks.ModelCheckpoint(dirpath=args.output_dir, filename='checkpoint_{epoch}',
                                                       save_top_k=1,
                                                       # every_n_train_steps=10000,
                                                       every_n_epochs=1
                                                       )
    ngpu = 8
    trainer = pl.Trainer(accelerator='gpu', devices=ngpu, precision='16-mixed', max_epochs=config['schedular']['epochs'],
                         callbacks=[checkpoint_callback], strategy=DDPStrategy(find_unused_parameters=False), limit_val_batches=0.,
                         # logger=WandbLogger(),
                         )
    trainer.fit(model, data_loader, None, ckpt_path=None)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--checkpoint', default='')
    parser.add_argument('--enc_checkpoint', default='./Pretrain/checkpoint_encoder.ckpt')
    parser.add_argument('--data_path', default='./data/pubchem_10m.txt')
    parser.add_argument('--resume', default=False, type=bool)
    parser.add_argument('--output_dir', default='./Pretrain')
    parser.add_argument('--vocab_filename', default='./vocab_bpe_300_sc.txt')
    parser.add_argument('--seed', default=42, type=int)
    args = parser.parse_args()

    pretrain_config = {
        'property_width': 768,
        'embed_dim': 256,
        'batch_size': 128,
        'temp': 0.07,
        'mlm_probability': 0.15,
        'momentum': 0.995,
        'alpha': 0.4,
        'bert_config_decoder': './config_decoder.json',
        'bert_config_encoder': './config_encoder.json',
        'schedular': {'sched': 'cosine', 'lr': 0.5e-4, 'epochs': 5, 'min_lr': 1e-5,
                      'decay_rate': 1, 'warmup_lr': 5e-5, 'warmup_epochs': 20, 'cooldown_epochs': 0},
        'optimizer': {'opt': 'adamW', 'lr': 0.5e-4, 'weight_decay': 0.02}
    }
    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
    main(args, pretrain_config)