Update finetune code

smi-ted/finetune/args.py

@@ -304,6 +304,7 @@ def get_parser(parser=None):
     # parser.add_argument("--patience_epochs", type=int, required=True)
     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("--restart_filename", type=str, default="")
     # 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)

smi-ted/finetune/finetune_classification.py

@@ -28,7 +28,7 @@ def main(config):
     elif config.smi_ted_version == 'v2':
         from smi_ted_large.load import load_smi_ted
 
-    model = load_smi_ted(folder=config.model_path, ckpt_filename=config.ckpt_filename, n_output=config.n_output)
+    model = load_smi_ted(folder=config.model_path, ckpt_filename=config.ckpt_filename, n_output=config.n_output, eval=False)
     model.net.apply(model._init_weights)
     print(model.net)
 
@@ -46,7 +46,9 @@ def main(config):
         hparams=config,
         target_metric=config.target_metric,
         seed=config.start_seed,
+        smi_ted_version=config.smi_ted_version,
         checkpoints_folder=config.checkpoints_folder,
+        restart_filename=config.restart_filename,
         device=device,
         save_every_epoch=bool(config.save_every_epoch),
         save_ckpt=bool(config.save_ckpt)

smi-ted/finetune/finetune_classification_multitask.py

@@ -48,6 +48,7 @@ def main(config):
             'MUV-692', 'MUV-712', 'MUV-713', 'MUV-733', 'MUV-737', 'MUV-810',
             'MUV-832', 'MUV-846', 'MUV-852', 'MUV-858', 'MUV-859'
         ]
+    config.n_output = len(targets)
 
     # load dataset
     df_train = pd.read_csv(f"{config.data_root}/train.csv")
@@ -60,7 +61,7 @@ def main(config):
     elif config.smi_ted_version == 'v2':
         from smi_ted_large.load import load_smi_ted
 
-    model = load_smi_ted(folder=config.model_path, ckpt_filename=config.ckpt_filename, n_output=len(targets))
+    model = load_smi_ted(folder=config.model_path, ckpt_filename=config.ckpt_filename, n_output=len(targets), eval=False)
     model.net.apply(model._init_weights)
     print(model.net)
 
@@ -78,7 +79,9 @@ def main(config):
         hparams=config,
         target_metric=config.target_metric,
         seed=config.start_seed,
+        smi_ted_version=config.smi_ted_version,
         checkpoints_folder=config.checkpoints_folder,
+        restart_filename=config.restart_filename,
         device=device,
         save_every_epoch=bool(config.save_every_epoch),
         save_ckpt=bool(config.save_ckpt)

smi-ted/finetune/finetune_regression.py

@@ -28,7 +28,7 @@ def main(config):
     elif config.smi_ted_version == 'v2':
         from smi_ted_large.load import load_smi_ted
 
-    model = load_smi_ted(folder=config.model_path, ckpt_filename=config.ckpt_filename, n_output=config.n_output)
+    model = load_smi_ted(folder=config.model_path, ckpt_filename=config.ckpt_filename, n_output=config.n_output, eval=False)
     model.net.apply(model._init_weights)
     print(model.net)
 
@@ -48,7 +48,9 @@ def main(config):
         hparams=config,
         target_metric=config.target_metric,
         seed=config.start_seed,
+        smi_ted_version=config.smi_ted_version,
         checkpoints_folder=config.checkpoints_folder,
+        restart_filename=config.restart_filename,
         device=device,
         save_every_epoch=bool(config.save_every_epoch),
         save_ckpt=bool(config.save_ckpt)

smi-ted/finetune/smi_ted_large/load.py

@@ -318,7 +318,7 @@ class Net(nn.Module):
 
 class MoLEncoder(nn.Module):
 
-    def __init__(self, config, n_vocab):
+    def __init__(self, config, n_vocab, eval=False):
         super(MoLEncoder, self).__init__()
 
         # embeddings
