Source code for ray.tune.search.zoopt.zoopt_search

import copy
import logging
from typing import Dict, List, Optional, Tuple

import ray
import ray.cloudpickle as pickle
from ray.tune.result import DEFAULT_METRIC
from ray.tune.search import (
    UNDEFINED_METRIC_MODE,
    UNDEFINED_SEARCH_SPACE,
    UNRESOLVED_SEARCH_SPACE,
    Searcher,
)
from ray.tune.search.sample import (
    Categorical,
    Domain,
    Float,
    Integer,
    Quantized,
    Uniform,
)
from ray.tune.search.variant_generator import parse_spec_vars
from ray.tune.utils.util import unflatten_dict

try:
    import zoopt
    from zoopt import Solution, ValueType
except ImportError:
    zoopt = None
    Solution = ValueType = None

logger = logging.getLogger(__name__)


[docs] class ZOOptSearch(Searcher): """A wrapper around ZOOpt to provide trial suggestions. ZOOptSearch is a library for derivative-free optimization. It is backed by the `ZOOpt <https://github.com/polixir/ZOOpt>`__ package. Currently, Asynchronous Sequential RAndomized COordinate Shrinking (ASRacos) is implemented in Tune. To use ZOOptSearch, install zoopt (>=0.4.1): ``pip install -U zoopt``. Tune automatically converts search spaces to ZOOpt"s format: .. code-block:: python from ray import train, tune from ray.tune.search.zoopt import ZOOptSearch "config": { "iterations": 10, # evaluation times "width": tune.uniform(-10, 10), "height": tune.uniform(-10, 10) } zoopt_search_config = { "parallel_num": 8, # how many workers to parallel } zoopt_search = ZOOptSearch( algo="Asracos", # only support Asracos currently budget=20, # must match `num_samples` in `tune.TuneConfig()`. dim_dict=dim_dict, metric="mean_loss", mode="min", **zoopt_search_config ) tuner = tune.Tuner( my_objective, tune_config=tune.TuneConfig( search_alg=zoopt_search, num_samples=20 ), run_config=train.RunConfig( name="zoopt_search", stop={"timesteps_total": 10} ), param_space=config ) tuner.fit() If you would like to pass the search space manually, the code would look like this: .. code-block:: python from ray import train, tune from ray.tune.search.zoopt import ZOOptSearch from zoopt import ValueType dim_dict = { "height": (ValueType.CONTINUOUS, [-10, 10], 1e-2), "width": (ValueType.DISCRETE, [-10, 10], False), "layers": (ValueType.GRID, [4, 8, 16]) } "config": { "iterations": 10, # evaluation times } zoopt_search_config = { "parallel_num": 8, # how many workers to parallel } zoopt_search = ZOOptSearch( algo="Asracos", # only support Asracos currently budget=20, # must match `num_samples` in `tune.TuneConfig()`. dim_dict=dim_dict, metric="mean_loss", mode="min", **zoopt_search_config ) tuner = tune.Tuner( my_objective, tune_config=tune.TuneConfig( search_alg=zoopt_search, num_samples=20 ), run_config=train.RunConfig( name="zoopt_search", stop={"timesteps_total": 10} ), ) tuner.fit() Parameters: algo: To specify an algorithm in zoopt you want to use. Only support ASRacos currently. budget: Number of samples. dim_dict: Dimension dictionary. For continuous dimensions: (continuous, search_range, precision); For discrete dimensions: (discrete, search_range, has_order); For grid dimensions: (grid, grid_list). More details can be found in zoopt package. metric: The training result objective value attribute. If None but a mode was passed, the anonymous metric `_metric` will be used per default. mode: One of {min, max}. Determines whether objective is minimizing or maximizing the metric attribute. points_to_evaluate: Initial parameter suggestions to be run first. This is for when you already have some good parameters you want to run first to help the algorithm make better suggestions for future parameters. Needs to be a list of dicts containing the configurations. parallel_num: How many workers to parallel. Note that initial phase may start less workers than this number. More details can be found in zoopt package. """ optimizer = None def __init__( self, algo: str = "asracos", budget: Optional[int] = None, dim_dict: Optional[Dict] = None, metric: Optional[str] = None, mode: Optional[str] = None, points_to_evaluate: Optional[List[Dict]] = None, parallel_num: int = 1, **kwargs ): assert ( zoopt is not None ), "ZOOpt not found - please install zoopt by `pip install -U zoopt`." assert budget is not None, "`budget` should not be None!" if mode: assert mode in ["min", "max"], "`mode` must be 'min' or 'max'." _algo = algo.lower() assert _algo in [ "asracos", "sracos", ], "`algo` must be in ['asracos', 'sracos'] currently" self._algo = _algo if isinstance(dim_dict, dict) and dim_dict: resolved_vars, domain_vars, grid_vars = parse_spec_vars(dim_dict) if domain_vars or grid_vars: logger.warning( UNRESOLVED_SEARCH_SPACE.format(par="dim_dict", cls=type(self)) ) dim_dict = self.convert_search_space(dim_dict, join=True) self._dim_dict = dim_dict self._budget = budget self._metric = metric if mode == "max": self._metric_op = -1.0 elif mode == "min": self._metric_op = 1.0 self._points_to_evaluate = copy.deepcopy(points_to_evaluate) self._live_trial_mapping = {} self._dim_keys = [] self.solution_dict = {} self.best_solution_list = [] self.optimizer = None self.kwargs = kwargs self.parallel_num = parallel_num super(ZOOptSearch, self).__init__(metric=self._metric, mode=mode) if self._dim_dict: self._setup_zoopt() def _setup_zoopt(self): if self._metric is None and self._mode: # If only a mode was passed, use anonymous metric self._metric = DEFAULT_METRIC _dim_list = [] for k in self._dim_dict: self._dim_keys.append(k) _dim_list.append(self._dim_dict[k]) init_samples = None if self._points_to_evaluate: logger.warning( "`points_to_evaluate` is ignored by ZOOpt in versions <= 0.4.1." ) init_samples = [ Solution(x=tuple(point[dim] for dim in self._dim_keys)) for point in self._points_to_evaluate ] dim = zoopt.Dimension2(_dim_list) par = zoopt.Parameter(budget=self._budget, init_samples=init_samples) if self._algo == "sracos" or self._algo == "asracos": from zoopt.algos.opt_algorithms.racos.sracos import SRacosTune self.optimizer = SRacosTune( dimension=dim, parameter=par, parallel_num=self.parallel_num, **self.kwargs ) def set_search_properties( self, metric: Optional[str], mode: Optional[str], config: Dict, **spec ) -> bool: if self._dim_dict: return False space = self.convert_search_space(config) self._dim_dict = space if metric: self._metric = metric if mode: self._mode = mode if self._mode == "max": self._metric_op = -1.0 elif self._mode == "min": self._metric_op = 1.0 self._setup_zoopt() return True def suggest(self, trial_id: str) -> Optional[Dict]: if not self._dim_dict or not self.optimizer: raise RuntimeError( UNDEFINED_SEARCH_SPACE.format( cls=self.__class__.__name__, space="dim_dict" ) ) if not self._metric or not self._mode: raise RuntimeError( UNDEFINED_METRIC_MODE.format( cls=self.__class__.__name__, metric=self._metric, mode=self._mode ) ) _solution = self.optimizer.suggest() if _solution == "FINISHED": if ray.__version__ >= "0.8.7": return Searcher.FINISHED else: return None if _solution: self.solution_dict[str(trial_id)] = _solution _x = _solution.get_x() new_trial = dict(zip(self._dim_keys, _x)) self._live_trial_mapping[trial_id] = new_trial return unflatten_dict(new_trial)
[docs] def on_trial_complete( self, trial_id: str, result: Optional[Dict] = None, error: bool = False ): """Notification for the completion of trial.""" if result: _solution = self.solution_dict[str(trial_id)] _best_solution_so_far = self.optimizer.complete( _solution, self._metric_op * result[self._metric] ) if _best_solution_so_far: self.best_solution_list.append(_best_solution_so_far) del self._live_trial_mapping[trial_id]
def save(self, checkpoint_path: str): save_object = self.__dict__ with open(checkpoint_path, "wb") as outputFile: pickle.dump(save_object, outputFile) def restore(self, checkpoint_path: str): with open(checkpoint_path, "rb") as inputFile: save_object = pickle.load(inputFile) self.__dict__.update(save_object) @staticmethod def convert_search_space(spec: Dict, join: bool = False) -> Dict[str, Tuple]: spec = copy.deepcopy(spec) resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec) if not domain_vars and not grid_vars: return {} if grid_vars: raise ValueError( "Grid search parameters cannot be automatically converted " "to a ZOOpt search space." ) def resolve_value(domain: Domain) -> Tuple: quantize = None sampler = domain.get_sampler() if isinstance(sampler, Quantized): quantize = sampler.q sampler = sampler.sampler if isinstance(domain, Float): precision = quantize or 1e-12 if isinstance(sampler, Uniform): return ( ValueType.CONTINUOUS, [domain.lower, domain.upper], precision, ) elif isinstance(domain, Integer): if isinstance(sampler, Uniform): return (ValueType.DISCRETE, [domain.lower, domain.upper - 1], True) elif isinstance(domain, Categorical): # Categorical variables would use ValueType.DISCRETE with # has_partial_order=False, however, currently we do not # keep track of category values and cannot automatically # translate back and forth between them. if isinstance(sampler, Uniform): return (ValueType.GRID, domain.categories) raise ValueError( "ZOOpt does not support parameters of type " "`{}` with samplers of type `{}`".format( type(domain).__name__, type(domain.sampler).__name__ ) ) conv_spec = { "/".join(path): resolve_value(domain) for path, domain in domain_vars } if join: spec.update(conv_spec) conv_spec = spec return conv_spec