lower max learning rate value to avoid gradient explode. set verbose …

…and error_score to both Searcher
alan-turing-institute · Mar 13, 2024 · 06ba365 · 06ba365
1 parent f70d018
commit 06ba365
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 9 deletions.
diff --git a/autoemulate/emulators/neural_networks/mlp.py b/autoemulate/emulators/neural_networks/mlp.py
@@ -50,18 +50,18 @@ def get_grid_params(search_type: str = "random"):
                 nn.Sigmoid,
                 nn.GELU,
             ],
-            "optimizer": [torch.optim.AdamW, torch.optim.LBFGS],
+            "optimizer": [torch.optim.SGD, torch.optim.AdamW, torch.optim.LBFGS],
             "optimizer__weight_decay": (1 / 10 ** np.arange(1, 9)).tolist(),
         }
         match search_type:
             case "random":
                 param_space |= {
-                    "lr": loguniform(1e-06, 1e-2),
+                    "lr": loguniform(1e-6, 1e-4),
                 }
             case "bayes":
                 param_space |= {
                     "optimizer": Categorical(param_space["optimizer"]),
-                    "lr": Real(1e-06, 1e-2, prior="log-uniform"),
+                    "lr": Real(1e-6, 1e-4, prior="log-uniform"),
                 }
             case _:
                 raise ValueError(f"Invalid search type: {search_type}")

diff --git a/autoemulate/emulators/neural_networks/rbf.py b/autoemulate/emulators/neural_networks/rbf.py
@@ -324,18 +324,18 @@ def get_grid_params(search_type: str = "random"):
                 rbf_inverse_quadratic,
                 rbf_inverse_multiquadric,
             ],
-            "optimizer": [torch.optim.AdamW, torch.optim.SGD],
+            "optimizer": [torch.optim.SGD, torch.optim.AdamW, torch.optim.LBFGS],
             "optimizer__weight_decay": (1 / 10 ** np.arange(1, 9)).tolist(),
         }
         match search_type:
             case "random":
                 param_space |= {
-                    "lr": loguniform(1e-06, 1e-2),
+                    "lr": loguniform(1e-06, 1e-3),
                 }
             case "bayes":
                 param_space |= {
                     "optimizer": Categorical(param_space["optimizer"]),
-                    "lr": Real(1e-06, 1e-2, prior="log-uniform"),
+                    "lr": Real(1e-06, 1e-3, prior="log-uniform"),
                 }
             case _:
                 raise ValueError(f"Invalid search type: {search_type}")

diff --git a/autoemulate/hyperparam_searching.py b/autoemulate/hyperparam_searching.py
@@ -23,6 +23,8 @@ def _optimize_params(
     param_space=None,
     n_jobs=None,
     logger=None,
+    error_score=np.nan,
+    verbose=0,
 ):
     """Performs hyperparameter search for the provided model.
 
@@ -49,7 +51,11 @@ def _optimize_params(
     n_jobs : int
         Number of jobs to run in parallel.
     logger : logging.Logger
-        Logger instance.
+        Logger instance
+    error_score: 'raise' or numeric
+        Value to assign to the score if an error occurs in estimator fitting.
+    verbose: int
+        Verbosity level for the searcher
 
     Returns
     -------
@@ -68,7 +74,8 @@ def _optimize_params(
             cv=cv,
             n_jobs=n_jobs,
             refit=True,
-            verbose=0,
+            error_score=error_score,
+            verbose=verbose,
         )
     # Bayes search
     elif search_type == "bayes":
@@ -79,7 +86,8 @@ def _optimize_params(
             cv=cv,
             n_jobs=n_jobs,
             refit=True,
-            verbose=0,
+            error_score=error_score,
+            verbose=verbose,
         )
     elif search_type == "grid":
         raise NotImplementedError("Grid search not available yet.")