Update trainer.py

trainer.py

@@ -6,7 +6,7 @@ import torch
 from utils import *
 from models import Generator, Generator2, simple_disc
 import torch_geometric.utils as geoutils
-#import #wandb
+#import wandb
 import re
 from torch_geometric.loader import DataLoader
 from new_dataloader import DruggenDataset
@@ -19,7 +19,7 @@ RDLogger.DisableLog('rdApp.*')
 from loss import discriminator_loss, generator_loss, discriminator2_loss, generator2_loss
 from training_data import load_data
 import random
-
+from tqdm import tqdm
 
 class Trainer(object):
     
@@ -27,6 +27,19 @@ class Trainer(object):
 
     def __init__(self, config):
         
+        if config.set_seed:
+            np.random.seed(config.seed)
+            random.seed(config.seed)
+            torch.manual_seed(config.seed)
+            torch.cuda.manual_seed(config.seed)
+
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+
+            os.environ["PYTHONHASHSEED"] = str(config.seed)
+
+            print(f'Using seed {config.seed}')
+
         self.device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')
         """Initialize configurations."""
         self.submodel = config.submodel
@@ -57,7 +70,10 @@ class Trainer(object):
         
         self.inf_drugs_dataset_file = config.inf_drug_dataset_file  # Drug dataset file name for the second GAN. 
                                                             # Contains drug molecules only. (In this case AKT1 inhibitors.)
-
+        self.inference_iterations = config.inference_iterations
+        
+        self.inf_batch_size = config.inf_batch_size
+        
         self.mol_data_dir = config.mol_data_dir  # Directory where the dataset files are stored.
         
         self.drug_data_dir = config.drug_data_dir  # Directory where the drug dataset files are stored.
@@ -219,6 +235,14 @@ class Trainer(object):
         self.clipping_value = config.clipping_value
         # Miscellaneous.
         
+        # resume training
+
+        self.resume = config.resume
+        self.resume_epoch = config.resume_epoch
+        self.resume_iter = config.resume_iter
+        self.resume_directory = config.resume_directory
+
+
         self.mode = config.mode
 
         self.noise_strength_0 = torch.nn.Parameter(torch.zeros([]))
@@ -398,7 +422,7 @@ class Trainer(object):
         
         ''' Loading the atom and bond decoders'''
         
-        with open("data/decoders/" + dictionary_name + "_" + self.dataset_name + '.pkl', 'rb') as f:
+        with open("DrugGEN/data/decoders/" + dictionary_name + "_" + self.dataset_name + '.pkl', 'rb') as f:
             
             return pickle.load(f)    
 
@@ -406,7 +430,7 @@ class Trainer(object):
         
         ''' Loading the atom and bond decoders'''
         
-        with open("data/decoders/" + dictionary_name +"_" + self.drugs_name +'.pkl', 'rb') as f:
+        with open("DrugGEN/data/decoders/" + dictionary_name +"_" + self.drugs_name +'.pkl', 'rb') as f:
             
             return pickle.load(f)    
  
@@ -429,17 +453,17 @@ class Trainer(object):
         print('Loading the trained models from epoch / iteration {}-{}...'.format(epoch, iteration))
         
         G_path = os.path.join(model_directory, '{}-{}-G.ckpt'.format(epoch, iteration))
-        #D_path = os.path.join(model_directory, '{}-{}-D.ckpt'.format(epoch, iteration))
+        D_path = os.path.join(model_directory, '{}-{}-D.ckpt'.format(epoch, iteration))
         
         self.G.load_state_dict(torch.load(G_path, map_location=lambda storage, loc: storage))
-        #self.D.load_state_dict(torch.load(D_path, map_location=lambda storage, loc: storage))
+        self.D.load_state_dict(torch.load(D_path, map_location=lambda storage, loc: storage))
       
         
         G2_path = os.path.join(model_directory, '{}-{}-G2.ckpt'.format(epoch, iteration))
