Source code for orion.algo.nevergradoptimizer

"""
Nevergrad Optimizer
===================

Wraps the nevergrad library to expose its algorithm to orion

"""
from __future__ import annotations

import logging
import pickle
from typing import Callable, Iterable, Sequence, SupportsInt

from orion.core.worker.trial import Trial

try:
    import nevergrad as ng
    from nevergrad.parametrization.container import Instrumentation
    from nevergrad.parametrization.core import Parameter

    IMPORT_ERROR = None

except ImportError as err:
    IMPORT_ERROR = err


from orion.algo.base import BaseAlgorithm
from orion.algo.space import Categorical, Dimension, Fidelity, Integer, Real, Space
from orion.core.utils.format_trials import dict_to_trial

logger = logging.getLogger(__name__)

registry: dict[tuple[str, str], Callable[[Dimension], Parameter]] = {}


[docs]def register(dimension_type: str, prior: str): """Register a conversion function for the given type and prior.""" def deco(func): registry[dimension_type, prior] = func return func return deco
[docs]def to_ng_space(orion_space: Space) -> Instrumentation: """Convert an orion space to a nevergrad space.""" if IMPORT_ERROR: raise IMPORT_ERROR converted_dimensions: dict[str, Parameter] = {} for name, dim in orion_space.items(): try: converted_dimensions[name] = registry[dim.type, dim.prior_name](dim) except KeyError as exc: raise RuntimeError( f"Dimension with type and prior: {exc.args[0]} cannot be converted to nevergrad." ) from exc return ng.p.Instrumentation(**converted_dimensions)
def _intshape(shape: Iterable[SupportsInt]) -> tuple[int, ...]: # ng.p.Array does not accept np.int64 in shapes, they have to be ints return tuple(int(x) for x in shape) @register("categorical", "choices") def _(dim: Categorical): if dim.shape: raise NotImplementedError("Array of Categorical cannot be converted.") if len(set(dim.original_dimension.prior.pk)) != 1: raise NotImplementedError( "All categories in Categorical must have the same probability." ) return ng.p.Choice(dim.interval()) @register("real", "uniform") def _(dim: Dimension): lower, upper = dim.interval() if dim.shape: # Temporary fix pending [#800] # ng.p.Array expects an array with a shape or a float # an array with an empty shape is not valid # so we cast it to a float if hasattr(lower, "shape") and lower.shape == (): lower = float(lower) if hasattr(upper, "shape") and upper.shape == (): upper = float(upper) return ng.p.Array(lower=lower, upper=upper, shape=_intshape(dim.shape)) else: return ng.p.Scalar(lower=lower, upper=upper) @register("integer", "int_uniform") def _(dim: Integer): return registry["real", "uniform"](dim).set_integer_casting() @register("integer", "int_reciprocal") def _(dim: Integer): return registry["real", "reciprocal"](dim).set_integer_casting() @register("real", "reciprocal") def _(dim: Real): if dim.shape: raise NotImplementedError("Array with reciprocal prior cannot be converted.") lower, upper = dim.interval() return ng.p.Log(lower=lower, upper=upper, exponent=2) @register("real", "norm") @register("real", "normal") def _(dim: Real): if dim.shape: raise NotImplementedError("Array with normal prior cannot be converted.") return ng.p.Scalar(init=dim.original_dimension.prior.mean()).set_mutation( sigma=dim.original_dimension.prior.std() ) @register("fidelity", "None") def _(dim: Fidelity): if dim.shape: raise NotImplementedError("Array of Fidelity cannot be converted.") _, upper = dim.interval() # No equivalent to Fidelity space, so we always use the upper value return upper NOT_WORKING = { "BO", "BOSplit", "BayesOptimBO", "DiscreteDoerrOnePlusOne", "NelderMead", "PCABO", "SPSA", "SQP", "NoisyBandit", "NoisyOnePlusOne", "OptimisticDiscreteOnePlusOne", "OptimisticNoisyOnePlusOne", "Powell", "CmaFmin2", "RPowell", "RSQP", "PymooNSGA2", }
[docs]class NevergradOptimizer(BaseAlgorithm): """Wraps the nevergrad library to expose its algorithm to orion Parameters ---------- space: `orion.algo.space.Space` Optimisation space with priors for each dimension. model_name: str Nevergrad model to use as optimizer budget: int Maximal number of trial to generated num_workers: int Number of worker to use seed: None, int or sequence of int Seed for the random number generator used to sample new trials. Default: ``None`` """ requires_type = None requires_dist = None requires_shape = None def __init__( self, space: Space, model_name: str = "NGOpt", seed: int | Sequence[int] | None = None, budget: int = 100, num_workers: int = 10, ): if IMPORT_ERROR: raise IMPORT_ERROR super().__init__(space) self.model_name = model_name self.seed = seed self.budget = budget self.num_workers = num_workers if model_name in NOT_WORKING: raise ValueError(f"Model {model_name} is not supported.") param = to_ng_space(space) self.algo = ng.optimizers.registry[model_name]( parametrization=param, budget=budget, num_workers=num_workers ) self.algo.enable_pickling() self._trial_mapping = {} self._fresh = True self._is_done = False self.seed = seed if seed is not None: self.seed_rng(seed)
[docs] def seed_rng(self, seed): """Seed the state of the random number generator. Parameters ---------- seed: int Integer seed for the random number generator. """ self.algo.parametrization.random_state.seed(seed)
@property def state_dict(self): """Return a state dict that can be used to reset the state of the algorithm.""" state_dict = super().state_dict state_dict["algo"] = pickle.dumps(self.algo) # type: ignore state_dict["_is_done"] = self._is_done state_dict["_fresh"] = self._fresh state_dict["_trial_mapping"] = { trial_id: list(suggestions) for trial_id, suggestions in self._trial_mapping.items() } return state_dict
[docs] def set_state(self, state_dict): """Reset the state of the algorithm based on the given state_dict Parameters ---------- state_dict: dict Dictionary representing state of an algorithm """ super().set_state(state_dict) self.algo = pickle.loads(state_dict["algo"]) self._is_done = state_dict["_is_done"] self._fresh = state_dict["_fresh"] self._trial_mapping = state_dict["_trial_mapping"]
def _associate_trial(self, trial: Trial, suggestion: Parameter): """Associate a trial with a Nevergrad suggestion. Returns ------- True if the trial was not seen before, false otherwise. """ trial_id = self.get_id(trial) seen = trial_id in self._trial_mapping if seen: orig_trial, existing = self._trial_mapping[trial_id] if orig_trial.status == "completed": self.algo.tell(suggestion, orig_trial.objective.value) else: existing.append(suggestion) else: self._trial_mapping[trial_id] = (trial, [suggestion]) return not seen def _ask(self): suggestion = self.algo.ask() if suggestion.args: raise RuntimeError( "Nevergrad sampled a trial with args but this should never happen." " Please report this issue at" " https://github.com/Epistimio/orion.algo.nevergrad/issues" ) new_trial = dict_to_trial(suggestion.kwargs, self.space) if self._associate_trial(new_trial, suggestion): self.register(new_trial) return new_trial else: logger.debug("Ignoring duplicated trial") return None def _can_produce(self): if self.is_done: return False algo = self.algo is_sequential = algo.no_parallelization if not is_sequential and hasattr(algo, "optim"): is_sequential = algo.optim.no_parallelization if is_sequential and algo.num_ask > (algo.num_tell - algo.num_tell_not_asked): logger.debug( "Cannot produce new trials because %d > (%d - %d)", algo.num_ask, algo.num_tell, algo.num_tell_not_asked, ) return False return True
[docs] def suggest(self, num: int) -> list[Trial]: """Suggest a number of new sets of parameters. Parameters ---------- num: int, optional Number of trials to suggest. The algorithm may return less than the number of trials requested. Returns ------- list of trials A list of trials representing values suggested by the algorithm. The algorithm may opt out if it cannot make a good suggestion at the moment (it may be waiting for other trials to complete), in which case it will return None. """ attempts = 0 max_attempts = num + 100 trials: list[Trial] = [] while len(trials) < num and attempts < max_attempts and self._can_produce(): attempts += 1 trial = self._ask() if trial is not None: trials.append(trial) if not trials and self._can_produce(): self._is_done = True self._fresh = False return trials
[docs] def observe(self, trials: list[Trial]) -> None: """Observe the trials new state of result. Parameters ---------- trials: list of ``orion.core.worker.trial.Trial`` Trials from a `orion.algo.space.Space`. """ for trial in trials: if trial.status == "completed": tid = self.get_id(trial) if tid in self._trial_mapping: _, suggestions = self._trial_mapping[tid] else: sugg = self.algo.parametrization.spawn_child(((), trial.params)) suggestions = [sugg] for suggestion in suggestions: self.algo.tell(suggestion, trial.objective.value) self._trial_mapping[tid] = (trial, []) self._fresh = True super().observe(trials)
@property def is_done(self) -> bool: return self._is_done or super().is_done