Repository: facebookresearch/nevergrad

Issue:
Make optimizers use less memory
Optimizer's archive is a dict of type `Dict[Tuple[float,...], Value]` where `Value` is a light-weighted class. It keeps track of all function evaluations.
This is quite resourceful in term of memory for several reasons:
 1. it's a dict.
 2. the keys are tuples of float of the function dimension size (which could be millions).
 3. it keeps track of all the estimations.

While I cannot think of a way to get rid of the dict, we can deal with the other 2 issues:
 2. use some kind of hashable array instead of tuple of floats. This can lead to a factor 4 improvement at least. See some examples here https://stackoverflow.com/questions/16589791/most-efficient-property-to-hash-for-numpy-array
 3. implement pruning on the archive, getting rid of the lowest values for instance.



Discussion:

Resolving patch:
diff --git a/nevergrad/optimization/base.py b/nevergrad/optimization/base.py
--- a/nevergrad/optimization/base.py
+++ b/nevergrad/optimization/base.py
@@ -19,7 +19,6 @@
 
 
 registry = Registry()
-
 _OptimCallBack = Union[Callable[["Optimizer", ArrayLike, float], None], Callable[["Optimizer"], None]]
 
 
@@ -74,9 +73,13 @@ def __init__(self, dimension: int, budget: Optional[int] = None, num_workers: in
         self.dimension = dimension
         self.name = self.__class__.__name__  # printed name in repr
         # keep a record of evaluations, and current bests which are updated at each new evaluation
-        self.archive = utils.Archive()  # dict like structure taking np.ndarray as keys and Value as values
+        self.archive = utils.Archive[utils.Value]()  # dict like structure taking np.ndarray as keys and Value as values
         self.current_bests = {x: utils.Point(np.zeros(dimension, dtype=np.float), utils.Value(np.inf))
                               for x in ["optimistic", "pessimistic", "average"]}
+        # pruning function, called at each "tell"
+        # this can be desactivated or modified by each implementation
+        self.pruning: Optional[Callable[[utils.Archive[utils.Value]], utils.Archive[utils.Value]]] = None
+        self.pruning = utils.Pruning.sensible_default(num_workers=num_workers, dimension=dimension)
         # instance state
         self._num_ask = 0
         self._num_tell = 0
@@ -181,6 +184,8 @@ def _update_archive_and_bests(self, x: ArrayLike, value: float) -> None:
                         y = np.frombuffer(
                             min(self.archive.bytesdict, key=lambda z, n=name: self.archive.bytesdict[z].get_estimation(n)))
                         assert self.current_bests[name].x in self.archive, "Best value should exist in the archive"
+        if self.pruning is not None:
+            self.archive = self.pruning(self.archive)
 
     def ask(self) -> ArrayLike:
         """Provides a point to explore.
diff --git a/nevergrad/optimization/mutations.py b/nevergrad/optimization/mutations.py
--- a/nevergrad/optimization/mutations.py
+++ b/nevergrad/optimization/mutations.py
@@ -66,7 +66,7 @@ def crossover(parent: ArrayLike, donor: ArrayLike) -> ArrayLike:
     return discrete_mutation(mix)
 
 
-def get_roulette(archive: utils.Archive, num: Optional[int] = None) -> Any:
+def get_roulette(archive: utils.Archive[utils.Value], num: Optional[int] = None) -> Any:
     """Apply a roulette tournament selection.
     """
     if num is None:
diff --git a/nevergrad/optimization/utils.py b/nevergrad/optimization/utils.py
--- a/nevergrad/optimization/utils.py
+++ b/nevergrad/optimization/utils.py
@@ -3,6 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+import warnings
 import operator
 from uuid import uuid4
 from collections import OrderedDict, defaultdict
@@ -44,6 +45,7 @@ def pessimistic_confidence_bound(self) -> float:
         return float(self.mean + .1 * np.sqrt((self.variance) / (1 + self.count)))
 
     def get_estimation(self, name: str) -> float:
+        # Note: pruning below relies on the fact than only 3 modes exist. If a new mode is added, update pruning
         if name == "optimistic":
             return self.optimistic_confidence_bound
         elif name == "pessimistic":
@@ -172,31 +174,34 @@ def _tobytes(x: ArrayLike) -> bytes:
               "but it is less efficient.")
 
 
-class Archive:
+Y = TypeVar("Y")
+
+
+class Archive(Generic[Y]):
     """A dict-like object with numpy arrays as keys.
     The underlying `bytesdict` dict stores the arrays as bytes since arrays are not hashable.
     Keys can be converted back with np.frombuffer(key)
     """
 
     def __init__(self) -> None:
-        self.bytesdict: Dict[bytes, Value] = {}
+        self.bytesdict: Dict[bytes, Y] = {}
 
-    def __setitem__(self, x: ArrayLike, value: Value) -> None:
+    def __setitem__(self, x: ArrayLike, value: Y) -> None:
         self.bytesdict[_tobytes(x)] = value
 
-    def __getitem__(self, x: ArrayLike) -> Value:
+    def __getitem__(self, x: ArrayLike) -> Y:
         return self.bytesdict[_tobytes(x)]
 
     def __contains__(self, x: ArrayLike) -> bool:
         return _tobytes(x) in self.bytesdict
 
-    def get(self, x: ArrayLike, default: Optional[Value] = None) -> Optional[Value]:
+    def get(self, x: ArrayLike, default: Optional[Y] = None) -> Optional[Y]:
         return self.bytesdict.get(_tobytes(x), default)
 
     def __len__(self) -> int:
         return len(self.bytesdict)
 
-    def values(self) -> ValuesView[Value]:
+    def values(self) -> ValuesView[Y]:
         return self.bytesdict.values()
 
     def keys(self) -> None:
@@ -205,7 +210,7 @@ def keys(self) -> None:
     def items(self) -> None:
         raise RuntimeError(_ERROR_STR)
 
-    def items_as_array(self) -> Iterator[Tuple[np.ndarray, Value]]:
+    def items_as_array(self) -> Iterator[Tuple[np.ndarray, Y]]:
         """Functions that iterates on key-values but transforms keys
         to np.ndarray. This is to simplify interactions, but should not
         be used in an algorithm since the conversion can be inefficient.
@@ -233,6 +238,61 @@ def __iter__(self) -> None:
         raise RuntimeError(_ERROR_STR)
 
 
+class Pruning:
+    """Callable for pruning archives in the optimizer class.
+    See Optimizer.pruning attribute, called at each "tell".
+
+    Parameters
+    ----------
+    min_len: int
+        minimum length of the pruned archive.
+    max_len: int
+        length at which pruning is activated (maximum allowed length for the archive).
+
+    Note
+    ----
+    For each of the 3 criteria (optimistic, pessimistic and average), the min_len best (lowest)
+    points will be kept, which can lead to at most 3 * min_len points.
+    """
+
+    def __init__(self, min_len: int, max_len: int):
+        self.min_len = min_len
+        self.max_len = max_len
+
+    def __call__(self, archive: Archive[Value]) -> Archive[Value]:
+        if len(archive) < self.max_len:
+            return archive
+        warnings.warn("Pruning archive to save memory")
+        quantiles: Dict[str, float] = {}
+        threshold = float(self.min_len) / len(archive)
+        names = ["optimistic", "pessimistic", "average"]
+        for name in names:
+            quantiles[name] = np.quantile([v.get_estimation(name) for v in archive.values()], threshold)
+        new_archive = Archive[Value]()
+        new_archive.bytesdict = {b: v for b, v in archive.bytesdict.items() if any(v.get_estimation(n) <= quantiles[n] for n in names)}
+        return new_archive
+
+    @classmethod
+    def sensible_default(cls, num_workers: int, dimension: int) -> 'Pruning':
+        """ Very conservative pruning
+        - keep at least min_len 3 times num_workers
+        - keep at most 30 times min_len or up to 1GB of array memory (whatever is biggest)
+
+        Parameters
+        ----------
+        num_workers: int
+            number of evaluations which will be run in parallel at once
+        dimension: int
+            dimension of the optimization space
+        """
+        # safer to keep at least 3 time the workers
+        min_len = 3 * num_workers
+        max_len = 10 * 3 * min_len  # len after pruning can be up to 3 min_len, amortize with an order of magnitude
+        max_len_1gb = 1024**3 // (dimension * 8)
+        print(max_len, max_len_1gb)
+        return cls(min_len, max(max_len, max_len_1gb))
+
+
 class Particule:
 
     def __init__(self) -> None:


Test patch:
diff --git a/nevergrad/optimization/test_mutations.py b/nevergrad/optimization/test_mutations.py
--- a/nevergrad/optimization/test_mutations.py
+++ b/nevergrad/optimization/test_mutations.py
@@ -45,7 +45,7 @@ def test_run_with_array(func: Callable[..., Any]) -> None:
 )
 def test_get_roulette(num: int, expected: str) -> None:
     np.random.seed(24)