-        #D2_path = os.path.join(model_directory, '{}-{}-D2.ckpt'.format(epoch, iteration))
+        D2_path = os.path.join(model_directory, '{}-{}-D2.ckpt'.format(epoch, iteration))
         
         self.G2.load_state_dict(torch.load(G2_path, map_location=lambda storage, loc: storage))
-        #self.D2.load_state_dict(torch.load(D2_path, map_location=lambda storage, loc: storage))
+        self.D2.load_state_dict(torch.load(D2_path, map_location=lambda storage, loc: storage))
 
    
     def save_model(self, model_directory, idx,i):
@@ -507,16 +531,19 @@ class Trainer(object):
 
 
         # protein data
-        full_smiles = [line for line in open("data/chembl_train.smi", 'r').read().splitlines()]
-        drug_smiles = [line for line in open("data/akt_train.smi", 'r').read().splitlines()]
+        full_smiles = [line for line in open("DrugGEN/data/chembl_train.smi", 'r').read().splitlines()]
+        drug_smiles = [line for line in open("DrugGEN/data/akt_train.smi", 'r').read().splitlines()]
         
         drug_mols = [Chem.MolFromSmiles(smi) for smi in drug_smiles]
         drug_scaf = [MurckoScaffold.GetScaffoldForMol(x) for x in drug_mols]
         fps_r = [Chem.RDKFingerprint(x) for x in drug_scaf]
 
-        akt1_human_adj = torch.load("data/akt/AKT1_human_adj.pt").reshape(1,-1).to(self.device).float() 
-        akt1_human_annot = torch.load("data/akt/AKT1_human_annot.pt").reshape(1,-1).to(self.device).float() 
-      
+        akt1_human_adj = torch.load("DrugGEN/data/akt/AKT1_human_adj.pt").reshape(1,-1).to(self.device).float() 
+        akt1_human_annot = torch.load("DrugGEN/data/akt/AKT1_human_annot.pt").reshape(1,-1).to(self.device).float() 
+
+        if self.resume:
+            self.restore_model(self.resume_epoch, self.resume_iter, self.resume_directory)
+
         # Start training.
         
         print('Start training...')
@@ -577,8 +604,8 @@ class Trainer(object):
                     GAN2_disc_e = drugs_a_tensor
                     GAN2_disc_x = drugs_x_tensor        
                 elif self.submodel == "RL":
-                    GAN1_input_e = z_edge
-                    GAN1_input_x = z_node
+                    GAN1_input_e = a_tensor
+                    GAN1_input_x = x_tensor
                     GAN1_disc_e = a_tensor
                     GAN1_disc_x = x_tensor
                     GAN2_input_e = drugs_a_tensor
@@ -586,8 +613,8 @@ class Trainer(object):
                     GAN2_disc_e = drugs_a_tensor
                     GAN2_disc_x = drugs_x_tensor    
                 elif self.submodel == "NoTarget":
-                    GAN1_input_e = z_edge 
-                    GAN1_input_x = z_node 
+                    GAN1_input_e = a_tensor 
+                    GAN1_input_x = x_tensor 
                     GAN1_disc_e = a_tensor
                     GAN1_disc_x = x_tensor                                                  
                     
@@ -639,9 +666,10 @@ class Trainer(object):
                                                     GAN1_input_x,
                                                     self.batch_size,
                                                     sim_reward,
-                                                    self.dataset.matrices2mol_drugs,
+                                                    self.dataset.matrices2mol,
                                                     fps_r,
-                                                    self.submodel)        
+                                                    self.submodel,
+                                                    self.dataset_name)        
                 
                 g_loss, fake_mol, g_edges_hat_sample, g_nodes_hat_sample, node, edge = generator_output    
             
@@ -659,7 +687,8 @@ class Trainer(object):
                                                 fps_r,
                                                 GAN2_input_e,
                                                 GAN2_input_x,
-                                                self.submodel)
+                                                self.submodel,
+                                                self.drugs_name)
                 
                     g2_loss, fake_mol_g, dr_g_edges_hat_sample, dr_g_nodes_hat_sample = output     
                 
@@ -695,31 +724,31 @@ class Trainer(object):
         
         # Load the trained generator.
         self.G.to(self.device)
