Add files via upload

Author: jinhojsk515

SPMM_models.py

@@ -5,6 +5,7 @@
 import torch.distributed
 import pytorch_lightning as pl
 from scheduler import create_scheduler
+import random
 
 
 class AttrDict(dict):
@@ -16,6 +17,7 @@ def __init__(self, *args, **kwargs):
 class SPMM(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
@@ -82,13 +84,16 @@ def forward(self, property_original, text_input_ids, text_attention_mask, alpha=
         property_feature = self.property_embed(property_original.unsqueeze(2))
 
         unk_tokens = self.property_mask.expand(property_original.size(0), property_original.size(1), -1)
-        mpm_mask = torch.bernoulli(torch.ones_like(property_original) * 0.5)
+        if random.random() < 0.05:
+            mpm_mask = torch.ones_like(property_original)                           # all mask
+        else:
+            mpm_mask = torch.bernoulli(torch.ones_like(property_original) * 0.5)    # 1 for mask, 0 for keep
         mpm_mask_expand = mpm_mask.unsqueeze(2).repeat(1, 1, unk_tokens.size(2))
         property_masked = property_feature * (1 - mpm_mask_expand) + unk_tokens * mpm_mask_expand
-        property = torch.cat([self.property_cls.expand(property_original.size(0), -1, -1), property_masked], dim=1)
+        properties = torch.cat([self.property_cls.expand(property_original.size(0), -1, -1), property_masked], dim=1)
 
-        prop_embeds = self.property_encoder(inputs_embeds=property, return_dict=True).last_hidden_state
-        prop_atts = torch.ones(prop_embeds.size()[:-1], dtype=torch.long).to(property.device)
+        prop_embeds = self.property_encoder(inputs_embeds=properties, return_dict=True).last_hidden_state
+        prop_atts = torch.ones(prop_embeds.size()[:-1], dtype=torch.long).to(properties.device)
         prop_feat = F.normalize(self.property_proj(prop_embeds[:, 0, :]), dim=-1)
 
         text_embeds = self.text_encoder.bert(text_input_ids, attention_mask=text_attention_mask, return_dict=True, mode='text').last_hidden_state
@@ -97,7 +102,7 @@ def forward(self, property_original, text_input_ids, text_attention_mask, alpha=
 
         with torch.no_grad():
             self._momentum_update()
-            prop_embeds_m = self.property_encoder_m(inputs_embeds=property, return_dict=True).last_hidden_state
+            prop_embeds_m = self.property_encoder_m(inputs_embeds=properties, return_dict=True).last_hidden_state
             prop_feat_m = F.normalize(self.property_proj_m(prop_embeds_m[:, 0, :]), dim=-1)
             prop_feat_all = torch.cat([prop_feat_m.t(), self.prop_queue.clone().detach()], dim=1)
 
@@ -110,7 +115,7 @@ def forward(self, property_original, text_input_ids, text_attention_mask, alpha=
             sim_i2i_m = prop_feat_m @ prop_feat_all / self.temp
             sim_t2t_m = text_feat_m @ text_feat_all / self.temp
 
-            sim_targets = torch.zeros(sim_i2t_m.size()).to(property.device)
+            sim_targets = torch.zeros(sim_i2t_m.size()).to(properties.device)
             sim_targets.fill_diagonal_(1)
 
             sim_i2t_targets = alpha * F.softmax(sim_i2t_m, dim=1) + (1 - alpha) * sim_targets
@@ -268,8 +273,8 @@ def _momentum_update(self):
 
     @torch.no_grad()
     def _dequeue_and_enqueue(self, img_feat, text_feat):
-        img_feats = img_feat
-        text_feats = text_feat
+        img_feats = concat_all_gather(img_feat)
+        text_feats = concat_all_gather(text_feat)
 
         batch_size = img_feats.shape[0]
 
@@ -354,3 +359,16 @@ def on_train_epoch_end(self):    # outputs: collection of returns from 'training
         if self.global_rank == 0:
             print(f'\n mean loss: {tmp[0]:.4f}, {tmp[1]:.4f}, {tmp[2]:.4f}, {tmp[3]:.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

SPMM_pretrain.py

@@ -12,7 +12,7 @@
 def main(args, config):
     # data
     print("Creating dataset")
-    dataset = SMILESDataset_pretrain(args.data_path)
+    dataset = SMILESDataset_pretrain(args.data_path, data_length=[0, 10000])
     print('#data:', len(dataset))
     data_loader = DataLoader(dataset, batch_size=config['batch_size'], num_workers=8, shuffle=True, pin_memory=True, drop_last=True)
     tokenizer = BertTokenizer(vocab_file=args.vocab_filename, do_lower_case=False, do_basic_tokenize=False)
@@ -27,17 +27,18 @@ def main(args, config):
     # training
     checkpoint_callback = pl.callbacks.ModelCheckpoint(dirpath=args.output_dir, filename='checkpoint_{epoch}',
                                                        save_top_k=config['schedular']['epochs'], monitor='loss_mlm')
-    trainer = pl.Trainer(accelerator='gpu', devices=[0], precision=16, max_epochs=config['schedular']['epochs'],
+    trainer = pl.Trainer(accelerator='gpu', devices=[0, 1], precision=16, 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=args.checkpoint if args.checkpoint else None)
 
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('--checkpoint', default='')
-    parser.add_argument('--data_path', default='./data/1_Pretrain/pretrain_20m.txt')
+    # parser.add_argument('--data_path', default='./data/1_Pretrain/pretrain_20m.txt')
+    parser.add_argument('--data_path', default='../VLP_chem/data/pubchem-100m-simple-shuffle.txt')
     parser.add_argument('--resume', default=False, type=bool)
-    parser.add_argument('--output_dir', default='./checkpoints')
+    parser.add_argument('--output_dir', default='./Pretrain')
     parser.add_argument('--vocab_filename', default='./vocab_bpe_300.txt')
     parser.add_argument('--seed', default=42, type=int)
     args = parser.parse_args()

d_pv2smiles_stochastic.py

@@ -14,7 +14,9 @@
 
 
 @torch.no_grad()
-def generate_with_property(model, property, tokenizer, device, n_sample, prop_mask):
+def generate_with_property(model, property, n_sample, prop_mask, stochastic=True):
+    device = model.device
+    tokenizer = model.tokenizer
     # test
     model.eval()
     print("PV-to-SMILES generation in stochastic manner...")
@@ -40,7 +42,7 @@ def generate_with_property(model, property, tokenizer, device, n_sample, prop_ma
         text_input = torch.tensor([tokenizer.cls_token_id]).expand(prop.size(0), 1).to(device)
         end_count = torch.zeros_like(text_input).to(bool)
         for _ in range(100):
-            output = generate(model, prop_embeds, text_input, stochastic=True)
+            output = generate(model, prop_embeds, text_input, stochastic=stochastic)
             end_count = torch.logical_or(end_count, (output == tokenizer.sep_token_id))
             if end_count.all():
                 break
@@ -160,7 +162,7 @@ def main(args, config):
     # prop_input = torch.zeros(53)
 
     print("=" * 50)
-    samples = generate_with_property(model, prop_input, tokenizer, device, args.n_generate, prop_mask)
+    samples = generate_with_property(model, prop_input, args.n_generate, prop_mask)
     metric_eval(prop_input, samples, prop_mask)
     print("=" * 50)
 

d_smiles2pv.py

@@ -28,10 +28,33 @@ def generate(model, prop_input, text_embeds, text_atts):
 
 
 @torch.no_grad()
-def evaluate(model, data_loader, tokenizer, device):
+def pv_generate(model, data_loader):
     # test
+    with open('./normalize.pkl', 'rb') as w:
+        mean, std = pickle.load(w)
+    device = model.device
+    tokenizer = model.tokenizer
     model.eval()
     print("SMILES-to-PV generation...")
+    # convert list of string to dataloader
+    if isinstance(data_loader, list):
+        gather = []
+        text_input = tokenizer(data_loader, padding='longest', truncation=True, max_length=100, return_tensors="pt").to(device)
+        text_embeds = model.text_encoder.bert(text_input.input_ids[:, 1:], attention_mask=text_input.attention_mask[:, 1:],
+                                              return_dict=True, mode='text').last_hidden_state
+        prop_input = model.property_cls.expand(len(data_loader), -1, -1)
+        prediction = []
+        for _ in range(53):
+            output = generate(model, prop_input, text_embeds, text_input.attention_mask[:, 1:])
+            prediction.append(output)
+            output = model.property_embed(output.unsqueeze(2))
+            prop_input = torch.cat([prop_input, output], dim=1)
+
+        prediction = torch.stack(prediction, dim=-1)
+        for i in range(len(data_loader)):
+            gather.append(prediction[i].cpu()*std + mean)
+        return gather
+
     reference, candidate = [], []
     for (prop, text) in data_loader:
         text_input = tokenizer(text, padding='longest', truncation=True, max_length=100, return_tensors="pt").to(device)
@@ -139,7 +162,7 @@ def main(args, config):
     model = model.to(device)
 
     print("=" * 50)
-    r_test, c_test = evaluate(model, test_loader, tokenizer, device)
+    r_test, c_test = pv_generate(model, test_loader)
     metric_eval(r_test, c_test)
     print("=" * 50)
 

dataset.py

@@ -12,15 +12,23 @@
 
 class SMILESDataset_pretrain(Dataset):
     def __init__(self, data_path, data_length=None, shuffle=False):
-        with open(data_path, 'r') as f:
-            lines = f.readlines()
+        if data_length is not None:
+            with open(data_path, 'r') as f:
+                for _ in range(data_length[0]):
+                    f.readline()
+                lines = []
+                for _ in range(data_length[1] - data_length[0]):
+                    lines.append(f.readline())
+        else:
+            with open(data_path, 'r') as f:
+                lines = f.readlines()
         self.data = [l.strip() for l in lines]
         with open('./normalize.pkl', 'rb') as w:
             norm = pickle.load(w)
         self.property_mean, self.property_std = norm
 
-        if shuffle: random.shuffle(self.data)
-        if data_length is not None: self.data = self.data[data_length[0]:data_length[1]]
+        if shuffle:
+            random.shuffle(self.data)
 
     def __len__(self):
         return len(self.data)