@@ -337,7 +337,7 @@ class MoLEncoder(nn.Module):
             # unless we do deterministic_eval here, we will have random outputs
             feature_map=partial(GeneralizedRandomFeatures, 
                                 n_dims=config['num_feats'], 
-                                deterministic_eval=False),
+                                deterministic_eval=eval),
             activation='gelu'
         )
         self.blocks = builder.get()
@@ -361,7 +361,7 @@ class MoLDecoder(nn.Module):
 class Smi_ted(nn.Module):
     """materials.smi-ted-Large 738M Parameters"""
 
-    def __init__(self, tokenizer, config=None):
+    def __init__(self, tokenizer, config=None, eval=False):
         super(Smi_ted, self).__init__()
 
         # configuration
@@ -373,11 +373,11 @@ class Smi_ted(nn.Module):
 
         # instantiate modules
         if self.config:
-            self.encoder = MoLEncoder(self.config, self.n_vocab)
+            self.encoder = MoLEncoder(self.config, self.n_vocab, eval=eval)
             self.decoder = MoLDecoder(self.n_vocab, self.config['max_len'], self.config['n_embd'])
             self.net = Net(self.config['n_embd'], n_output=self.config['n_output'], dropout=self.config['dropout'])
     
-    def load_checkpoint(self, ckpt_path, n_output):
+    def load_checkpoint(self, ckpt_path, n_output, eval=False):
         # load checkpoint file
         checkpoint = torch.load(ckpt_path, map_location=torch.device('cpu'))
 
@@ -388,7 +388,7 @@ class Smi_ted(nn.Module):
         self._set_seed(self.config['seed'])
 
         # instantiate modules
-        self.encoder = MoLEncoder(self.config, self.n_vocab)
+        self.encoder = MoLEncoder(self.config, self.n_vocab, eval=eval)
         self.decoder = MoLDecoder(self.n_vocab, self.max_len, self.n_embd)
         self.net = Net(self.n_embd, n_output=self.config['n_output'] if 'n_output' in self.config else n_output, dropout=self.config['dropout'])
 
@@ -493,11 +493,12 @@ class Smi_ted(nn.Module):
 def load_smi_ted(folder="./smi_ted_large", 
               ckpt_filename="smi-ted-Large_30.pt",
               vocab_filename="bert_vocab_curated.txt",
-              n_output=1
+              n_output=1,
+              eval=False
               ):
     tokenizer = MolTranBertTokenizer(os.path.join(folder, vocab_filename))
     model = Smi_ted(tokenizer)
-    model.load_checkpoint(os.path.join(folder, ckpt_filename), n_output)
+    model.load_checkpoint(os.path.join(folder, ckpt_filename), n_output, eval=eval)
     print('Vocab size:', len(tokenizer.vocab))
     print(f'[FINETUNE MODE - {str(model)}]')
     return model

smi-ted/finetune/smi_ted_light/load.py

@@ -318,7 +318,7 @@ class Net(nn.Module):
 
 class MoLEncoder(nn.Module):
 
-    def __init__(self, config, n_vocab):
+    def __init__(self, config, n_vocab, eval=False):
         super(MoLEncoder, self).__init__()
 
         # embeddings
@@ -337,7 +337,7 @@ class MoLEncoder(nn.Module):
             # unless we do deterministic_eval here, we will have random outputs
             feature_map=partial(GeneralizedRandomFeatures, 
                                 n_dims=config['num_feats'], 
-                                deterministic_eval=False),
+                                deterministic_eval=eval),
             activation='gelu'
         )
         self.blocks = builder.get()
@@ -361,7 +361,7 @@ class MoLDecoder(nn.Module):
 class Smi_ted(nn.Module):
     """materials.smi-ted-Light 289M Parameters"""
 
-    def __init__(self, tokenizer, config=None):
+    def __init__(self, tokenizer, config=None, eval=False):
         super(Smi_ted, self).__init__()
 
         # configuration
@@ -373,11 +373,11 @@ class Smi_ted(nn.Module):
 
         # instantiate modules
         if self.config:
-            self.encoder = MoLEncoder(self.config, self.n_vocab)
+            self.encoder = MoLEncoder(self.config, self.n_vocab, eval=eval)
             self.decoder = MoLDecoder(self.n_vocab, self.config['max_len'], self.config['n_embd'])
             self.net = Net(self.config['n_embd'], n_output=self.config['n_output'], dropout=self.config['dropout'])
     
-    def load_checkpoint(self, ckpt_path, n_output):
+    def load_checkpoint(self, ckpt_path, n_output, eval=False):
         # load checkpoint file
         checkpoint = torch.load(ckpt_path, map_location=torch.device('cpu'))
 
@@ -388,7 +388,7 @@ class Smi_ted(nn.Module):
         self._set_seed(self.config['seed'])
 
         # instantiate modules
-        self.encoder = MoLEncoder(self.config, self.n_vocab)
+        self.encoder = MoLEncoder(self.config, self.n_vocab, eval=eval)
         self.decoder = MoLDecoder(self.n_vocab, self.max_len, self.n_embd)
         self.net = Net(self.n_embd, n_output=self.config['n_output'] if 'n_output' in self.config else n_output, dropout=self.config['dropout'])
 
@@ -493,11 +493,12 @@ class Smi_ted(nn.Module):
 def load_smi_ted(folder="./smi_ted_light", 
               ckpt_filename="smi-ted-Light_40.pt",
               vocab_filename="bert_vocab_curated.txt",
-              n_output=1
+              n_output=1,
+              eval=False
               ):
     tokenizer = MolTranBertTokenizer(os.path.join(folder, vocab_filename))
     model = Smi_ted(tokenizer)
-    model.load_checkpoint(os.path.join(folder, ckpt_filename), n_output)
+    model.load_checkpoint(os.path.join(folder, ckpt_filename), n_output, eval=eval)
     print('Vocab size:', len(tokenizer.vocab))
     print(f'[FINETUNE MODE - {str(model)}]')
     return model

smi-ted/finetune/trainers.py

@@ -14,6 +14,7 @@ import numpy as np
 import random
 import args
 import os
+import shutil
 from tqdm import tqdm
 
 # Machine Learning
@@ -25,7 +26,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_every_epoch=False, save_ckpt=True, device='cpu'):
+                 target_metric='rmse', seed=0, smi_ted_version=None, checkpoints_folder='./checkpoints', restart_filename=None, save_every_epoch=False, save_ckpt=True, device='cpu'):
         # data
         self.df_train = raw_data[0]
         self.df_valid = raw_data[1]
@@ -39,10 +40,15 @@ class Trainer:
         # config
         self.target_metric = target_metric
         self.seed = seed
+        self.smi_ted_version = smi_ted_version
         self.checkpoints_folder = checkpoints_folder
+        self.restart_filename = restart_filename
+        self.start_epoch = 1
         self.save_every_epoch = save_every_epoch
         self.save_ckpt = save_ckpt
         self.device = device
+        self.best_vloss = float('inf')
+        self.last_filename = None
         self._set_seed(seed)
 
     def _prepare_data(self):
@@ -80,11 +86,12 @@ class Trainer:
         self.optimizer = optimizer
         self.loss_fn = loss_fn
         self._print_configuration()
+        if self.restart_filename:
+            self._load_checkpoint(self.restart_filename)
+            print('Checkpoint restored!')
 
     def fit(self, max_epochs=500):
-        best_vloss = float('inf')
-
-        for epoch in range(1, max_epochs+1):
+        for epoch in range(self.start_epoch, max_epochs+1):
             print(f'\n=====Epoch [{epoch}/{max_epochs}]=====')
 
             # training
@@ -99,44 +106,68 @@ class Trainer:
                 print(f"[VALID] Evaluation {m.upper()}: {round(val_metrics[m], 4)}")
 
             ############################### Save Finetune checkpoint #######################################
-            if ((val_loss < best_vloss) or self.save_every_epoch) and self.save_ckpt:
+            if ((val_loss < self.best_vloss) or self.save_every_epoch) and self.save_ckpt:
                 # remove old checkpoint