-        #self.D.to(self.device)
         self.G2.to(self.device)
-        #self.D2.to(self.device)        
         
         G_path = os.path.join(self.inference_model, '{}-G.ckpt'.format(self.submodel))
         self.G.load_state_dict(torch.load(G_path, map_location=lambda storage, loc: storage))
-        G2_path = os.path.join(self.inference_model, '{}-G2.ckpt'.format(self.submodel))
-        self.G2.load_state_dict(torch.load(G2_path, map_location=lambda storage, loc: storage))        
+        if self.submodel != "NoTarget" and self.submodel != "CrossLoss":
+            G2_path = os.path.join(self.inference_model, '{}-G2.ckpt'.format(self.submodel))
+            self.G2.load_state_dict(torch.load(G2_path, map_location=lambda storage, loc: storage))           
         
-        print(G_path)
-        drug_smiles = [line for line in open("data/akt_test.smi", 'r').read().splitlines()]
+        
+        drug_smiles = [line for line in open("DrugGEN/data/akt_test.smi", 'r').read().splitlines()]
         
         drug_mols = [Chem.MolFromSmiles(smi) for smi in drug_smiles]
         drug_scaf = [MurckoScaffold.GetScaffoldForMol(x) for x in drug_mols]
         fps_r = [Chem.RDKFingerprint(x) for x in drug_scaf]
 
-        akt1_human_adj = torch.load("data/akt/AKT1_human_adj.pt").reshape(1,-1).to(self.device).float() 
-        akt1_human_annot = torch.load("data/akt/AKT1_human_annot.pt").reshape(1,-1).to(self.device).float() 
+        akt1_human_adj = torch.load("DrugGEN/data/akt/AKT1_human_adj.pt").reshape(1,-1).to(self.device).float() 
+        akt1_human_annot = torch.load("DrugGEN/data/akt/AKT1_human_annot.pt").reshape(1,-1).to(self.device).float() 
         
         self.G.eval()
         #self.D.eval()
         self.G2.eval()
         #self.D2.eval()
+
+        step = self.inference_iterations
         
-        self.inf_batch_size =256
         self.inf_dataset = DruggenDataset(self.mol_data_dir,
                                       self.inf_dataset_file, 
                                       self.inf_raw_file, 
@@ -753,24 +782,25 @@ class Trainer(object):
         #metric_calc_mol = []
         metric_calc_dr = []
         date = time.time()
-        if not os.path.exists("experiments/inference/{}".format(self.submodel)):
-            os.makedirs("experiments/inference/{}".format(self.submodel))
+        if not os.path.exists("DrugGEN/experiments/inference/{}".format(self.submodel)):
+            os.makedirs("DrugGEN/experiments/inference/{}".format(self.submodel))
         with torch.inference_mode():
             
-            dataloader_iterator = iter(self.drugs_loader)
-            
-            for i, data in enumerate(self.loader):   
+            dataloader_iterator = iter(self.inf_drugs_loader)
+            pbar = tqdm(range(self.inference_sample_num))
+            pbar.set_description('Inference mode for {} model started'.format(self.submodel))
+            for i, data in enumerate(self.inf_loader):   
                 try:
                     drugs = next(dataloader_iterator)
                 except StopIteration:
-                    dataloader_iterator = iter(self.drugs_loader)
+                    dataloader_iterator = iter(self.inf_drugs_loader)
                     drugs = next(dataloader_iterator)
 
                 # Preprocess both dataset 
                 
                 bulk_data = load_data(data,
                                      drugs,
-                                     self.batch_size, 
+                                     self.inf_batch_size, 
                                      self.device,
                                      self.b_dim,
                                      self.m_dim,
@@ -809,8 +839,8 @@ class Trainer(object):
                     GAN2_disc_e = drugs_a_tensor
                     GAN2_disc_x = drugs_x_tensor        
                 elif self.submodel == "RL":
-                    GAN1_input_e = z_edge
-                    GAN1_input_x = z_node
+                    GAN1_input_e = a_tensor
+                    GAN1_input_x = x_tensor
                     GAN1_disc_e = a_tensor
                     GAN1_disc_x = x_tensor
                     GAN2_input_e = drugs_a_tensor
@@ -818,8 +848,8 @@ class Trainer(object):
                     GAN2_disc_e = drugs_a_tensor
                     GAN2_disc_x = drugs_x_tensor    
                 elif self.submodel == "NoTarget":
-                    GAN1_input_e = z_edge
-                    GAN1_input_x = z_node
+                    GAN1_input_e = a_tensor
+                    GAN1_input_x = x_tensor
                     GAN1_disc_e = a_tensor
                     GAN1_disc_x = x_tensor      
                 # =================================================================================== #
@@ -830,53 +860,50 @@ class Trainer(object):
                                                     self.V,
                                                     GAN1_input_e,
                                                     GAN1_input_x,
-                                                    self.batch_size,
+                                                    self.inf_batch_size,
                                                     sim_reward,
-                                                    self.dataset.matrices2mol_drugs,
+                                                    self.dataset.matrices2mol,
                                                     fps_r,
-                                                    self.submodel)   
+                                                    self.submodel,
+                                                    self.dataset_name)   
                 