-    archive = utils.Archive()
+    archive = utils.Archive[utils.Value]()
     for k in range(4):
         archive[np.array([k + .5])] = utils.Value(k)
     output = mutations.get_roulette(archive, num)
diff --git a/nevergrad/optimization/test_utils.py b/nevergrad/optimization/test_utils.py
--- a/nevergrad/optimization/test_utils.py
+++ b/nevergrad/optimization/test_utils.py
@@ -3,6 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+import pytest
 import numpy as np
 from ..common import testing
 from .test_base import CounterFunction
@@ -59,10 +60,10 @@ def test_get_nash() -> None:
 
 def test_archive() -> None:
     data = [1, 4.5, 12, 0]
-    archive = utils.Archive()
-    archive[np.array(data)] = utils.Value(12)
-    np.testing.assert_equal(archive[np.array(data)].mean, 12)
-    np.testing.assert_equal(archive.get(data).mean, 12)  # type: ignore
+    archive = utils.Archive[int]()
+    archive[np.array(data)] = 12
+    np.testing.assert_equal(archive[np.array(data)], 12)
+    np.testing.assert_equal(archive.get(data), 12)
     np.testing.assert_equal(archive.get([0, 12.]), None)
     y = np.frombuffer(next(iter(archive.bytesdict.keys())))
     assert data in archive