-                if best_vloss != float('inf') and not self.save_every_epoch:
+                if (self.last_filename != None) 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'
-                self.last_filename = f"{model_name}_epoch={epoch}_{self.dataset_name}_seed{self.seed}_valloss={round(val_loss, 4)}.pt"
+                self.last_filename = f"{model_name}_seed{self.seed}_{self.dataset_name}_epoch={epoch}_valloss={round(val_loss, 4)}.pt"
+
+                # update best loss
+                self.best_vloss = val_loss
 
                 # save checkpoint
                 print('Saving checkpoint...')
                 self._save_checkpoint(epoch, self.last_filename)
 
-                # update best loss
-                best_vloss = val_loss
-
-    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 evaluate(self, verbose=True):
+        if verbose:
+            print("\n=====Test Evaluation=====")
+
+        if self.smi_ted_version == 'v1':
+            import smi_ted_light.load as load
+        elif self.smi_ted_version == 'v2':
+            import smi_ted_large.load as load
+        else:
+            raise Exception('Please, specify the SMI-TED version: `v1` or `v2`.')
+
+        # copy vocabulary to checkpoint folder
+        if not os.path.exists(os.path.join(self.checkpoints_folder, 'bert_vocab_curated.txt')):
+            smi_ted_path = os.path.dirname(load.__file__)
+            shutil.copy(os.path.join(smi_ted_path, 'bert_vocab_curated.txt'), self.checkpoints_folder)
+
+        # load model for inference
+        model_inf = load.load_smi_ted(
+            folder=self.checkpoints_folder,
+            ckpt_filename=self.last_filename,
+            eval=True,
+        ).to(self.device)
+
+        # set model evaluation mode
+        model_inf.eval()
+
+        # evaluate on test set
+        tst_preds, tst_loss, tst_metrics = self._validate_one_epoch(self.test_loader, model_inf)
+
+        if verbose:
+            # 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
 
-    def _validate_one_epoch(self, data_loader):
+    def _validate_one_epoch(self, data_loader, model=None):
         raise NotImplementedError
 
     def _print_configuration(self):
@@ -157,6 +188,8 @@ class Trainer:
         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'])
+        self.start_epoch = ckpt_dict['EPOCHS_RUN'] + 1
+        self.best_vloss = ckpt_dict['finetune_info']['best_vloss']
 
     def _save_checkpoint(self, current_epoch, filename):
         if not os.path.exists(self.checkpoints_folder):
@@ -177,6 +210,7 @@ class Trainer:
                 'train_size': self.df_train.shape[0],
                 'valid_size': self.df_valid.shape[0],
                 'test_size': self.df_test.shape[0],
+                'best_vloss': self.best_vloss,
             },
             'seed': self.seed,
         }
@@ -203,9 +237,9 @@ 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_every_epoch=False, save_ckpt=True, device='cpu'):
+                 target_metric='rmse', seed=0, smi_ted_version=None, checkpoints_folder='./checkpoints', restart_filename=None, 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_every_epoch, save_ckpt, device) 
+                         target_metric, seed, smi_ted_version, checkpoints_folder, restart_filename, save_every_epoch, save_ckpt, device) 
 
     def _train_one_epoch(self):
         running_loss = 0.0
@@ -239,11 +273,13 @@ class TrainerRegressor(Trainer):
 
         return running_loss / len(self.train_loader)
 
-    def _validate_one_epoch(self, data_loader):
+    def _validate_one_epoch(self, data_loader, model=None):
         data_targets = []
         data_preds = []
         running_loss = 0.0
 
+        model = self.model if model is None else model
+
         with torch.no_grad():
             for idx, data in enumerate(pbar := tqdm(data_loader)):
                 # Every data instance is an input + label pair
@@ -251,8 +287,8 @@ class TrainerRegressor(Trainer):
                 targets = targets.clone().detach().to(self.device)
 
                 # Make predictions for this batch
-                embeddings = self.model.extract_embeddings(smiles).to(self.device)
-                predictions = self.model.net(embeddings).squeeze()
+                embeddings = model.extract_embeddings(smiles).to(self.device)
+                predictions = model.net(embeddings).squeeze()
 
                 # Compute the loss
                 loss = self.loss_fn(predictions, targets)
