Add PyTorch Lightning integration (#203)

This pull request is to merge a new experiment
`torch-lightning-experiment` for the issue `195`. The new class
`TorchExperiment` is an experiment adapter for PyTorch Lightning
experiments and is used to perform experiments using PyTorch Lightning
modules. It allows for hyperparameter tuning and evaluation of the
model's performance using specified metrics.

The `TorchExperiment` class accepts a `LightningModule`, `DataModule`,
`Trainer`, and an `objective_matric` with default value `val_loss`. The
`_evaluate` function internally performs a training run and returns a
score of the `objective_matric`.

Fixes #195

Author: amitsubhashchejara

pyproject.toml

@@ -74,6 +74,7 @@ test_parallel_backends = [
 all_extras = [
   "hyperactive[integrations]",
   "optuna<5",
+  "lightning",
 ]
 
 

src/hyperactive/experiment/integrations/__init__.py

@@ -11,10 +11,14 @@
 from hyperactive.experiment.integrations.sktime_forecasting import (
     SktimeForecastingExperiment,
 )
+from hyperactive.experiment.integrations.torch_lightning_experiment import (
+    TorchExperiment,
+)
 
 __all__ = [
     "SklearnCvExperiment",
     "SkproProbaRegExperiment",
     "SktimeClassificationExperiment",
     "SktimeForecastingExperiment",
+    "TorchExperiment",
 ]

src/hyperactive/experiment/integrations/torch_lightning_experiment.py

@@ -0,0 +1,373 @@
+"""Experiment adapter for PyTorch Lightning experiments."""
+
+# copyright: hyperactive developers, MIT License (see LICENSE file)
+
+__author__ = ["amitsubhashchejara"]
+
+import numpy as np
+
+from hyperactive.base import BaseExperiment
+
+
+class TorchExperiment(BaseExperiment):
+    """Experiment adapter for PyTorch Lightning experiments.
+
+    This class is used to perform experiments using PyTorch Lightning modules.
+    It allows for hyperparameter tuning and evaluation of the model's performance
+    using specified metrics.
+
+    The experiment trains a Lightning module with given hyperparameters and returns
+    the validation metric value for optimization.
+
+    Parameters
+    ----------
+    datamodule : L.LightningDataModule
+        A PyTorch Lightning DataModule that handles data loading and preparation.
+    lightning_module : type
+        A PyTorch Lightning Module class (not an instance) that will be instantiated
+        with hyperparameters during optimization.
+    trainer_kwargs : dict, optional (default=None)
+        A dictionary of keyword arguments to pass to the PyTorch Lightning Trainer.
+    objective_metric : str, optional (default='val_loss')
+        The metric used to evaluate the model's performance. This should correspond
+        to a metric logged in the LightningModule during validation.
+
+    Examples
+    --------
+    >>> from hyperactive.experiment.integrations import TorchExperiment
+    >>> import torch
+    >>> import lightning as L
+    >>> from torch import nn
+    >>> from torch.utils.data import DataLoader
+    >>>
+    >>> # Define a simple Lightning Module
+    >>> class SimpleLightningModule(L.LightningModule):
+    ...     def __init__(self, input_dim=10, hidden_dim=16, lr=1e-3):
+    ...         super().__init__()
+    ...         self.save_hyperparameters()
+    ...         self.model = nn.Sequential(
+    ...             nn.Linear(input_dim, hidden_dim),
+    ...             nn.ReLU(),
+    ...             nn.Linear(hidden_dim, 2)
+    ...         )
+    ...         self.lr = lr
+    ...
+    ...     def forward(self, x):
+    ...         return self.model(x)
+    ...
+    ...     def training_step(self, batch, batch_idx):
+    ...         x, y = batch
+    ...         y_hat = self(x)
+    ...         loss = nn.functional.cross_entropy(y_hat, y)
+    ...         self.log("train_loss", loss)
+    ...         return loss
+    ...
+    ...     def validation_step(self, batch, batch_idx):
+    ...         x, y = batch
+    ...         y_hat = self(x)
+    ...         val_loss = nn.functional.cross_entropy(y_hat, y)
+    ...         self.log("val_loss", val_loss, on_epoch=True)
+    ...         return val_loss
+    ...
+    ...     def configure_optimizers(self):
+    ...         return torch.optim.Adam(self.parameters(), lr=self.lr)
+    >>>
+    >>> # Create DataModule
+    >>> class RandomDataModule(L.LightningDataModule):
+    ...     def __init__(self, batch_size=32):
+    ...         super().__init__()
+    ...         self.batch_size = batch_size
+    ...
+    ...     def setup(self, stage=None):
+    ...         dataset = torch.utils.data.TensorDataset(
+    ...             torch.randn(100, 10),
+    ...             torch.randint(0, 2, (100,))
+    ...         )
+    ...         self.train, self.val = torch.utils.data.random_split(
+    ...             dataset, [80, 20]
+    ...         )
+    ...
+    ...     def train_dataloader(self):
+    ...         return DataLoader(self.train, batch_size=self.batch_size)
+    ...
+    ...     def val_dataloader(self):
+    ...         return DataLoader(self.val, batch_size=self.batch_size)
+    >>>
+    >>> datamodule = RandomDataModule(batch_size=16)
+    >>> datamodule.setup()
+    >>>
+    >>> # Create Experiment
+    >>> experiment = TorchExperiment(
+    ...     datamodule=datamodule,
+    ...     lightning_module=SimpleLightningModule,
+    ...     trainer_kwargs={'max_epochs': 3},
+    ...     objective_metric="val_loss"
+    ... )
+    >>>
+    >>> params = {"input_dim": 10, "hidden_dim": 16, "lr": 1e-3}
+    >>>
+    >>> val_result, metadata = experiment._evaluate(params)
+    """
+
+    _tags = {
+        "property:randomness": "random",
+        "property:higher_or_lower_is_better": "lower",
+        "authors": ["amitsubhashchejara"],
+        "python_dependencies": ["torch", "lightning"],
+    }
+
+    def __init__(
+        self,
+        datamodule,
+        lightning_module,
+        trainer_kwargs=None,
+        objective_metric: str = "val_loss",
+    ):
+        self.datamodule = datamodule
+        self.lightning_module = lightning_module
+        self.trainer_kwargs = trainer_kwargs or {}
+        self.objective_metric = objective_metric
+
+        super().__init__()
+
+        self._trainer_kwargs = {
+            "max_epochs": 10,
+            "enable_checkpointing": False,
+            "logger": False,
+            "enable_progress_bar": False,
+            "enable_model_summary": False,
+        }
+        if trainer_kwargs is not None:
+            self._trainer_kwargs.update(trainer_kwargs)
+
+    def _paramnames(self):
+        """Return the parameter names of the search.
+
+        Returns
+        -------
+        list of str, or None
+            The parameter names of the search parameters.
+            If not known or arbitrary, return None.
+        """
+        import inspect
+
+        sig = inspect.signature(self.lightning_module.__init__)
+        return [p for p in sig.parameters.keys() if p != "self"]
+
+    def _evaluate(self, params):
+        """Evaluate the parameters.
+
+        Parameters
+        ----------
+        params : dict with string keys
+            Parameters to evaluate.
+
+        Returns
+        -------
+        float
+            The value of the parameters as per evaluation.
+        dict
+            Additional metadata about the search.
+        """
+        import lightning as L
+
+        try:
+            model = self.lightning_module(**params)
+            trainer = L.Trainer(**self._trainer_kwargs)
+            trainer.fit(model, self.datamodule)
+
+            val_result = trainer.callback_metrics.get(self.objective_metric)
+            metadata = {}
+
+            if val_result is None:
+                available_metrics = list(trainer.callback_metrics.keys())
+                raise ValueError(
+                    f"Metric '{self.objective_metric}' not found. "
+                    f"Available: {available_metrics}"
+                )
+            if hasattr(val_result, "item"):
+                val_result = np.float64(val_result.detach().cpu().item())
+            elif isinstance(val_result, (int, float)):
+                val_result = np.float64(val_result)
+            else:
+                val_result = np.float64(float(val_result))
+
+            return val_result, metadata
+
+        except Exception as e:
+            print(f"Training failed with params {params}: {e}")
+            return np.float64(float("inf")), {}
+
+    @classmethod
+    def get_test_params(cls, parameter_set="default"):
+        """Return testing parameter settings for the estimator.
+
+        Parameters
+        ----------
+        parameter_set : str, default="default"
+            Name of the set of test parameters to return, for use in tests.
+
+        Returns
+        -------
+        params : dict or list of dict, default = {}
+            Parameters to create testing instances of the class.
+        """
+        import lightning as L
+        import torch
+        from torch import nn
+        from torch.utils.data import DataLoader
+
+        class SimpleLightningModule(L.LightningModule):
+            def __init__(self, input_dim=10, hidden_dim=16, lr=1e-3):
+                super().__init__()
+                self.save_hyperparameters()
+                self.model = nn.Sequential(
+                    nn.Linear(input_dim, hidden_dim),
+                    nn.ReLU(),
+                    nn.Linear(hidden_dim, 2),
+                )
+                self.lr = lr
+
+            def forward(self, x):
+                return self.model(x)
+
+            def training_step(self, batch, batch_idx):
+                x, y = batch
+                y_hat = self(x)
+                loss = nn.functional.cross_entropy(y_hat, y)
+                self.log("train_loss", loss)
+                return loss
+
+            def validation_step(self, batch, batch_idx):
+                x, y = batch
+                y_hat = self(x)
+                val_loss = nn.functional.cross_entropy(y_hat, y)
+                self.log("val_loss", val_loss, on_epoch=True)
+                return val_loss
+
+            def configure_optimizers(self):
+                return torch.optim.Adam(self.parameters(), lr=self.lr)
+
+        class RandomDataModule(L.LightningDataModule):
+            def __init__(self, batch_size=32):
+                super().__init__()
+                self.batch_size = batch_size
+
+            def setup(self, stage=None):
+                dataset = torch.utils.data.TensorDataset(
+                    torch.randn(100, 10), torch.randint(0, 2, (100,))
+                )
+                self.train, self.val = torch.utils.data.random_split(dataset, [80, 20])
+
+            def train_dataloader(self):
+                return DataLoader(self.train, batch_size=self.batch_size)
+
+            def val_dataloader(self):
+                return DataLoader(self.val, batch_size=self.batch_size)
+
+        datamodule = RandomDataModule(batch_size=16)
+
+        params = {
+            "datamodule": datamodule,
+            "lightning_module": SimpleLightningModule,
+            "trainer_kwargs": {
+                "max_epochs": 1,
+                "enable_progress_bar": False,
+                "enable_model_summary": False,
+                "logger": False,
+            },
+            "objective_metric": "val_loss",
+        }
+
+        class RegressionModule(L.LightningModule):
+            def __init__(self, num_layers=2, hidden_size=32, dropout=0.1):
+                super().__init__()
+                self.save_hyperparameters()
+                layers = []
+                input_size = 20
+                for _ in range(num_layers):
+                    layers.extend(
+                        [
+                            nn.Linear(input_size, hidden_size),
+                            nn.ReLU(),
+                            nn.Dropout(dropout),
+                        ]
+                    )
+                    input_size = hidden_size
+                layers.append(nn.Linear(hidden_size, 1))
+                self.model = nn.Sequential(*layers)
+
+            def forward(self, x):
+                return self.model(x)
+
+            def training_step(self, batch, batch_idx):
+                x, y = batch
+                y_hat = self(x).squeeze()
+                loss = nn.functional.mse_loss(y_hat, y)
+                self.log("train_loss", loss)
+                return loss
+
+            def validation_step(self, batch, batch_idx):
+                x, y = batch
+                y_hat = self(x).squeeze()
+                val_loss = nn.functional.mse_loss(y_hat, y)
+                self.log("val_loss", val_loss, on_epoch=True)
+                return val_loss
+
+            def configure_optimizers(self):
+                return torch.optim.SGD(self.parameters(), lr=0.01)
+
+        class RegressionDataModule(L.LightningDataModule):
+            def __init__(self, batch_size=16, num_samples=150):
+                super().__init__()
+                self.batch_size = batch_size
+                self.num_samples = num_samples
+
+            def setup(self, stage=None):
+                X = torch.randn(self.num_samples, 20)
+                y = torch.randn(self.num_samples)
+                dataset = torch.utils.data.TensorDataset(X, y)
+                train_size = int(0.8 * self.num_samples)
+                val_size = self.num_samples - train_size
+                self.train, self.val = torch.utils.data.random_split(
+                    dataset, [train_size, val_size]
+                )
+
+            def train_dataloader(self):
+                return DataLoader(self.train, batch_size=self.batch_size)
+
+            def val_dataloader(self):
+                return DataLoader(self.val, batch_size=self.batch_size)
+
+        datamodule2 = RegressionDataModule(batch_size=16, num_samples=150)
+
+        params2 = {
+            "datamodule": datamodule2,
+            "lightning_module": RegressionModule,
+            "trainer_kwargs": {
+                "max_epochs": 1,
+                "enable_progress_bar": False,
+                "enable_model_summary": False,
+                "logger": False,
+            },
+            "objective_metric": "val_loss",
+        }
+
+        return [params, params2]
+
+    @classmethod
+    def _get_score_params(cls):
+        """Return settings for testing score/evaluate functions.
+
+        Returns a list, the i-th element should be valid arguments for
+        self.evaluate and self.score, of an instance constructed with
+        self.get_test_params()[i].
+
+        Returns
+        -------
+        list of dict
+            The parameters to be used for scoring.
+        """
+        score_params1 = {"input_dim": 10, "hidden_dim": 20, "lr": 0.001}
+        score_params2 = {"num_layers": 3, "hidden_size": 64, "dropout": 0.2}
+        return [score_params1, score_params2]