Added LR scheduler + set default sum of embeddings

protac_degradation_predictor/config.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass, field
 class Config:
     # Embeddings information
     morgan_radius: int = 15
-    fingerprint_size: int = 224
+    fingerprint_size: int = 256 # 224
     protein_embedding_size: int = 1024
     cell_embedding_size: int = 768
 

protac_degradation_predictor/optuna_utils.py

@@ -118,7 +118,6 @@ def pytorch_model_objective(
     hidden_dim = trial.suggest_categorical('hidden_dim', hidden_dim_options)
     batch_size = trial.suggest_categorical('batch_size', batch_size_options)
     learning_rate = trial.suggest_float('learning_rate', *learning_rate_options, log=True)
-    join_embeddings = trial.suggest_categorical('join_embeddings', ['beginning', 'concat', 'sum'])
     smote_k_neighbors = trial.suggest_categorical('smote_k_neighbors', smote_k_neighbors_options)
     use_smote = trial.suggest_categorical('use_smote', [True, False])
     apply_scaling = trial.suggest_categorical('apply_scaling', [True, False])
@@ -161,7 +160,6 @@ def pytorch_model_objective(
             test_df=test_df,
             hidden_dim=hidden_dim,
             batch_size=batch_size,
-            join_embeddings=join_embeddings,
             learning_rate=learning_rate,
             dropout=dropout,
             max_epochs=max_epochs,
@@ -177,7 +175,6 @@ def pytorch_model_objective(
         if test_df is not None:
             _, trainer, metrics, val_pred, test_pred = ret
             test_preds.append(test_pred)
-            logging.info(f'Test predictions: {test_pred}')
         else:
             _, trainer, metrics, val_pred = ret
         train_metrics = {m: v.item() for m, v in trainer.callback_metrics.items() if 'train' in m}
@@ -190,7 +187,7 @@ def pytorch_model_objective(
     trial.set_user_attr('report', report)
 
     # Get the majority vote for the test predictions
-    if test_df is not None:
+    if test_df is not None and not fast_dev_run:
         # Get the majority vote for the test predictions
         test_preds = torch.stack(test_preds)
         test_preds, _ = torch.mode(test_preds, dim=0)
@@ -340,27 +337,28 @@ def hyperparameter_tuning_and_training(
     test_report = pd.DataFrame(test_report)
 
     # Get the majority vote for the test predictions
-    test_preds = torch.stack(test_preds)
-    test_preds, _ = torch.mode(test_preds, dim=0)
-    y = torch.tensor(test_df[active_label].tolist())
-    # Measure the test accuracy and ROC AUC
-    majority_vote_metrics = {
-        'cv_models': False,
-        'test_acc': Accuracy(task='binary')(test_preds, y).item(),
-        'test_roc_auc': AUROC(task='binary')(test_preds, y).item(),
-        'test_precision': Precision(task='binary')(test_preds, y).item(),
-        'test_recall': Recall(task='binary')(test_preds, y).item(),
-        'test_f1': F1Score(task='binary')(test_preds, y).item(),
-    }
-    majority_vote_metrics.update(get_dataframe_stats(train_val_df, test_df=test_df, active_label=active_label))
-    majority_vote_metrics_cv = study.best_trial.user_attrs['majority_vote_metrics']
-    majority_vote_metrics_cv['cv_models'] = True
-    majority_vote_report = pd.DataFrame([
-        majority_vote_metrics,
-        majority_vote_metrics_cv,
-    ])
-    majority_vote_report['model_type'] = 'Pytorch'
-    majority_vote_report['split_type'] = split_type
+    if not fast_dev_run:
+        test_preds = torch.stack(test_preds)
+        test_preds, _ = torch.mode(test_preds, dim=0)
+        y = torch.tensor(test_df[active_label].tolist())
+        # Measure the test accuracy and ROC AUC
+        majority_vote_metrics = {
+            'cv_models': False,
+            'test_acc': Accuracy(task='binary')(test_preds, y).item(),
+            'test_roc_auc': AUROC(task='binary')(test_preds, y).item(),
+            'test_precision': Precision(task='binary')(test_preds, y).item(),
+            'test_recall': Recall(task='binary')(test_preds, y).item(),
+            'test_f1': F1Score(task='binary')(test_preds, y).item(),
+        }
+        majority_vote_metrics.update(get_dataframe_stats(train_val_df, test_df=test_df, active_label=active_label))
+        majority_vote_metrics_cv = study.best_trial.user_attrs['majority_vote_metrics']
+        majority_vote_metrics_cv['cv_models'] = True
+        majority_vote_report = pd.DataFrame([
+            majority_vote_metrics,
+            majority_vote_metrics_cv,
+        ])
+        majority_vote_report['model_type'] = 'Pytorch'
+        majority_vote_report['split_type'] = split_type
 
     # Ablation study: disable embeddings at a time
     ablation_report = []
@@ -407,8 +405,9 @@ def hyperparameter_tuning_and_training(
         'hparam_report': hparam_report,
         'test_report': test_report,
         'ablation_report': ablation_report,
-        'majority_vote_report': majority_vote_report,
     }
+    if not fast_dev_run:
+        ret['majority_vote_report'] = majority_vote_report
     return ret
 
 

protac_degradation_predictor/pytorch_models.py

@@ -36,7 +36,7 @@ class PROTAC_Predictor(nn.Module):
         e3_emb_dim: int = config.protein_embedding_size,
         cell_emb_dim: int = config.cell_embedding_size,
         dropout: float = 0.2,
-        join_embeddings: Literal['beginning', 'concat', 'sum'] = 'concat',
+        join_embeddings: Literal['beginning', 'concat', 'sum'] = 'sum',
         disabled_embeddings: list = [],
     ):
         """ Initialize the PROTAC model.
@@ -140,7 +140,7 @@ class PROTAC_Model(pl.LightningModule):
         batch_size: int = 32,
         learning_rate: float = 1e-3,
         dropout: float = 0.2,
-        join_embeddings: Literal['beginning', 'concat', 'sum'] = 'concat',
+        join_embeddings: Literal['beginning', 'concat', 'sum'] = 'sum',
         train_dataset: PROTAC_Dataset = None,
         val_dataset: PROTAC_Dataset = None,
         test_dataset: PROTAC_Dataset = None,
@@ -308,7 +308,19 @@ class PROTAC_Model(pl.LightningModule):
         return self.step(batch, batch_idx, 'test')
 
     def configure_optimizers(self):
-        return optim.Adam(self.parameters(), lr=self.learning_rate)
+        optimizer = optim.Adam(self.parameters(), lr=self.learning_rate)
+        return {
+            'optimizer': optimizer,
+            'lr_scheduler': optim.lr_scheduler.ReduceLROnPlateau(
+                optimizer=optimizer,
+                mode='min',
+                factor=0.5,
+                patience=2,
+            ),
+            'interval': 'step',  # or 'epoch'
+            'frequency': 1,
+            'monitor': 'val_loss',
+        }
 
     def predict_step(self, batch, batch_idx):
         poi_emb = batch['poi_emb']
@@ -384,7 +396,7 @@ def train_model(
         poi_emb_dim: int = config.protein_embedding_size,
         e3_emb_dim: int = config.protein_embedding_size,
         cell_emb_dim: int = config.cell_embedding_size,
-        join_embeddings: Literal['beginning', 'concat', 'sum'] = 'concat',
+        join_embeddings: Literal['beginning', 'concat', 'sum'] = 'sum',
         smote_k_neighbors:int = 5,
         use_smote: bool = True,
         apply_scaling: bool = False,
@@ -482,6 +494,8 @@ def train_model(
             verbose=False,
         ),
     ]
+    if use_logger:
+        callbacks.append(pl.callbacks.LearningRateMonitor(logging_interval='step'))
     if enable_checkpointing:
         callbacks.append(pl.callbacks.ModelCheckpoint(
             monitor='val_acc',

src/run_experiments.py

@@ -309,15 +309,8 @@ def main(
 
         # Save the reports to file
         for report_name, report in optuna_reports.items():
-            report.to_csv(f'../reports/report_{report_name}_{experiment_name}.csv', index=False)
+            report.to_csv(f'../reports/{report_name}_{experiment_name}.csv', index=False)
             reports[report_name].append(report.copy())
-    
-    # Save the reports to file after concatenating them
-    for report_name, report in reports.items():
-        report = pd.concat(report)
-        report.to_csv(f'../reports/report_{report_name}_{active_name}_test_split_{test_split}.csv', index=False)
-
-
 
         # # Start the CV over the folds
         # X = train_val_df.drop(columns=active_col)