@@ -292,9 +328,9 @@ class TrainerRegressor(Trainer):
 class TrainerClassifier(Trainer):
 
     def __init__(self, raw_data, dataset_name, target, batch_size, hparams,
-                 target_metric='roc-auc', seed=0, checkpoints_folder='./checkpoints', save_every_epoch=False, save_ckpt=True, device='cpu'):
+                 target_metric='roc-auc', seed=0, smi_ted_version=None, checkpoints_folder='./checkpoints', restart_filename=None, 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_every_epoch, save_ckpt, device) 
+                         target_metric, seed, smi_ted_version, checkpoints_folder, restart_filename, save_every_epoch, save_ckpt, device) 
 
     def _train_one_epoch(self):
         running_loss = 0.0
@@ -328,11 +364,13 @@ class TrainerClassifier(Trainer):
 
         return running_loss / len(self.train_loader)
 
-    def _validate_one_epoch(self, data_loader):
+    def _validate_one_epoch(self, data_loader, model=None):
         data_targets = []
         data_preds = []
         running_loss = 0.0
 
+        model = self.model if model is None else model
+
         with torch.no_grad():
             for idx, data in enumerate(pbar := tqdm(data_loader)):
                 # Every data instance is an input + label pair
@@ -340,8 +378,8 @@ class TrainerClassifier(Trainer):
                 targets = targets.clone().detach().to(self.device)
 
                 # Make predictions for this batch
-                embeddings = self.model.extract_embeddings(smiles).to(self.device)
-                predictions = self.model.net(embeddings).squeeze()
+                embeddings = model.extract_embeddings(smiles).to(self.device)
+                predictions = model.net(embeddings).squeeze()
 
                 # Compute the loss
                 loss = self.loss_fn(predictions, targets.long())
@@ -397,9 +435,9 @@ class TrainerClassifier(Trainer):
 class TrainerClassifierMultitask(Trainer):
 
     def __init__(self, raw_data, dataset_name, target, batch_size, hparams,
-                 target_metric='roc-auc', seed=0, checkpoints_folder='./checkpoints', save_every_epoch=False, save_ckpt=True, device='cpu'):
+                 target_metric='roc-auc', seed=0, smi_ted_version=None, checkpoints_folder='./checkpoints', restart_filename=None, 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_every_epoch, save_ckpt, device)
+                         target_metric, seed, smi_ted_version, checkpoints_folder, restart_filename, save_every_epoch, save_ckpt, device)
 
     def _prepare_data(self):
         # normalize dataset
@@ -464,12 +502,14 @@ class TrainerClassifierMultitask(Trainer):
 
         return running_loss / len(self.train_loader)
 
-    def _validate_one_epoch(self, data_loader):
+    def _validate_one_epoch(self, data_loader, model=None):
         data_targets = []
         data_preds = []
         data_masks = []
         running_loss = 0.0
 
+        model = self.model if model is None else model
+
         with torch.no_grad():
             for idx, data in enumerate(pbar := tqdm(data_loader)):
                 # Every data instance is an input + label pair + mask
@@ -477,8 +517,8 @@ class TrainerClassifierMultitask(Trainer):
                 targets = targets.clone().detach().to(self.device)
 
                 # Make predictions for this batch
-                embeddings = self.model.extract_embeddings(smiles).to(self.device)
-                predictions = self.model.net(embeddings, multitask=True).squeeze()
+                embeddings = model.extract_embeddings(smiles).to(self.device)
+                predictions = model.net(embeddings, multitask=True).squeeze()
                 predictions = predictions * target_masks.to(self.device)
 
                 # Compute the loss
@@ -548,4 +588,4 @@ class TrainerClassifierMultitask(Trainer):
             'specificity': average_sp.item(),
         }
 
-        return preds, running_loss / len(data_loader), metrics
+        return preds.cpu().numpy(), running_loss / len(data_loader), metrics