-                _, fake_mol, _, _, node, edge = generator_output  
+                _, fake_mol_g, _, _, node, edge = generator_output  
 
                 # =================================================================================== #
                 #                             3. GAN2 Inference                                       #
                 # =================================================================================== #   
                         
-                output = generator2_loss(self.G2,
-                                            self.D2,
-                                            self.V2,
-                                            edge,
-                                            node,
-                                            self.batch_size,
-                                            sim_reward,
-                                            self.dataset.matrices2mol_drugs,
-                                            fps_r,
-                                            GAN2_input_e,
-                                            GAN2_input_x,
-                                            self.submodel)
+                if self.submodel != "NoTarget" and self.submodel != "CrossLoss":      
+                    output = generator2_loss(self.G2,
+                                                self.D2,
+                                                self.V2,
+                                                edge,
+                                                node,
+                                                self.inf_batch_size,
+                                                sim_reward,
+                                                self.dataset.matrices2mol_drugs,
+                                                fps_r,
+                                                GAN2_input_e,
+                                                GAN2_input_x,
+                                                self.submodel,
+                                                self.drugs_name)
             
-                _, fake_mol_g, _, _ = output     
+                    _, fake_mol_g, edges, nodes = output      
 
                 inference_drugs = [Chem.MolToSmiles(line) for line in fake_mol_g if line is not None]   
+                inference_drugs = [None if x is None else max(x.split('.'), key=len) for x in inference_drugs]
 
-
-                
-                #inference_smiles = [Chem.MolToSmiles(line) for line in fake_mol]  
-                 
-                
-
-                print("molecule batch {} inferred".format(i))
-
-                with open("experiments/inference/{}/inference_drugs.txt".format(self.submodel), "a") as f:
+                with open("DrugGEN/experiments/inference/{}/inference_drugs.txt".format(self.submodel), "a") as f:
                     for molecules in inference_drugs:
                         
                         f.write(molecules)
                         f.write("\n")
                         metric_calc_dr.append(molecules)
             
-                 
-                                            
-                if i == 120:
+                if len(inference_drugs) > 0:
+                    pbar.update(1)
+                              
+                if len(metric_calc_dr) == self.inference_sample_num:
                     break
         
         et = time.time() - start_time
@@ -885,8 +912,8 @@ class Trainer(object):
         
         print("Metrics calculation started using MOSES.")
                    
-        print("Validity: ", fraction_valid(inference_drugs), "\n")
-        print("Uniqueness: ", fraction_unique(inference_drugs), "\n")
-        print("Validity: ", novelty(inference_drugs, drug_smiles), "\n")
+        print("Validity: ", fraction_valid(metric_calc_dr), "\n")
+        print("Uniqueness: ", fraction_unique(metric_calc_dr), "\n")
+        print("Validity: ", novelty(metric_calc_dr, drug_smiles), "\n")
 
-        print("Metrics are calculated.")
+        print("Metrics are calculated.")