scaler for cross validation

Author: ahinterl94-th

src/spotPython/data/lightcrossvalidationdatamodule.py

@@ -2,7 +2,9 @@
 from torch.utils.data import DataLoader, Subset
 from typing import Optional
 from sklearn.model_selection import KFold
-from torch.utils.data import Dataset
+from torch.utils.data import Dataset, TensorDataset
+import torch
+
 
 
 class LightCrossValidationDataModule(L.LightningDataModule):
@@ -44,6 +46,7 @@ def __init__(
         data_dir: str = "./data",
         num_workers: int = 0,
         pin_memory: bool = False,
+        scaler: Optional[object] = None,
     ):
         super().__init__()
         self.batch_size = batch_size
@@ -54,6 +57,7 @@ def __init__(
         self.split_seed = split_seed
         self.num_splits = num_splits
         self.pin_memory = pin_memory
+        self.scaler = scaler
         self.save_hyperparameters(logger=False)
         assert 0 <= self.k < self.num_splits, "incorrect fold number"
 
@@ -85,6 +89,21 @@ def setup(self, stage: Optional[str] = None) -> None:
             print(f"Train Dataset Size: {len(self.data_train)}")
             self.data_val = Subset(dataset_full, val_indexes)
             print(f"Val Dataset Size: {len(self.data_val)}")
+        
+        if self.scaler is not None:
+            # Fit the scaler on training data and transform both train and val data
+            scaler_train_data = torch.stack([self.data_train[i][0] for i in range(len(self.data_train))]).squeeze(1)
+            self.scaler.fit(scaler_train_data)
+            self.data_train = [(self.scaler.transform(data), target) for data, target in self.data_train]
+            data_tensors_train = [data.clone().detach() for data, target in self.data_train]
+            target_tensors_train = [target.clone().detach() for data, target in self.data_train]
+            self.data_train = TensorDataset(
+                torch.stack(data_tensors_train).squeeze(1), torch.stack(target_tensors_train)
+            )
+            self.data_val = [(self.scaler.transform(data), target) for data, target in self.data_val]
+            data_tensors_val = [data.clone().detach() for data, target in self.data_val]
+            target_tensors_val = [target.clone().detach() for data, target in self.data_val]
+            self.data_val = TensorDataset(torch.stack(data_tensors_val).squeeze(1), torch.stack(target_tensors_val))
 
     def train_dataloader(self) -> DataLoader:
         """

src/spotPython/light/cvmodel.py

@@ -71,7 +71,9 @@ def cv_model(config: dict, fun_control: dict) -> float:
             num_workers=fun_control["num_workers"],
             batch_size=config["batch_size"],
             data_dir=fun_control["DATASET_PATH"],
+            scaler=fun_control["scaler"],
         )
+        dm.setup()
         dm.prepare_data()
 
         # TODO: Check if this is necessary: