Added model tests + checkpointing of the scaler object

protac_degradation_predictor/__init__.py

@@ -5,6 +5,8 @@ from .data_utils import (
     is_active,
 )
 from .pytorch_models import (
+    PROTAC_Predictor,
+    PROTAC_Model,
     train_model,
 )
 from .sklearn_models import (

protac_degradation_predictor/optuna_utils.py

@@ -73,7 +73,7 @@ def pytorch_model_objective(
     dropout = trial.suggest_float('dropout', *dropout_options)
 
     # Start the CV over the folds
-    X = train_val_df.drop(columns=active_label)
+    X = train_val_df.copy().drop(columns=active_label)
     y = train_val_df[active_label].tolist()
     report = []
     for k, (train_index, val_index) in enumerate(kf.split(X, y, groups)):
@@ -108,11 +108,11 @@ def pytorch_model_objective(
         # At each fold, train and evaluate the Pytorch model
         # Train the model with the current set of hyperparameters
         _, _, metrics = train_model(
-            protein2embedding,
-            cell2embedding,
-            smiles2fp,
-            train_df,
-            val_df,
+            protein2embedding=protein2embedding,
+            cell2embedding=cell2embedding,
+            smiles2fp=smiles2fp,
+            train_df=train_df,
+            val_df=val_df,
             hidden_dim=hidden_dim,
             batch_size=batch_size,
             join_embeddings=join_embeddings,
@@ -223,7 +223,7 @@ def hyperparameter_tuning_and_training(
     test_report = []
     # Retrain N models with the best hyperparameters (measure model uncertainty)
     for i in range(n_models_for_test):
-        pl.seed_everything(42 + i)
+        pl.seed_everything(42 + i + 1)
         _, _, metrics = train_model(
             protein2embedding=protein2embedding,
             cell2embedding=cell2embedding,
@@ -235,9 +235,9 @@ def hyperparameter_tuning_and_training(
             active_label=active_label,
             max_epochs=max_epochs,
             disabled_embeddings=[],
-            logger_name=f'{logger_name}_best_model_{i}',
+            logger_name=f'{logger_name}_best_model_n{i}',
             enable_checkpointing=True,
-            checkpoint_model_name=f'best_model_{split_type}_{i}',
+            checkpoint_model_name=f'best_model_n{i}_{split_type}',
             **study.best_params,
         )
         # Rename the keys in the metrics dictionary
@@ -245,6 +245,9 @@ def hyperparameter_tuning_and_training(
         metrics = {k.replace('train_', 'train_val_'): v for k, v in metrics.items()}
         metrics['model_type'] = 'Pytorch'
         metrics['test_model_id'] = i
+        metrics['test_len'] = len(test_df)
+        metrics['test_active_perc'] = test_df[active_label].sum() / len(test_df)
+        metrics['test_inactive_perc'] = (len(test_df) - test_df[active_label].sum()) / len(test_df)
         test_report.append(metrics.copy())
     test_report = pd.DataFrame(test_report)
 

protac_degradation_predictor/pytorch_models.py

@@ -1,4 +1,6 @@
 import warnings
+import pickle
+import logging
 from typing import Literal, List, Tuple, Optional, Dict
 
 from .protac_dataset import PROTAC_Dataset
@@ -125,7 +127,6 @@ class PROTAC_Predictor(nn.Module):
         return x
 
 
-
 class PROTAC_Model(pl.LightningModule):
 
     def __init__(
@@ -218,13 +219,26 @@ class PROTAC_Model(pl.LightningModule):
             '''
         
         # Apply scaling in datasets
-        if self.apply_scaling:
-            use_single_scaler = True if self.join_embeddings == 'beginning' else False
+        self.scalers = None
+        if self.apply_scaling and self.train_dataset is not None:
+            self.initialize_scalers()
+
+    def initialize_scalers(self):
+        """Initialize or reinitialize scalers based on dataset properties."""
+        if self.scalers is None:
+            use_single_scaler = self.join_embeddings == 'beginning'
             self.scalers = self.train_dataset.fit_scaling(use_single_scaler)
+            self.apply_scalers()
+
+    def apply_scalers(self):
+        """Apply scalers to all datasets."""
+        use_single_scaler = self.join_embeddings == 'beginning'
+        if self.train_dataset:
             self.train_dataset.apply_scaling(self.scalers, use_single_scaler)
+        if self.val_dataset:
             self.val_dataset.apply_scaling(self.scalers, use_single_scaler)
-            if self.test_dataset:
-                self.test_dataset.apply_scaling(self.scalers, use_single_scaler)
+        if self.test_dataset:
+            self.test_dataset.apply_scaling(self.scalers, use_single_scaler)
 
     def forward(self, poi_emb, e3_emb, cell_emb, smiles_emb):
         return self.model(poi_emb, e3_emb, cell_emb, smiles_emb)
@@ -316,6 +330,23 @@ class PROTAC_Model(pl.LightningModule):
             batch_size=self.batch_size,
             shuffle=False,
         )
+    
+    def on_save_checkpoint(self, checkpoint):
+        """ Serialize the scalers to the checkpoint. """
+        checkpoint['scalers'] = pickle.dumps(self.scalers)
+    
+    def on_load_checkpoint(self, checkpoint):
+        """Deserialize the scalers from the checkpoint."""
+        if 'scalers' in checkpoint:
+            self.scalers = pickle.loads(checkpoint['scalers'])
+        else:
+            self.scalers = None
+        if self.apply_scaling:
+            if self.scalers is not None:
+                # Re-apply scalers to ensure datasets are scaled
+                self.apply_scalers()
+            else:
+                logging.warning("Scalers not found in checkpoint. Consider re-fitting scalers if necessary.")
 
 
 def train_model(
@@ -421,7 +452,7 @@ def train_model(
             monitor='val_acc',
             mode='max',
             verbose=False,
-            filename=checkpoint_model_name + '-{epoch}-{val_metrics_opt_score:.4f}',
+            filename=checkpoint_model_name + '-{epoch}-{val_acc:.2f}-{val_roc_auc:.3f}',
         ))
     # Define Trainer
     trainer = pl.Trainer(
@@ -455,6 +486,9 @@ def train_model(
         warnings.simplefilter("ignore")
         trainer.fit(model)
     metrics = trainer.validate(model, verbose=False)[0]
+
+    # Add train metrics to metrics
+
     if test_df is not None:
         test_metrics = trainer.test(model, verbose=False)[0]
         metrics.update(test_metrics)
@@ -472,6 +506,15 @@ def load_model(
     Returns:
         PROTAC_Model: The loaded model.
     """
-    model = PROTAC_Model.load_from_checkpoint(ckpt_path)
+    # NOTE: The `map_locat` argument is automatically handled in newer versions
+    # of PyTorch Lightning, but we keep it here for compatibility with older ones.
+    model = PROTAC_Model.load_from_checkpoint(
+        ckpt_path,
+        map_location=torch.device('cpu') if not torch.cuda.is_available() else None,
+    )
+    # NOTE: The following is left as example for eventually re-applying scaling
+    # with other datasets...
+    # if model.apply_scaling:
+    #     model.apply_scalers()
     model.eval()
     return model

src/run_experiments.py

@@ -207,10 +207,11 @@ def get_target_split_indices(active_df: pd.DataFrame, active_col: str, test_spli
 
 def main(
     active_col: str = 'Active (Dmax 0.6, pDC50 6.0)',
-    n_trials: int = 50,
+    n_trials: int = 100,
     fast_dev_run: bool = False,
-    test_split: float = 0.2,
+    test_split: float = 0.1,
     cv_n_splits: int = 5,
+    max_epochs: int = 100,
     run_sklearn: bool = False,
 ):
     """ Train a PROTAC model using the given datasets and hyperparameters.
@@ -287,7 +288,7 @@ def main(
             n_models_for_test=3,
             fast_dev_run=fast_dev_run,
             n_trials=n_trials,
-            max_epochs=10,
+            max_epochs=max_epochs,
             logger_name=f'logs_{experiment_name}',
             active_label=active_col,
             study_filename=f'../reports/study_{experiment_name}.pkl',

tests/test_pytorch_model.py

@@ -0,0 +1,80 @@
+import pytest
+import os
+import sys
+import logging
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from protac_degradation_predictor import PROTAC_Model, PROTAC_Predictor
+
+import torch
+
+
+def test_protac_model():
+    model = PROTAC_Model(hidden_dim=128)
+    assert model.hidden_dim == 128
+    assert model.smiles_emb_dim == 224
+    assert model.poi_emb_dim == 1024
+    assert model.e3_emb_dim == 1024
+    assert model.cell_emb_dim == 768
+    assert model.batch_size == 32
+    assert model.learning_rate == 0.001
+    assert model.dropout == 0.2
+    assert model.join_embeddings == 'concat'
+    assert model.train_dataset is None
+    assert model.val_dataset is None
+    assert model.test_dataset is None
+    assert model.disabled_embeddings == []
+    assert model.apply_scaling == False
+
+def test_protac_predictor():
+    predictor = PROTAC_Predictor(hidden_dim=128)
+    assert predictor.hidden_dim == 128
+    assert predictor.smiles_emb_dim == 224
+    assert predictor.poi_emb_dim == 1024
+    assert predictor.e3_emb_dim == 1024
+    assert predictor.cell_emb_dim == 768
+    assert predictor.join_embeddings == 'concat'
+    assert predictor.disabled_embeddings == []
+
+def test_load_model(caplog):
+    caplog.set_level(logging.WARNING)
+
+    model = PROTAC_Model.load_from_checkpoint(
+        'data/test_model.ckpt',
+        map_location=torch.device("cpu") if not torch.cuda.is_available() else None,
+    )
+    # apply_scaling: true
+    # batch_size: 8
+    # cell_emb_dim: 768
+    # disabled_embeddings: []
+    # dropout: 0.1498104322091649
+    # e3_emb_dim: 1024
+    # hidden_dim: 768
+    # join_embeddings: concat
+    # learning_rate: 4.881387978425994e-05
+    # poi_emb_dim: 1024
+    # smiles_emb_dim: 224
+    assert model.hidden_dim == 768
+    assert model.smiles_emb_dim == 224
+    assert model.poi_emb_dim == 1024
+    assert model.e3_emb_dim == 1024
+    assert model.cell_emb_dim == 768
+    assert model.batch_size == 8
+    assert model.learning_rate == 4.881387978425994e-05
+    assert model.dropout == 0.1498104322091649
+    assert model.join_embeddings == 'concat'
+    assert model.disabled_embeddings == []
+    assert model.apply_scaling == True
+
+
+def test_checkpoint_file():
+    checkpoint = torch.load(
+        'data/test_model.ckpt',
+        map_location=torch.device("cpu") if not torch.cuda.is_available() else None,
+    )
+    print(checkpoint.keys())
+    print(checkpoint["hyper_parameters"])
+    print([k for k, v in checkpoint["state_dict"].items()])
+
+pytest.main()