Add evaluate method and option to save for each epoch in finetune

smi-ted/finetune/args.py

@@ -305,6 +305,7 @@ def get_parser(parser=None):
     parser.add_argument("--model_path", type=str, default="./smi_ted/")
     parser.add_argument("--ckpt_filename", type=str, default="smi_ted_Light_40.pt")
     # parser.add_argument('--n_output', type=int, default=1)
+    parser.add_argument("--save_every_epoch", type=int, default=0)
     parser.add_argument("--save_ckpt", type=int, default=1)
     parser.add_argument("--start_seed", type=int, default=0)
     parser.add_argument("--smi_ted_version", type=str, default="v1")

smi-ted/finetune/finetune_classification.py

@@ -48,6 +48,7 @@ def main(config):
         seed=config.start_seed,
         checkpoints_folder=config.checkpoints_folder,
         device=device,
+        save_every_epoch=bool(config.save_every_epoch),
         save_ckpt=bool(config.save_ckpt)
     )
     trainer.compile(
@@ -56,6 +57,7 @@ def main(config):
         loss_fn=loss_function
     )
     trainer.fit(max_epochs=config.max_epochs)
+    trainer.evaluate()
 
 
 if __name__ == '__main__':

smi-ted/finetune/finetune_classification_multitask.py

@@ -80,6 +80,7 @@ def main(config):
         seed=config.start_seed,
         checkpoints_folder=config.checkpoints_folder,
         device=device,
+        save_every_epoch=bool(config.save_every_epoch),
         save_ckpt=bool(config.save_ckpt)
     )
     trainer.compile(
@@ -88,6 +89,7 @@ def main(config):
         loss_fn=loss_function
     )
     trainer.fit(max_epochs=config.max_epochs)
+    trainer.evaluate()
 
 
 if __name__ == '__main__':

smi-ted/finetune/finetune_regression.py

@@ -50,6 +50,7 @@ def main(config):
         seed=config.start_seed,
         checkpoints_folder=config.checkpoints_folder,
         device=device,
+        save_every_epoch=bool(config.save_every_epoch),
         save_ckpt=bool(config.save_ckpt)
     )
     trainer.compile(
@@ -58,6 +59,7 @@ def main(config):
         loss_fn=loss_function
     )
     trainer.fit(max_epochs=config.max_epochs)
+    trainer.evaluate()
 
 
 if __name__ == '__main__':

smi-ted/finetune/trainers.py

@@ -25,7 +25,7 @@ from utils import RMSE, sensitivity, specificity
 class Trainer:
 
     def __init__(self, raw_data, dataset_name, target, batch_size, hparams,
-                 target_metric='rmse', seed=0, checkpoints_folder='./checkpoints', save_ckpt=True, device='cpu'):
+                 target_metric='rmse', seed=0, checkpoints_folder='./checkpoints', save_every_epoch=False, save_ckpt=True, device='cpu'):
         # data
         self.df_train = raw_data[0]
         self.df_valid = raw_data[1]
@@ -40,6 +40,7 @@ class Trainer:
         self.target_metric = target_metric
         self.seed = seed
         self.checkpoints_folder = checkpoints_folder
+        self.save_every_epoch = save_every_epoch
         self.save_ckpt = save_ckpt
         self.device = device
         self._set_seed(seed)
@@ -81,8 +82,7 @@ class Trainer:
         self._print_configuration()
 
     def fit(self, max_epochs=500):
-        best_vloss = 1000
-        best_vmetric = -1
+        best_vloss = float('inf')
 
         for epoch in range(1, max_epochs+1):
             print(f'\n=====Epoch [{epoch}/{max_epochs}]=====')
@@ -91,47 +91,47 @@ class Trainer:
             self.model.to(self.device)
             self.model.train()
             train_loss = self._train_one_epoch()
-            print(f'Training loss: {round(train_loss, 6)}')
 
-            # Evaluate the model
+            # validation
             self.model.eval()
             val_preds, val_loss, val_metrics = self._validate_one_epoch(self.valid_loader)
-            tst_preds, tst_loss, tst_metrics = self._validate_one_epoch(self.test_loader)
-
-            print(f"Valid loss: {round(val_loss, 6)}")
             for m in val_metrics.keys():
                 print(f"[VALID] Evaluation {m.upper()}: {round(val_metrics[m], 4)}")
-            print("-"*32)
-            print(f"Test loss: {round(tst_loss, 6)}")
-            for m in tst_metrics.keys():
-                print(f"[TEST] Evaluation {m.upper()}: {round(tst_metrics[m], 4)}")
 
             ############################### Save Finetune checkpoint #######################################
-            if (val_loss < best_vloss) and self.save_ckpt:
+            if ((val_loss < best_vloss) or self.save_every_epoch) and self.save_ckpt:
                 # remove old checkpoint
-                if best_vmetric != -1:
-                    os.remove(os.path.join(self.checkpoints_folder, filename))
+                if best_vloss != float('inf') and not self.save_every_epoch:
+                    os.remove(os.path.join(self.checkpoints_folder, self.last_filename))
 
                 # filename
                 model_name = f'{str(self.model)}-Finetune'
