Add a simple early stopping mechanism

CHANGELOG.md

@@ -26,6 +26,8 @@
 
 - the new `nevergrad.errors` module gathers errors and warnings used throughout the package (WIP) [#1031](https://github.com/facebookresearch/nevergrad/pull/1031).
 - `EvolutionStrategy` now defaults to NSGA2 selection in the multiobjective case
+- A new experimental callback adds an early stopping mechanism
+  [#1054](https://github.com/facebookresearch/nevergrad/pull/1054).
 
 ## 0.4.3 (2021-01-28)
 

docs/optimizers_ref.rst

@@ -20,7 +20,7 @@ Callbacks can be registered through the :code:`optimizer.register_callback` for
 `ng.callbacks` namespace.
 
 .. automodule:: nevergrad.callbacks
-    :members: OptimizerDump, ParametersLogger, ProgressBar
+    :members: OptimizerDump, ParametersLogger, ProgressBar, EarlyStopping
 
 Configurable optimizers
 -----------------------

docs/parametrization_ref.rst

@@ -38,3 +38,4 @@ Parameter API
 
 .. autoclass:: nevergrad.p.Parameter
     :members:
+    :inherited-members:

nevergrad/common/errors.py

@@ -20,6 +20,10 @@ class NevergradWarning(Warning):
 # errors
 
 
+class NevergradEarlyStopping(StopIteration, NevergradError):
+    """Stops the minimization loop if raised"""
+
+
 class NevergradRuntimeError(RuntimeError, NevergradError):
     """Runtime error raised by Nevergrad"""
 

nevergrad/optimization/base.py

@@ -599,13 +599,18 @@ def minimize(
                 if verbosity and new_sugg:
                     print(f"Launching {new_sugg} jobs with new suggestions")
                 for _ in range(new_sugg):
-                    args = self.ask()
+                    try:
+                        args = self.ask()
+                    except errors.NevergradEarlyStopping:
+                        remaining_budget = 0
+                        break
                     self._running_jobs.append(
                         (args, executor.submit(objective_function, *args.args, **args.kwargs))
                     )
                 if new_sugg:
                     sleeper.start_timer()
-            remaining_budget = self.budget - self.num_ask
+            if remaining_budget > 0:  # early stopping sets it to 0
+                remaining_budget = self.budget - self.num_ask
             # split (repopulate finished and runnings in only one loop to avoid
             # weird effects if job finishes in between two list comprehensions)
             tmp_runnings, tmp_finished = [], deque()

nevergrad/optimization/callbacks.py

@@ -11,6 +11,7 @@
 from pathlib import Path
 import numpy as np
 import nevergrad.common.typing as tp
+from nevergrad.common import errors
 from nevergrad.parametrization import parameter as p
 from nevergrad.parametrization import helpers
 from . import base
@@ -246,3 +247,36 @@ def __getstate__(self) -> tp.Dict[str, tp.Any]:
         state = dict(self.__dict__)
         state["_progress_bar"] = None
         return state
+
+
+class EarlyStopping:
+    """Callback for stopping the :code:`minimize` method before the budget is
+    fully used.
+
+    Parameters
+    ----------
+    stopping_criterion: func(optimizer) -> bool
+        function that takes the current optimizer as input and returns True
+        if the minimization must be stopped
+
+    Note
+    ----
+    This callback must be register on the "ask" method only.
+
+    Example
+    -------
+    In the following code, the :code:`minimize` method will be stopped at the 4th "ask"
+
+    >>> early_stopping = ng.callbacks.EarlyStopping(lambda opt: opt.num_ask > 3)
+    >>> optimizer.register_callback("ask", early_stopping)
+    >>> optimizer.minimize(_func, verbosity=2)
+    """
+
+    def __init__(self, stopping_criterion: tp.Callable[[base.Optimizer], bool]) -> None:
+        self.stopping_criterion = stopping_criterion
+
+    def __call__(self, optimizer: base.Optimizer, *args: tp.Any, **kwargs: tp.Any) -> None:
+        if args or kwargs:
+            raise errors.NevergradRuntimeError("EarlyStopping must be registered on ask method")
+        if self.stopping_criterion(optimizer):
+            raise errors.NevergradEarlyStopping("Early stopping criterion is reached")

nevergrad/optimization/test_callbacks.py

@@ -81,3 +81,14 @@ def test_progressbar_dump(tmp_path: Path) -> None:
     for _ in range(12):
         cand = optimizer.ask()
         optimizer.tell(cand, 0)
+
+
+def test_early_stopping() -> None:
+    instrum = ng.p.Instrumentation(
+        None, 2.0, blublu="blublu", array=ng.p.Array(shape=(3, 2)), multiobjective=True
+    )
+    optimizer = optimizerlib.OnePlusOne(parametrization=instrum, budget=100)
+    early_stopping = ng.callbacks.EarlyStopping(lambda opt: opt.num_ask > 3)
+    optimizer.register_callback("ask", early_stopping)
+    optimizer.minimize(_func, verbosity=2)
+    assert optimizer.num_ask == 4