from xbert import BertConfig from transformers import BertModel import torch import torch.nn.functional as F from torch import nn import torch.distributed import pytorch_lightning as pl from scheduler import create_scheduler import argparse from pathlib import Path from dataset import SMILESDataset_pretrain from pytorch_lightning.strategies import DDPStrategy from rdkit import Chem import random from torch.utils.data import DataLoader from pysmilesutils.augment import MolAugmenter from rdkit.Chem.EnumerateStereoisomers import EnumerateStereoisomers from utils import regexTokenizer class AttrDict(dict): def __init__(self, *args, **kwargs): super(AttrDict, self).__init__(*args, **kwargs) self.__dict__ = self class ldmol_encoder(pl.LightningModule): def __init__(self, tokenizer=None, config=None, loader_len=0, no_train=False): super().__init__() self.save_hyperparameters() self.automatic_optimization = False self.config = config self.tokenizer = tokenizer self.training_step_outputs = [] embed_dim = config['embed_dim'] bert_config = BertConfig.from_json_file(config['bert_config_encoder']) self.text_encoder = BertModel(config=bert_config) text_width = self.text_encoder.config.hidden_size self.text_proj = nn.Linear(text_width, embed_dim) self.aug = MolAugmenter() # create momentum models self.text_encoder_m = BertModel(config=bert_config) self.text_proj_m = nn.Linear(text_width, embed_dim) for p in self.text_encoder_m.parameters(): p.requires_grad = False for p in self.text_proj_m.parameters(): p.requires_grad = False self.model_pairs = [[self.text_encoder, self.text_encoder_m], [self.text_proj, self.text_proj_m], ] self.copy_params() # create the queue if not no_train: self.temp = nn.Parameter(torch.ones([]) * config['temp']) self.warmup_steps = config['schedular']['warmup_epochs'] self.loader_len = loader_len self.momentum = config['momentum'] self.queue_size = config['queue_size'] self.register_buffer("text1_queue", torch.randn(embed_dim, self.queue_size)) self.register_buffer("text2_queue", torch.randn(embed_dim, self.queue_size)) self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long)) self.text1_queue = nn.functional.normalize(self.text1_queue, dim=0) self.text2_queue = nn.functional.normalize(self.text2_queue, dim=0) def forward(self, text1_input_ids, text1_attention_mask, text2_input_ids, text2_attention_mask, alpha=0): with torch.no_grad(): self.temp.clamp_(0.07, 0.5) text1_embeds = self.text_encoder(text1_input_ids, attention_mask=text1_attention_mask, return_dict=True).last_hidden_state text1_feat = F.normalize(self.text_proj(text1_embeds[:, 0, :]), dim=-1) text2_embeds = self.text_encoder(text2_input_ids, attention_mask=text2_attention_mask, return_dict=True).last_hidden_state text2_feat = F.normalize(self.text_proj(text2_embeds[:, 0, :]), dim=-1) # get momentum features with torch.no_grad(): self._momentum_update() text1_embeds_m = self.text_encoder_m(text1_input_ids, attention_mask=text1_attention_mask, return_dict=True).last_hidden_state text1_feat_m = F.normalize(self.text_proj(text1_embeds_m[:, 0, :]), dim=-1) text1_feat_all = torch.cat([text1_feat_m.t(), self.text1_queue.clone().detach()], dim=1) text2_embeds_m = self.text_encoder_m(text2_input_ids, attention_mask=text2_attention_mask, return_dict=True).last_hidden_state text2_feat_m = F.normalize(self.text_proj(text2_embeds_m[:, 0, :]), dim=-1) text2_feat_all = torch.cat([text2_feat_m.t(), self.text2_queue.clone().detach()], dim=1) sim_21_m = text2_feat_m @ text1_feat_all / self.temp sim_12_m = text1_feat_m @ text2_feat_all / self.temp # sim_11_m = text1_feat_m @ text1_feat_all / self.temp # sim_22_m = text2_feat_m @ text2_feat_all / self.temp sim_targets = torch.zeros(sim_21_m.size()).to(self.device) sim_targets.fill_diagonal_(1) sim_21_targets = alpha * F.softmax(sim_21_m, dim=1) + (1 - alpha) * sim_targets sim_12_targets = alpha * F.softmax(sim_12_m, dim=1) + (1 - alpha) * sim_targets # sim_11_targets = alpha * F.softmax(sim_11_m, dim=1) + (1 - alpha) * sim_targets # sim_22_targets = alpha * F.softmax(sim_22_m, dim=1) + (1 - alpha) * sim_targets sim_21 = text2_feat @ text1_feat_all / self.temp sim_12 = text1_feat @ text2_feat_all / self.temp # sim_11 = text1_feat @ text1_feat_all / self.temp # sim_22 = text2_feat @ text2_feat_all / self.temp loss_21 = -torch.sum(F.log_softmax(sim_21, dim=1) * sim_21_targets, dim=1).mean() loss_12 = -torch.sum(F.log_softmax(sim_12, dim=1) * sim_12_targets, dim=1).mean() # loss_11 = -torch.sum(F.log_softmax(sim_11, dim=1) * sim_11_targets, dim=1).mean() # loss_22 = -torch.sum(F.log_softmax(sim_22, dim=1) * sim_22_targets, dim=1).mean() loss_ita = loss_21 + loss_12 # + loss_11 + loss_22 self._dequeue_and_enqueue(text1_feat_m, text2_feat_m) return loss_ita @torch.no_grad() def copy_params(self): for model_pair in self.model_pairs: for param, param_m in zip(model_pair[0].parameters(), model_pair[1].parameters()): param_m.data.copy_(param.data) # initialize param_m.requires_grad = False # not update by gradient @torch.no_grad() def _momentum_update(self): for model_pair in self.model_pairs: for param, param_m in zip(model_pair[0].parameters(), model_pair[1].parameters()): param_m.data = param_m.data * self.momentum + param.data * (1. - self.momentum) @torch.no_grad() def _dequeue_and_enqueue(self, text1_feat, text2_feat): text1_feats = concat_all_gather(text1_feat) text2_feats = concat_all_gather(text2_feat) # print(text1_feats.shape, text2_feats.shape) text1_feats = text1_feats[:64] text2_feats = text2_feats[:64] batch_size = text1_feats.shape[0] ptr = int(self.queue_ptr) assert self.queue_size % batch_size == 0 # for simplicity # replace the keys at ptr (dequeue and enqueue) self.text1_queue[:, ptr:ptr + batch_size] = text1_feats.T self.text2_queue[:, ptr:ptr + batch_size] = text2_feats.T ptr = (ptr + batch_size) % self.queue_size # move pointer self.queue_ptr[0] = ptr 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() text1 = train_batch text1 = [t.split('Q') for t in text1] tmp = [] for t in text1: tmp += t text2 = [] text1 = [] for t in tmp: try: t2 = '[CLS]' + Chem.MolToSmiles(self.aug([Chem.MolFromSmiles(t[5:])])[0], canonical=False, isomericSmiles=True) text1.append(t) text2.append(t2) except: print('err', t) continue # text_input1 = self.tokenizer(text1, padding='longest', truncation=True, max_length=128, return_tensors="pt").to(self.device) # text_input2 = self.tokenizer(text2, padding='longest', truncation=True, max_length=128, return_tensors="pt").to(self.device) text_input_ids = self.tokenizer(text1, truncation='longest').to(self.device) text_attention_mask = torch.where(text_input_ids == 0, 0, 1).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) alpha = self.config['alpha'] if self.current_epoch > 0 else self.config['alpha'] * min(1., batch_idx / self.loader_len) # loss = self(text_input1.input_ids[:, 1:], text_input1.attention_mask[:, 1:], text_input2.input_ids[:, 1:], text_input2.attention_mask[:, 1:], alpha=alpha) loss = self(text_input_ids, text_attention_mask, text2_input_ids, text2_attention_mask, alpha=alpha) 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', loss, prog_bar=True) self.log('temp', self.temp, 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])) return torch.tensor([loss]) 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() @torch.no_grad() def concat_all_gather(tensor): """ Performs all_gather operation on the provided tensors. *** Warning ***: torch.distributed.all_gather has no gradient. """ tensors_gather = [torch.ones_like(tensor) for _ in range(torch.distributed.get_world_size())] torch.distributed.all_gather(tensors_gather, tensor, async_op=False) output = torch.cat(tensors_gather, dim=0) return output def main(args, config): # data print("Creating dataset") dataset = SMILESDataset_pretrain(args.data_path, data_length=[0, 10000000]) print('#data:', len(dataset)) data_loader = DataLoader(dataset, batch_size=config['batch_size'], num_workers=16, shuffle=False, pin_memory=True, drop_last=True) tokenizer = regexTokenizer(vocab_path=args.vocab_filename, max_len=127)#newtkn # model print("Creating model") ngpu=1 model = ldmol_encoder(config=config, tokenizer=tokenizer, loader_len=len(data_loader) // ngpu) 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'] msg = model.load_state_dict(state_dict, strict=False) print('load checkpoint from %s' % args.checkpoint) print(msg) # 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 ) trainer = pl.Trainer(accelerator='gpu', devices=ngpu, precision='16-mixed', max_epochs=config['schedular']['epochs'], callbacks=[checkpoint_callback], strategy=DDPStrategy(find_unused_parameters=True), limit_val_batches=0.) trainer.fit(model, data_loader, None, ckpt_path=None) @torch.no_grad() def evaluate(args, config): # data print("Creating dataset") dataset = SMILESDataset_pretrain(args.data_path, data_length=[0, 10], is_train=False) print('#data:', len(dataset)) data_loader = DataLoader(dataset, batch_size=config['batch_size'], num_workers=8, shuffle=True, pin_memory=True, drop_last=False) tokenizer = regexTokenizer(vocab_path=args.vocab_filename, max_len=127)#newtkn # model print("Creating model") ngpu = 1 model = ldmol_encoder(config=config, tokenizer=tokenizer, loader_len=len(data_loader) // ngpu) 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'] msg = model.load_state_dict(state_dict, strict=False) print('load checkpoint from %s' % args.checkpoint) print(msg) model = model.to('cuda') model.eval() for text1 in data_loader: # text1 = ['c1ccccc1', 'c1ccccc1O', 'c1ccccc1F', 'c1ccccc1C'] text1 = [Chem.MolToSmiles(Chem.MolFromSmiles(t[5:]), canonical=True, isomericSmiles=False) for t in text1] text1_sc = [Chem.MolToSmiles(random.choice(list(EnumerateStereoisomers(Chem.MolFromSmiles(t)))), canonical=True, isomericSmiles=True) for t in text1] text1 = [Chem.MolToSmiles(random.choice(list(EnumerateStereoisomers(Chem.MolFromSmiles(t)))), canonical=True, isomericSmiles=True) for t in text1] text1_combined = [] for i in range(len(text1)): text1_combined.append(text1[i]) text1_combined.append(text1_sc[i]) print(text1_combined) text1 = ['[CLS]' + t for t in text1_combined] text2 = ['[CLS]' + Chem.MolToSmiles(model.aug([Chem.MolFromSmiles(t[5:])])[0], canonical=False, isomericSmiles=True) for t in text1] # text_input1 = model.tokenizer(text1, padding='longest', truncation=True, max_length=128, return_tensors="pt").to(model.device) # text_input2 = model.tokenizer(text2, padding='longest', truncation=True, max_length=128, return_tensors="pt").to(model.device) text_input_ids = model.tokenizer(text1, truncation='longest').to(model.device) text_attention_mask = torch.where(text_input_ids == 0, 0, 1).to(model.device) text2_input_ids = model.tokenizer(text2, truncation='longest').to(model.device) text2_attention_mask = torch.where(text2_input_ids == 0, 0, 1).to(model.device) # text1_embeds = model.text_encoder(text_input1.input_ids[:, 1:], attention_mask=text_input1.attention_mask[:, 1:], return_dict=True).last_hidden_state text1_embeds = model.text_encoder(text_input_ids, attention_mask=text_attention_mask, return_dict=True).last_hidden_state text1_feat = F.normalize(model.text_proj(text1_embeds[:, 0, :]), dim=-1) # text2_embeds = model.text_encoder(text_input2.input_ids[:, 1:], attention_mask=text_input2.attention_mask[:, 1:], return_dict=True).last_hidden_state text2_embeds = model.text_encoder(text2_input_ids, attention_mask=text2_attention_mask, return_dict=True).last_hidden_state text2_feat = F.normalize(model.text_proj(text2_embeds[:, 0, :]), dim=-1) sim = text1_feat @ text2_feat.T print(sim) break if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--checkpoint', default="") # "./Pretrain/checkpoint_smauCLIP_sc.ckpt" parser.add_argument('--data_path', default='./data/unpaired_200k.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 = { 'embed_dim': 256, 'batch_size': 64, 'temp': 0.07, 'queue_size': 16384, 'momentum': 0.995, 'alpha': 0.4, 'bert_config_encoder': './config_encoder.json', 'schedular': {'sched': 'cosine', 'lr': 1e-4, 'epochs': 5, 'min_lr': 1e-5, 'decay_rate': 1, 'warmup_lr': 5e-5, 'warmup_epochs': 20, 'cooldown_epochs': 0}, 'optimizer': {'opt': 'adamW', 'lr': 1e-4, 'weight_decay': 0.02} } Path(args.output_dir).mkdir(parents=True, exist_ok=True) main(args, pretrain_config) # evaluate(args, pretrain_config)