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:
jrapin: TODO: add docstrings + update documentation

Patch:
diff --git a/docs/adding_an_algorithm.md b/docs/adding_an_algorithm.md
index 8ca3684d..5936c62b 100644
--- a/docs/adding_an_algorithm.md
+++ b/docs/adding_an_algorithm.md
@@ -31,9 +31,10 @@ The key tuple if the point location, and `Value` is a class with attributes:
 - `mean`: mean value of the evaluations at this point.
 - `variance`: variance of the evaluations at this point.
 
-For more detauls, see the implementation in [utils.py](../nevergrad/optimization/utils.py).
+For more details, see the implementation in [utils.py](../nevergrad/optimization/utils.py).
 
-Through the archive, you can therefore access most useful information about past evaluations.
+Through the archive, you can therefore access most useful information about past evaluations. A pruning mechanism makes sure this archive does
+not grow too much. This pruning can be tuned through the `pruning` attribute of the optimizer (the default is very conservative).
 
 The base `Optimizer` class also tracks the best optimistic and pessimistic points through the `current_bests` attribute which is of type:
 ```
diff --git a/nevergrad/optimization/base.py b/nevergrad/optimization/base.py
index df8a801d..13653a05 100644
--- a/nevergrad/optimization/base.py
+++ b/nevergrad/optimization/base.py
@@ -19,7 +19,6 @@ from . import utils
 
 
 registry = Registry()
-
 _OptimCallBack = Union[Callable[["Optimizer", ArrayLike, float], None], Callable[["Optimizer"], None]]
 
 
@@ -74,9 +73,13 @@ class Optimizer(abc.ABC):  # pylint: disable=too-many-instance-attributes
         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 @@ class Optimizer(abc.ABC):  # pylint: disable=too-many-instance-attributes
                         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
index b9b1900c..16b45493 100644
--- 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/optimizerlib.py b/nevergrad/optimization/optimizerlib.py
index f8fdccbe..5b8c114d 100644
--- a/nevergrad/optimization/optimizerlib.py
+++ b/nevergrad/optimization/optimizerlib.py
@@ -434,6 +434,14 @@ class MEDA(EDA):
             self.evaluated_population_fitness = []
 
 
+class ParticuleTBPSA:
+
+    def __init__(self, position: np.ndarray, sigma: float, loss: Optional[float] = None) -> None:
+        self.position = np.array(position, copy=False)
+        self.sigma = sigma
+        self.loss = loss
+
+
 @registry.register
 class TBPSA(base.Optimizer):
     """Test-based population-size adaptation.
@@ -451,36 +459,28 @@ class TBPSA(base.Optimizer):
         if num_workers is not None:
             self.llambda = max(self.llambda, num_workers)
         self.current_center: np.ndarray = np.zeros(dimension)
-        # Evaluated population
-        self.evaluated_population: List[base.ArrayLike] = []
-        self.evaluated_population_sigma: List[float] = []
-        self.evaluated_population_fitness: List[float] = []
-        # Unevaluated population
-        self.unevaluated_population: List[base.ArrayLike] = []
-        self.unevaluated_population_sigma: List[float] = []
-        # Archive
-        self.archive_fitness: List[float] = []
+        self._loss_record: List[float] = []
+        # population
+        self._evaluated_population: List[ParticuleTBPSA] = []
+        self._unevaluated_population: Dict[bytes, ParticuleTBPSA] = {}
 
     def _internal_provide_recommendation(self) -> base.ArrayLike:  # This is NOT the naive version. We deal with noise.
         return self.current_center
 
     def _internal_ask(self) -> base.ArrayLike:
         mutated_sigma = self.sigma * np.exp(np.random.normal(0, 1) / np.sqrt(self.dimension))
-        individual = tuple(self.current_center + mutated_sigma * np.random.normal(0, 1, self.dimension))
-        self.unevaluated_population_sigma += [mutated_sigma]
-        self.unevaluated_population += [tuple(individual)]
-        return individual
+        individual = self.current_center + mutated_sigma * np.random.normal(0, 1, self.dimension)
+        self._unevaluated_population[individual.tobytes()] = ParticuleTBPSA(individual, sigma=mutated_sigma)
+        return individual  # type: ignore
 
     def _internal_tell(self, x: base.ArrayLike, value: float) -> None:
-        self.archive_fitness += [value]
-        if len(self.archive_fitness) >= 5 * self.llambda:
-            first_fifth = [self.archive_fitness[i] for i in range(self.llambda)]
-            last_fifth = [self.archive_fitness[i] for i in range(4*self.llambda, 5*self.llambda)]
-            mean1 = sum(first_fifth) / float(self.llambda)
-            std1 = np.std(first_fifth) / np.sqrt(self.llambda - 1)
-            mean2 = sum(last_fifth) / float(self.llambda)
-            std2 = np.std(last_fifth) / np.sqrt(self.llambda - 1)
-            z = (mean1 - mean2) / (np.sqrt(std1**2 + std2**2))
+        self._loss_record += [value]
+        if len(self._loss_record) >= 5 * self.llambda:
+            first_fifth = self._loss_record[: self.llambda]
+            last_fifth = self._loss_record[-self.llambda:]
+            means = [sum(fitnesses) / float(self.llambda) for fitnesses in [first_fifth, last_fifth]]
+            stds = [np.std(fitnesses) / np.sqrt(self.llambda - 1) for fitnesses in [first_fifth, last_fifth]]
+            z = (means[0] - means[1]) / (np.sqrt(stds[0]**2 + stds[1]**2))
             if z < 2.:
                 self.mu *= 2
             else:
@@ -491,29 +491,27 @@ class TBPSA(base.Optimizer):
             if self.num_workers > 1:
                 self.llambda = max(self.llambda, self.num_workers)
                 self.mu = self.llambda // 4
-            self.archive_fitness = []
-        idx = self.unevaluated_population.index(tuple(x))
-        self.evaluated_population += [x]
-        self.evaluated_population_fitness += [value]
-        self.evaluated_population_sigma += [self.unevaluated_population_sigma[idx]]
-        del self.unevaluated_population[idx]
-        del self.unevaluated_population_sigma[idx]
-        if len(self.evaluated_population) >= self.llambda:
+            self._loss_record = []
+        x = np.array(x, copy=False)
+        x_bytes = x.tobytes()
+        particule = self._unevaluated_population[x_bytes]
+        particule.loss = value
+        self._evaluated_population.append(particule)
+        if len(self._evaluated_population) >= self.llambda:
             # Sorting the population.
-            sorted_pop_with_sigma_and_fitness = [(i, s, f) for f, i, s in sorted(
-                zip(self.evaluated_population_fitness, self.evaluated_population, self.evaluated_population_sigma))]
-            self.evaluated_population = [p[0] for p in sorted_pop_with_sigma_and_fitness]
-            self.evaluated_population_sigma = [p[1] for p in sorted_pop_with_sigma_and_fitness]
-            self.evaluated_population_fitness = [p[2] for p in sorted_pop_with_sigma_and_fitness]
+            self._evaluated_population.sort(key=lambda p: p.loss)
             # Computing the new parent.
-            self.current_center = sum([np.asarray(self.evaluated_population[i]) for i in range(self.mu)]) / self.mu  # type: ignore
-            self.sigma = np.exp(sum([np.log(self.evaluated_population_sigma[i]) for i in range(self.mu)]) / self.mu)
-            self.evaluated_population = []
-            self.evaluated_population_sigma = []
-            self.evaluated_population_fitness = []
+            self.current_center = sum(p.position for p in self._evaluated_population[:self.mu]) / self.mu  # type: ignore
+            self.sigma = np.exp(np.sum(np.log([p.sigma for p in self._evaluated_population[:self.mu]])) / self.mu)
+            self._evaluated_population = []
+        del self._unevaluated_population[x_bytes]
 
     def tell_not_asked(self, x: base.ArrayLike, value: float) -> None:
-        raise base.TellNotAskedNotSupportedError
+        x = np.array(x, copy=False)
+        sigma = np.linalg.norm(x - self.current_center) / np.sqrt(self.dimension)  # educated guess
+        self._unevaluated_population[x.tobytes()] = ParticuleTBPSA(x, sigma=sigma)
+        # go through standard pipeline so as to update the archive
+        self.tell(x, value)
 
 
 @registry.register
diff --git a/nevergrad/optimization/utils.py b/nevergrad/optimization/utils.py
index 0856563b..8e6e5f0b 100644
--- 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 @@ class Value:
         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 @@ _ERROR_STR = ("Generating numpy arrays from the bytes keys is inefficient, "
               "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 @@ class Archive:
     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 @@ class Archive:
         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
index a072a93e..48281c97 100644
--- 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_optimizerlib.py b/nevergrad/optimization/test_optimizerlib.py
index c5941065..8850ddbb 100644
--- a/nevergrad/optimization/test_optimizerlib.py
+++ b/nevergrad/optimization/test_optimizerlib.py
@@ -196,3 +196,14 @@ def test_tell_not_asked(name: str) -> None:
     if (0, 0, 0, 0) not in [tuple(x) for x in asked]:
         for value in opt.archive.values():
             assert value.count == 1
+
+
+def test_tbpsa_recom_with_update() -> None:
+    np.random.seed(12)
+    budget = 20
+    # set up problem
+    fitness = Fitness([.5, -.8, 0, 4])
+    optim = optimizerlib.TBPSA(dimension=4, budget=budget, num_workers=1)
+    optim.llambda = 3
+    output = optim.optimize(fitness)
+    np.testing.assert_almost_equal(output, [.037964, .0433031, -.4688667, .3633273])
diff --git a/nevergrad/optimization/test_utils.py b/nevergrad/optimization/test_utils.py
index 6284c328..1dc14de1 100644
--- 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