Optuna GridSearch

docs/source/whats_new.rst

@@ -17,7 +17,7 @@ Develop branch
 
 Enhancements
 ~~~~~~~~~~~~
-- None
+- - Add possibility to use RandomizedGridSearch (:gh:`630` by `Igor Carrara`_)
 
 Bugs
 ~~~~

moabb/evaluations/base.py

@@ -2,16 +2,20 @@
 from abc import ABC, abstractmethod
 
 import pandas as pd
+from optuna.integration import OptunaSearchCV
 from sklearn.base import BaseEstimator
 from sklearn.model_selection import GridSearchCV
 
 from moabb.analysis import Results
 from moabb.datasets.base import BaseDataset
+from moabb.evaluations.utils import _convert_sklearn_params_to_optuna
 from moabb.paradigms.base import BaseParadigm
 
 
 log = logging.getLogger(__name__)
 
+search_methods = {"grid": GridSearchCV, "optuna": OptunaSearchCV}
+
 
 class BaseEvaluation(ABC):
     """Base class that defines necessary operations for an evaluation.
@@ -53,6 +57,9 @@ class BaseEvaluation(ABC):
         Save model after training, for each fold of cross-validation if needed
     cache_config: bool, default=None
         Configuration for caching of datasets. See :class:`moabb.datasets.base.CacheConfig` for details.
+    optuna:bool, default=False
+        If optuna is enable it will change the GridSearch to a RandomizedGridSearch with 15 minutes of cut off time.
+        This option is compatible with list of entries of type None, bool, int, float and string
 
     Notes
     -----
@@ -77,6 +84,7 @@ def __init__(
         n_splits=None,
         save_model=False,
         cache_config=None,
+        optuna=False,
     ):
         self.random_state = random_state
         self.n_jobs = n_jobs
@@ -88,6 +96,7 @@ def __init__(
         self.n_splits = n_splits
         self.save_model = save_model
         self.cache_config = cache_config
+        self.optuna = optuna
         # check paradigm
         if not isinstance(paradigm, BaseParadigm):
             raise (ValueError("paradigm must be an Paradigm instance"))
@@ -261,19 +270,27 @@ def is_valid(self, dataset):
         """
 
     def _grid_search(self, param_grid, name, grid_clf, inner_cv):
+        extra_params = {}
         if param_grid is not None:
             if name in param_grid:
-                search = GridSearchCV(
+                if self.optuna:
+                    search = search_methods["optuna"]
+                    param_grid[name] = _convert_sklearn_params_to_optuna(param_grid[name])
+                    extra_params["timeout"] = 60 * 15  # 15 minutes
+                else:
+                    search = search_methods["grid"]
+
+                search = search(
                     grid_clf,
                     param_grid[name],
                     refit=True,
                     cv=inner_cv,
                     n_jobs=self.n_jobs,
                     scoring=self.paradigm.scoring,
                     return_train_score=True,
+                    **extra_params,
                 )
                 return search
-
             else:
                 return grid_clf
 

moabb/evaluations/utils.py

@@ -4,6 +4,7 @@
 from pickle import HIGHEST_PROTOCOL, dump
 from typing import Sequence
 
+import optuna
 from numpy import argmax
 from sklearn.pipeline import Pipeline
 
@@ -212,3 +213,24 @@ def create_save_path(
         return str(path_save)
     else:
         print("No hdf5_path provided, models will not be saved.")
+
+
+def _convert_sklearn_params_to_optuna(param_grid):
+    """
+    Function to convert the parameter in Optuna format. This function will create a integer distribution of values
+    between the max and minimum value of the parameter.
+    Parameters
+    ----------
+    param_grid
+
+    Returns
+    -------
+
+    """
+    optuna_params = {}
+    for key, value in param_grid.items():
+        if isinstance(value, list):
+            optuna_params[key] = optuna.distributions.CategoricalDistribution(value)
+        else:
+            optuna_params[key] = value
+    return optuna_params

moabb/tests/evaluations.py

@@ -58,6 +58,7 @@ def setUp(self):
             datasets=[dataset],
             hdf5_path="res_test",
             save_model=True,
+            optuna=False,
         )
 
     def test_mne_labels(self):
@@ -138,6 +139,30 @@ def test_eval_grid_search(self):
         # We should have 9 columns in the results data frame
         self.assertEqual(len(results[0].keys()), 9 if _carbonfootprint else 8)
 
+    def test_eval_grid_search_optuna(self):
+        # Test grid search
+        param_grid = {"C": {"csp__metric": ["euclid", "riemann"]}}
+        process_pipeline = self.eval.paradigm.make_process_pipelines(dataset)[0]
+
+        self.eval.optuna = True
+
+        results = [
+            r
+            for r in self.eval.evaluate(
+                dataset,
+                pipelines,
+                param_grid=param_grid,
+                process_pipeline=process_pipeline,
+            )
+        ]
+
+        self.eval.optuna = False
+
+        # We should get 4 results, 2 sessions 2 subjects
+        self.assertEqual(len(results), 4)
+        # We should have 9 columns in the results data frame
+        self.assertEqual(len(results[0].keys()), 9 if _carbonfootprint else 8)
+
     def test_within_session_evaluation_save_model(self):
         res_test_path = "./res_test"