-                metric = round(tst_metrics[self.target_metric], 4)
-                filename = f"{model_name}_epoch={epoch}_{self.dataset_name}_seed{self.seed}_{self.target_metric}={metric}.pt"
+                self.last_filename = f"{model_name}_epoch={epoch}_{self.dataset_name}_seed{self.seed}_valloss={round(val_loss, 4)}.pt"
 
                 # save checkpoint
                 print('Saving checkpoint...')
-                self._save_checkpoint(epoch, filename)
-
-                # save predictions
-                pd.DataFrame(tst_preds).to_csv(
-                    os.path.join(
-                        self.checkpoints_folder, 
-                        f'{self.dataset_name}_{self.target if isinstance(self.target, str) else self.target[0]}_predict_test_seed{self.seed}.csv'), 
-                    index=False
-                )
+                self._save_checkpoint(epoch, self.last_filename)
 
                 # update best loss
                 best_vloss = val_loss
-                best_vmetric = metric
+
+    def evaluate(self):
+        print("\n=====Test Evaluation=====")
+        self._load_checkpoint(self.last_filename)
+        self.model.eval()
+        tst_preds, tst_loss, tst_metrics = self._validate_one_epoch(self.test_loader)
+
+        # show metrics
+        for m in tst_metrics.keys():
+            print(f"[TEST] Evaluation {m.upper()}: {round(tst_metrics[m], 4)}")
+
+        # save predictions
+        pd.DataFrame(tst_preds).to_csv(
+            os.path.join(
+                self.checkpoints_folder, 
+                f'{self.dataset_name}_{self.target if isinstance(self.target, str) else self.target[0]}_predict_test_seed{self.seed}.csv'), 
+            index=False
+        )
 
     def _train_one_epoch(self):
         raise NotImplementedError
@@ -153,6 +153,11 @@ class Trainer:
         print('Valid size:\t', self.df_valid.shape[0])
         print('Test size:\t', self.df_test.shape[0])
 
+    def _load_checkpoint(self, filename):
+        ckpt_path = os.path.join(self.checkpoints_folder, filename)
+        ckpt_dict = torch.load(ckpt_path, map_location='cpu')
+        self.model.load_state_dict(ckpt_dict['MODEL_STATE'])
+
     def _save_checkpoint(self, current_epoch, filename):
         if not os.path.exists(self.checkpoints_folder):
             os.makedirs(self.checkpoints_folder)
@@ -198,14 +203,14 @@ class Trainer:
 class TrainerRegressor(Trainer):
 
     def __init__(self, raw_data, dataset_name, target, batch_size, hparams,
-                 target_metric='rmse', seed=0, checkpoints_folder='./checkpoints', save_ckpt=True, device='cpu'):
+                 target_metric='rmse', seed=0, checkpoints_folder='./checkpoints', save_every_epoch=False, save_ckpt=True, device='cpu'):
         super().__init__(raw_data, dataset_name, target, batch_size, hparams,
-                         target_metric, seed, checkpoints_folder, save_ckpt, device) 
+                         target_metric, seed, checkpoints_folder, save_every_epoch, save_ckpt, device) 
 
     def _train_one_epoch(self):
         running_loss = 0.0
 
-        for data in tqdm(self.train_loader):
+        for idx, data in enumerate(pbar := tqdm(self.train_loader)):
             # Every data instance is an input + label pair
             smiles, targets = data
             targets = targets.clone().detach().to(self.device)
@@ -227,6 +232,11 @@ class TrainerRegressor(Trainer):
             # print statistics
             running_loss += loss.item()
 
+            # progress bar
+            pbar.set_description('[TRAINING]')
+            pbar.set_postfix(loss=running_loss/(idx+1))
+            pbar.refresh()
+
         return running_loss / len(self.train_loader)
 
     def _validate_one_epoch(self, data_loader):
@@ -235,7 +245,7 @@ class TrainerRegressor(Trainer):
         running_loss = 0.0
 
         with torch.no_grad():
-            for data in tqdm(data_loader):
+            for idx, data in enumerate(pbar := tqdm(data_loader)):
                 # Every data instance is an input + label pair
                 smiles, targets = data
                 targets = targets.clone().detach().to(self.device)
@@ -253,6 +263,11 @@ class TrainerRegressor(Trainer):
                 # print statistics
                 running_loss += loss.item()
 
+                # progress bar
+                pbar.set_description('[EVALUATION]')
+                pbar.set_postfix(loss=running_loss/(idx+1))
+                pbar.refresh()
+
         # Put together predictions and labels from batches
         preds = torch.cat(data_preds, dim=0).cpu().numpy()
         tgts = torch.cat(data_targets, dim=0).cpu().numpy()
@@ -271,20 +286,20 @@ class TrainerRegressor(Trainer):
             'spearman': spearman,
         }
 
-        return preds, running_loss / len(self.train_loader), metrics
+        return preds, running_loss / len(data_loader), metrics
 
 
 class TrainerClassifier(Trainer):
 
     def __init__(self, raw_data, dataset_name, target, batch_size, hparams,
-                 target_metric='roc-auc', seed=0, checkpoints_folder='./checkpoints', save_ckpt=True, device='cpu'):
+                 target_metric='roc-auc', seed=0, checkpoints_folder='./checkpoints', save_every_epoch=False, save_ckpt=True, device='cpu'):
         super().__init__(raw_data, dataset_name, target, batch_size, hparams,
-                         target_metric, seed, checkpoints_folder, save_ckpt, device) 
+                         target_metric, seed, checkpoints_folder, save_every_epoch, save_ckpt, device) 
 
     def _train_one_epoch(self):
         running_loss = 0.0
 
-        for data in tqdm(self.train_loader):
+        for idx, data in enumerate(pbar := tqdm(self.train_loader)):
             # Every data instance is an input + label pair
             smiles, targets = data
             targets = targets.clone().detach().to(self.device)
@@ -306,6 +321,11 @@ class TrainerClassifier(Trainer):
             # print statistics
             running_loss += loss.item()
 
+            # progress bar
+            pbar.set_description('[TRAINING]')
+            pbar.set_postfix(loss=running_loss/(idx+1))
+            pbar.refresh()
+
         return running_loss / len(self.train_loader)
 
     def _validate_one_epoch(self, data_loader):
@@ -314,7 +334,7 @@ class TrainerClassifier(Trainer):
         running_loss = 0.0
 
         with torch.no_grad():
-            for data in tqdm(data_loader):
+            for idx, data in enumerate(pbar := tqdm(data_loader)):
                 # Every data instance is an input + label pair
                 smiles, targets = data
                 targets = targets.clone().detach().to(self.device)
@@ -332,6 +352,11 @@ class TrainerClassifier(Trainer):
                 # print statistics
                 running_loss += loss.item()
 
+                # progress bar
+                pbar.set_description('[EVALUATION]')
+                pbar.set_postfix(loss=running_loss/(idx+1))
+                pbar.refresh()
+
         # Put together predictions and labels from batches
         preds = torch.cat(data_preds, dim=0).cpu().numpy()
         tgts = torch.cat(data_targets, dim=0).cpu().numpy()
@@ -366,15 +391,15 @@ class TrainerClassifier(Trainer):
             'specificity': sp,
         }
 
-        return preds, running_loss / len(self.train_loader), metrics
+        return preds, running_loss / len(data_loader), metrics
 
 
 class TrainerClassifierMultitask(Trainer):
 
     def __init__(self, raw_data, dataset_name, target, batch_size, hparams,
-                 target_metric='roc-auc', seed=0, checkpoints_folder='./checkpoints', save_ckpt=True, device='cpu'):
+                 target_metric='roc-auc', seed=0, checkpoints_folder='./checkpoints', save_every_epoch=False, save_ckpt=True, device='cpu'):
         super().__init__(raw_data, dataset_name, target, batch_size, hparams,
-                         target_metric, seed, checkpoints_folder, save_ckpt, device)
+                         target_metric, seed, checkpoints_folder, save_every_epoch, save_ckpt, device)
 
     def _prepare_data(self):
         # normalize dataset
@@ -409,7 +434,7 @@ class TrainerClassifierMultitask(Trainer):
     def _train_one_epoch(self):
         running_loss = 0.0
 
-        for data in tqdm(self.train_loader):
+        for idx, data in enumerate(pbar := tqdm(self.train_loader)):
             # Every data instance is an input + label pair + mask
             smiles, targets, target_masks = data
             targets = targets.clone().detach().to(self.device)
@@ -432,6 +457,11 @@ class TrainerClassifierMultitask(Trainer):
             # print statistics
             running_loss += loss.item()
 
+            # progress bar
+            pbar.set_description('[TRAINING]')
+            pbar.set_postfix(loss=running_loss/(idx+1))
+            pbar.refresh()
+
         return running_loss / len(self.train_loader)
 
     def _validate_one_epoch(self, data_loader):
@@ -441,7 +471,7 @@ class TrainerClassifierMultitask(Trainer):
         running_loss = 0.0
 
         with torch.no_grad():
-            for data in tqdm(data_loader):
+            for idx, data in enumerate(pbar := tqdm(data_loader)):
                 # Every data instance is an input + label pair + mask
                 smiles, targets, target_masks = data
                 targets = targets.clone().detach().to(self.device)
@@ -461,6 +491,11 @@ class TrainerClassifierMultitask(Trainer):
                 # print statistics
                 running_loss += loss.item()
 
+                # progress bar
+                pbar.set_description('[EVALUATION]')
+                pbar.set_postfix(loss=running_loss/(idx+1))
+                pbar.refresh()
+
         # Put together predictions and labels from batches
         preds = torch.cat(data_preds, dim=0)
         tgts = torch.cat(data_targets, dim=0)
@@ -513,4 +548,4 @@ class TrainerClassifierMultitask(Trainer):
             'specificity': average_sp.item(),
         }
 
-        return preds, running_loss / len(self.train_loader), metrics
+        return preds, running_loss / len(data_loader), metrics