@@ -76,8 +77,8 @@ def test_archive() -> None:
 
 
 def test_archive_errors() -> None:
-    archive = utils.Archive()
-    archive[[12, 0.]] = utils.Value(12)
+    archive = utils.Archive[float]()
+    archive[[12, 0.]] = 12.
     np.testing.assert_raises(AssertionError, archive.__getitem__, [12, 0])  # int instead of float
     np.testing.assert_raises(AssertionError, archive.__getitem__, [[12], [0.]])  # int instead of float
     np.testing.assert_raises(RuntimeError, archive.keys)
@@ -128,3 +129,26 @@ def test_population_replace() -> None:
     for uuid in pop.uuids:
         # checks that it exists and correctly linked
         pop[uuid]  # pylint: disable= pointless-statement
+
+
+def test_pruning() -> None:
+    archive = utils.Archive[utils.Value]()
+    for k in range(3):
+        value = utils.Value(float(k))
+        archive[(float(k),)] = value
+    value = utils.Value(1.)
+    value.add_evaluation(1.)
+    archive[(3.,)] = value
+    # pruning
+    pruning = utils.Pruning(min_len=1, max_len=3)
+    # 0 is best optimistic and average, and 3 is best pessimistic (variance=0)
+    for k in range(2):
+        archive = pruning(archive)
+        testing.assert_set_equal([x[0] for x in archive.keys_as_array()], [0, 3], err_msg=f"Repetition #{k+1}")
+
+
+@pytest.mark.parametrize("dimension,expected_max", [(100, 1342177), (10000, 13421), (1000000, 1080)])  # type: ignore
+def test_pruning_sensible_default(dimension, expected_max) -> None:
+    pruning = utils.Pruning.sensible_default(num_workers=12, dimension=dimension)
+    assert pruning.min_len == 36
+    assert pruning.max_len == expected_max