diff --git a/examples/tune/ax/__init__.py b/examples/tune/ax/__init__.py new file mode 100644 index 00000000..40a96afc --- /dev/null +++ b/examples/tune/ax/__init__.py @@ -0,0 +1 @@ +# -*- coding: utf-8 -*- diff --git a/examples/tune/ax/tune.py b/examples/tune/ax/tune.py new file mode 100644 index 00000000..ddd4c1a5 --- /dev/null +++ b/examples/tune/ax/tune.py @@ -0,0 +1,98 @@ +# -*- coding: utf-8 -*- + +"""A simple example using sklearn and Ax support""" + +# Spock ONLY supports the service style API from Ax +# https://ax.dev/docs/api.html + + +from sklearn.datasets import load_iris +from sklearn.linear_model import LogisticRegression +from sklearn.model_selection import train_test_split + +from spock.addons.tune import ( + AxTunerConfig, + ChoiceHyperParameter, + RangeHyperParameter, + spockTuner, +) +from spock.builder import ConfigArgBuilder +from spock.config import spock + + +@spock +class BasicParams: + n_trials: int + max_iter: int + + +@spockTuner +class LogisticRegressionHP: + c: RangeHyperParameter + solver: ChoiceHyperParameter + + +def main(): + # Load the iris data + X, y = load_iris(return_X_y=True) + + # Split the Iris data + X_train, X_valid, y_train, y_valid = train_test_split(X, y) + + # Ax config -- this will internally spawn the AxClient service API style which will be returned + # by accessing the tuner_status property on the ConfigArgBuilder object + ax_config = AxTunerConfig(objective_name="accuracy", minimize=False) + + # Use the builder to setup + # Call tuner to indicate that we are going to do some HP tuning -- passing in an ax study object + attrs_obj = ( + ConfigArgBuilder( + LogisticRegressionHP, + BasicParams, + desc="Example Logistic Regression Hyper-Parameter Tuning -- Ax Backend", + ) + .tuner(tuner_config=ax_config) + .save(user_specified_path="/tmp/ax") + ) + + # Here we need some of the fixed parameters first so we can just call the generate fnc to grab all the fixed params + # prior to starting the sampling process + fixed_params = attrs_obj.generate() + + # Now we iterate through a bunch of ax trials + for _ in range(fixed_params.BasicParams.n_trials): + # The crux of spock support -- call save w/ the add_tuner_sample flag to write the current draw to file and + # then call sample to return the composed Spockspace of the fixed parameters and the sampled parameters + # Under the hood spock uses the AxClient Ax interface -- thus it handled the underlying call to get the next + # sample and returns the necessary AxClient object in the return dictionary to call 'complete_trial' with the + # associated metrics + hp_attrs = attrs_obj.save( + add_tuner_sample=True, user_specified_path="/tmp/ax" + ).sample() + # Use the currently sampled parameters in a simple LogisticRegression from sklearn + clf = LogisticRegression( + C=hp_attrs.LogisticRegressionHP.c, + solver=hp_attrs.LogisticRegressionHP.solver, + max_iter=hp_attrs.BasicParams.max_iter, + ) + clf.fit(X_train, y_train) + val_acc = clf.score(X_valid, y_valid) + # Get the status of the tuner -- this dict will contain all the objects needed to update + tuner_status = attrs_obj.tuner_status + # Pull the AxClient object and trial index out of the return dictionary and call 'complete_trial' on the + # AxClient object with the correct raw_data that contains the objective name + tuner_status["client"].complete_trial( + trial_index=tuner_status["trial_index"], + raw_data={"accuracy": (val_acc, 0.0)}, + ) + # Always save the current best set of hyper-parameters + attrs_obj.save_best(user_specified_path="/tmp/ax") + + # Grab the best config and metric + best_config, best_metric = attrs_obj.best + print(f"Best HP Config:\n{best_config}") + print(f"Best Metric: {best_metric}") + + +if __name__ == "__main__": + main() diff --git a/examples/tune/ax/tune.yaml b/examples/tune/ax/tune.yaml new file mode 100644 index 00000000..5a56e06d --- /dev/null +++ b/examples/tune/ax/tune.yaml @@ -0,0 +1,15 @@ +################ +# tune.yaml +################ +BasicParams: + n_trials: 10 + max_iter: 150 + +LogisticRegressionHP: + c: + type: float + bounds: [1E-07, 10.0] + log_scale: true + solver: + type: str + choices: ["lbfgs", "saga"] \ No newline at end of file diff --git a/examples/tune/optuna/tune.py b/examples/tune/optuna/tune.py index de645497..e1cd63fc 100644 --- a/examples/tune/optuna/tune.py +++ b/examples/tune/optuna/tune.py @@ -51,10 +51,10 @@ def main(): ConfigArgBuilder( LogisticRegressionHP, BasicParams, - desc="Example Logistic Regression Hyper-Parameter Tuning", + desc="Example Logistic Regression Hyper-Parameter Tuning -- Optuna Backend", ) .tuner(tuner_config=optuna_config) - .save(user_specified_path="/tmp") + .save(user_specified_path="/tmp/optuna") ) # Here we need some of the fixed parameters first so we can just call the generate fnc to grab all the fixed params @@ -68,7 +68,7 @@ def main(): # Under the hood spock uses the define-and-run Optuna interface -- thus it handled the underlying 'ask' call # and returns the necessary trial object in the return dictionary to call 'tell' with the study object hp_attrs = attrs_obj.save( - add_tuner_sample=True, user_specified_path="/tmp" + add_tuner_sample=True, user_specified_path="/tmp/optuna" ).sample() # Use the currently sampled parameters in a simple LogisticRegression from sklearn clf = LogisticRegression( @@ -84,7 +84,7 @@ def main(): # object tuner_status["study"].tell(tuner_status["trial"], val_acc) # Always save the current best set of hyper-parameters - attrs_obj.save_best(user_specified_path="/tmp") + attrs_obj.save_best(user_specified_path="/tmp/optuna") # Grab the best config and metric best_config, best_metric = attrs_obj.best diff --git a/requirements/TUNE_REQUIREMENTS.txt b/requirements/TUNE_REQUIREMENTS.txt index 8be60c6d..f66aed6e 100644 --- a/requirements/TUNE_REQUIREMENTS.txt +++ b/requirements/TUNE_REQUIREMENTS.txt @@ -1,4 +1,5 @@ +mypy_extensions==0.4.3; python_version < '3.8' optuna==2.9.1 -#torchvision -#torch -#ax-platform \ No newline at end of file +torchvision==0.9.1 +torch==1.8.1 +ax-platform==0.2.0 \ No newline at end of file diff --git a/spock/addons/tune/__init__.py b/spock/addons/tune/__init__.py index 6982e392..8a42a960 100644 --- a/spock/addons/tune/__init__.py +++ b/spock/addons/tune/__init__.py @@ -9,6 +9,7 @@ Please refer to the documentation provided in the README.md """ from spock.addons.tune.config import ( + AxTunerConfig, ChoiceHyperParameter, OptunaTunerConfig, RangeHyperParameter, @@ -19,6 +20,7 @@ "builder", "config", "spockTuner", + "AxTunerConfig", "RangeHyperParameter", "ChoiceHyperParameter", "OptunaTunerConfig", diff --git a/spock/addons/tune/ax.py b/spock/addons/tune/ax.py new file mode 100644 index 00000000..30b5de2b --- /dev/null +++ b/spock/addons/tune/ax.py @@ -0,0 +1,159 @@ +# -*- coding: utf-8 -*- + +# Copyright FMR LLC +# SPDX-License-Identifier: Apache-2.0 + +"""Handles the ax backend""" + +from ax.service.ax_client import AxClient + +from spock.addons.tune.config import AxTunerConfig +from spock.addons.tune.interface import BaseInterface + +try: + from typing import TypedDict +except ImportError: + from mypy_extensions import TypedDict + + +class AxTunerStatus(TypedDict): + """Tuner status return object for Ax -- supports the service style API from Ax + + *Attributes*: + + client: current AxClient instance + trial_index: current trial index + + """ + + client: AxClient + trial_index: int + + +class AxInterface(BaseInterface): + """Specific override to support the Ax backend -- supports the service style API from Ax""" + + def __init__(self, tuner_config: AxTunerConfig, tuner_namespace): + """AxInterface init call that maps variables, creates a map to fnc calls, and constructs the necessary + underlying objects + + *Args*: + + tuner_config: configuration object for the ax backend + tuner_namespace: tuner namespace that has attr classes that maps to an underlying library types + """ + super(AxInterface, self).__init__(tuner_config, tuner_namespace) + self._tuner_obj = AxClient( + generation_strategy=self._tuner_config.generation_strategy, + enforce_sequential_optimization=self._tuner_config.enforce_sequential_optimization, + random_seed=self._tuner_config.random_seed, + verbose_logging=self._tuner_config.verbose_logging, + ) + # Some variables to use later + self._trial_index = None + self._sample_hash = None + # Mapping spock underlying classes to ax distributions (search space) + self._map_type = { + "RangeHyperParameter": { + "int": self._ax_range, + "float": self._ax_range, + }, + "ChoiceHyperParameter": { + "int": self._ax_choice, + "float": self._ax_choice, + "str": self._ax_choice, + "bool": self._ax_choice, + }, + } + # Build the correct underlying dictionary object for Ax client create experiment + self._param_obj = self._construct() + # Create the AxClient experiment + self._tuner_obj.create_experiment( + parameters=self._param_obj, + name=self._tuner_config.name, + objective_name=self._tuner_config.objective_name, + minimize=self._tuner_config.minimize, + parameter_constraints=self._tuner_config.parameter_constraints, + outcome_constraints=self._tuner_config.outcome_constraints, + overwrite_existing_experiment=self._tuner_config.overwrite_existing_experiment, + tracking_metric_names=self._tuner_config.tracking_metric_names, + immutable_search_space_and_opt_config=self._tuner_config.immutable_search_space_and_opt_config, + is_test=self._tuner_config.is_test, + ) + + @property + def tuner_status(self) -> AxTunerStatus: + return AxTunerStatus(client=self._tuner_obj, trial_index=self._trial_index) + + @property + def best(self): + best_obj = self._tuner_obj.get_best_parameters() + rollup_dict, _ = self._sample_rollup(best_obj[0]) + return ( + self._gen_spockspace(rollup_dict), + best_obj[1][0][self._tuner_obj.objective_name], + ) + + @property + def _get_sample(self): + return self._tuner_obj.get_next_trial() + + def sample(self): + parameters, self._trial_index = self._get_sample + # Roll this back out into a Spockspace so it can be merged into the fixed parameter Spockspace + # Also need to un-dot the param names to rebuild the nested structure + rollup_dict, sample_hash = self._sample_rollup(parameters) + self._sample_hash = sample_hash + return self._gen_spockspace(rollup_dict) + + def _construct(self): + param_list = [] + # These will only be nested one level deep given the tuner syntax + for k, v in vars(self._tuner_namespace).items(): + for ik, iv in vars(v).items(): + param_fn = self._map_type[type(iv).__name__][iv.type] + param_list.append(param_fn(name=f"{k}.{ik}", val=iv)) + return param_list + + def _ax_range(self, name, val): + """Assemble the dictionary for ax range parameters + + *Args*: + + name: parameter name + val: current attr val + + *Returns*: + + dictionary that can be added to a parameter list + + """ + low, high = self._try_range_cast(val, type_string="RangeHyperParameter") + return { + "name": name, + "type": "range", + "bounds": [low, high], + "value_type": val.type, + "log_scale": val.log_scale, + } + + def _ax_choice(self, name, val): + """Assemble the dictionary for ax choice parameters + + *Args*: + + name: parameter name + val: current attr val + + *Returns*: + + dictionary that can be added to a parameter list + + """ + val = self._try_choice_cast(val, type_string="ChoiceHyperParameter") + return { + "name": name, + "type": "choice", + "values": val.choices, + "value_type": val.type, + } diff --git a/spock/addons/tune/config.py b/spock/addons/tune/config.py index 5974f789..982ffdb0 100644 --- a/spock/addons/tune/config.py +++ b/spock/addons/tune/config.py @@ -6,19 +6,39 @@ """Creates the spock config interface that wraps attr -- tune version for hyper-parameters""" import sys from typing import List, Optional, Sequence, Tuple, Union +from uuid import uuid4 import attr import optuna +from ax.modelbridge.generation_strategy import GenerationStrategy from spock.backend.config import _base_attr +@attr.s(auto_attribs=True) +class AxTunerConfig: + objective_name: str + tracking_metric_names: Optional[List[str]] = None + name: Optional[str] = f"spock_ax_{uuid4()}" + minimize: bool = True + parameter_constraints: Optional[List[str]] = None + outcome_constraints: Optional[List[str]] = None + support_intermediate_data: bool = False + overwrite_existing_experiment: bool = False + immutable_search_space_and_opt_config: bool = True + is_test: bool = False + generation_strategy: Optional[GenerationStrategy] = None + enforce_sequential_optimization: bool = True + random_seed: Optional[int] = None + verbose_logging: bool = True + + @attr.s(auto_attribs=True) class OptunaTunerConfig: storage: Optional[Union[str, optuna.storages.BaseStorage]] = None sampler: Optional[optuna.samplers.BaseSampler] = None pruner: Optional[optuna.pruners.BasePruner] = None - study_name: Optional[str] = None + study_name: Optional[str] = f"spock_optuna_{uuid4()}" direction: Optional[Union[str, optuna.study.StudyDirection]] = None load_if_exists: bool = False directions: Optional[Sequence[Union[str, optuna.study.StudyDirection]]] = None diff --git a/spock/addons/tune/interface.py b/spock/addons/tune/interface.py index eac45dbf..4a88a85c 100644 --- a/spock/addons/tune/interface.py +++ b/spock/addons/tune/interface.py @@ -4,28 +4,37 @@ # SPDX-License-Identifier: Apache-2.0 """Handles the base interface""" +import hashlib +import json from abc import ABC, abstractmethod -from typing import Dict +from typing import Dict, Union import attr +from spock.addons.tune.config import AxTunerConfig, OptunaTunerConfig from spock.backend.wrappers import Spockspace class BaseInterface(ABC): + """Base interface for the various hyper-parameter tuner backends + + *Attributes* + + _tuner_config: spock version of the tuner configuration + _tuner_namespace: tuner namespace that has attr classes that maps to an underlying library types + + """ + def __init__(self, tuner_config, tuner_namespace: Spockspace): """Base init call that maps a few variables *Args*: - _tuner_config: necessary object to determine the interface and sample correctly from the underlying library - _tuner_namespace: tuner namespace that has attr classes that maps to an underlying library types + tuner_config: necessary dict object to determine the interface and sample correctly from the underlying library + tuner_namespace: tuner namespace that has attr classes that maps to an underlying library types """ - - self._tuner_config = { - k: v for k, v in attr.asdict(tuner_config).items() if v is not None - } + self._tuner_config = tuner_config self._tuner_namespace = tuner_namespace @abstractmethod @@ -47,13 +56,65 @@ def _construct(self): *Returns*: - Any typed object needed for support + flat dictionary of all hyper-parameters named with dot notation (class.param_name) """ pass + @property + @abstractmethod + def _get_sample(self): + """Gets the sample parameter dictionary from the underlying backend""" + pass + + @property + @abstractmethod + def tuner_status(self): + """Returns a dictionary of all the necessary underlying tuner internals to report the result""" + pass + + @property + @abstractmethod + def best(self): + """Returns a Spockspace of the best hyper-parameter config and the associated metric value""" + @staticmethod - def _gen_attr_classes(tune_dict: Dict): + def _sample_rollup(params): + """Rollup the sample draw into a dictionary that can be converted to a spockspace with the correct names and + roots -- un-dots the name structure + + *Args*: + + params: current parameter dictionary -- named by dot notation + + *Returns*: + + dictionary of rolled up sampled parameters + md5 hash of the dictionary contents + + """ + key_set = {k.split(".")[0] for k in params.keys()} + rollup_dict = {val: {} for val in key_set} + for k, v in params.items(): + split_names = k.split(".") + rollup_dict[split_names[0]].update({split_names[1]: v}) + dict_hash = hashlib.md5( + json.dumps(rollup_dict, sort_keys=True).encode("utf-8") + ).digest() + return rollup_dict, dict_hash + + def _gen_spockspace(self, tune_dict: Dict): + """Converts a dictionary of dictionaries of parameters into a valid Spockspace + + *Args*: + + tune_dict: dictionary of current parameters + + Returns: + + tune_dict: Spockspace + + """ for k, v in tune_dict.items(): attrs_dict = { ik: attr.ib( @@ -63,7 +124,21 @@ def _gen_attr_classes(tune_dict: Dict): } obj = attr.make_class(name=k, attrs=attrs_dict, kw_only=True, frozen=True) tune_dict.update({k: obj(**v)}) - return tune_dict + return self._to_spockspace(tune_dict) + + @staticmethod + def _config_to_dict(tuner_config: Union[OptunaTunerConfig, AxTunerConfig]): + """Turns an attrs config object into a dictionary + + *Args*: + + tuner_config: attrs config object + + *Returns*: + + dictionary of the attrs config object + """ + return {k: v for k, v in attr.asdict(tuner_config).items() if v is not None} @staticmethod def _to_spockspace(tune_dict: Dict): @@ -95,13 +170,49 @@ def _get_caster(val): """ return __builtins__[val.type] - @property - @abstractmethod - def tuner_status(self): - """Returns a dictionary of all the necessary underlying tuner internals to report the result""" - pass + def _try_choice_cast(self, val, type_string: str): + """Try/except for casting choice parameters - @property - @abstractmethod - def best(self): - """Returns a Spockspace of the best hyper-parameter config and the associated metric value""" + *Args*: + + val: current attr val + type_string: spock hyper-parameter type name + + *Returns*: + + val: updated attr val + + """ + caster = self._get_caster(val) + # Just attempt to cast in a try except + try: + val.choices = [caster(v) for v in val.choices] + return val + except TypeError: + print( + f"Attempted to cast into type: {val.type} but failed -- check the inputs to {type_string}" + ) + + def _try_range_cast(self, val, type_string: str): + """Try/except for casting range parameters + + *Args*: + + val: current attr val + type_string: spock hyper-parameter type name + + *Returns*: + + low: low bound + high: high bound + + """ + caster = self._get_caster(val) + try: + low = caster(val.bounds[0]) + high = caster(val.bounds[1]) + return low, high + except TypeError: + print( + f"Attempted to cast into type: {val.type} but failed -- check the inputs to {type_string}" + ) diff --git a/spock/addons/tune/optuna.py b/spock/addons/tune/optuna.py index 9b80e298..ab00d08c 100644 --- a/spock/addons/tune/optuna.py +++ b/spock/addons/tune/optuna.py @@ -5,26 +5,43 @@ """Handles the optuna backend""" -import hashlib -import json -from warnings import warn -import attr import optuna from spock.addons.tune.config import OptunaTunerConfig from spock.addons.tune.interface import BaseInterface +try: + from typing import TypedDict +except ImportError: + from mypy_extensions import TypedDict + + +class OptunaTunerStatus(TypedDict): + """Tuner status return object for Optuna -- supports the define-and-run style interface from Optuna + + *Attributes*: + + trial: current ask trial sample + study: current optuna study object + + """ + + trial: optuna.Trial + study: optuna.Study + class OptunaInterface(BaseInterface): - """Specific override to support the optuna backend + """Specific override to support the optuna backend -- supports the define-and-run style interface from Optuna *Attributes*: _map_type: dictionary that maps class names and types to fns that create optuna distributions - _tuner_obj: necessary object to determine the interface and sample correctly from the underlying library + _trial: current trial object from the optuna backend + _tuner_obj: underlying optuna study object _tuner_namespace: tuner namespace that has attr classes that maps to an underlying library types _param_obj: underlying object that optuna study can sample from (flat dictionary) + _sample_hash: hash of the most recent sample draw """ @@ -34,15 +51,17 @@ def __init__(self, tuner_config: OptunaTunerConfig, tuner_namespace): *Args*: - tuner_config: necessary object to determine the interface and sample correctly from the underlying library + tuner_config: configuration object for the optuna backend tuner_namespace: tuner namespace that has attr classes that maps to an underlying library types """ super(OptunaInterface, self).__init__(tuner_config, tuner_namespace) - self._tuner_obj = optuna.create_study(**self._tuner_config) + self._tuner_obj = optuna.create_study( + **self._config_to_dict(self._tuner_config) + ) + # Some variables to use later self._trial = None self._sample_hash = None - self._trial_status_hash = None # Mapping spock underlying classes to optuna distributions (define-and-run interface) self._map_type = { "RangeHyperParameter": { @@ -60,54 +79,30 @@ def __init__(self, tuner_config: OptunaTunerConfig, tuner_namespace): self._param_obj = self._construct() @property - def tuner_status(self): - return {"trial": self._trial, "study": self._tuner_obj} + def tuner_status(self) -> OptunaTunerStatus: + return OptunaTunerStatus(trial=self._trial, study=self._tuner_obj) @property def best(self): - rollup_dict, _ = self._trial_rollup(self._tuner_obj.best_trial) + rollup_dict, _ = self._sample_rollup(self._tuner_obj.best_trial.params) return ( - self._to_spockspace(self._gen_attr_classes(rollup_dict)), + self._gen_spockspace(rollup_dict), self._tuner_obj.best_value, ) + @property + def _get_sample(self): + return self._tuner_obj.ask(self._param_obj) + def sample(self): - self._trial = self._tuner_obj.ask(self._param_obj) + self._trial = self._get_sample # Roll this back out into a Spockspace so it can be merged into the fixed parameter Spockspace # Also need to un-dot the param names to rebuild the nested structure - rollup_dict, sample_hash = self._trial_rollup(self._trial) + rollup_dict, sample_hash = self._sample_rollup(self._trial.params) self._sample_hash = sample_hash - return self._to_spockspace(self._gen_attr_classes(rollup_dict)) - - @staticmethod - def _trial_rollup(trial): - """Rollup the trial into a dictionary that can be converted to a spockspace with the correct names and roots - - *Returns*: - - dictionary of rolled up sampled parameters - md5 hash of the dictionary contents - - """ - key_set = {k.split(".")[0] for k in trial.params.keys()} - rollup_dict = {val: {} for val in key_set} - for k, v in trial.params.items(): - split_names = k.split(".") - rollup_dict[split_names[0]].update({split_names[1]: v}) - dict_hash = hashlib.md5( - json.dumps(rollup_dict, sort_keys=True).encode("utf-8") - ).digest() - return rollup_dict, dict_hash + return self._gen_spockspace(rollup_dict) def _construct(self): - """Constructs the base object needed by the underlying library to construct the correct object that allows - for hyper-parameter sampling - - *Returns*: - - flat dictionary of all hyper-parameters named with dot notation (class.param_name) - - """ optuna_dict = {} # These will only be nested one level deep given the tuner syntax for k, v in vars(self._tuner_namespace).items(): @@ -116,8 +111,7 @@ def _construct(self): optuna_dict.update({f"{k}.{ik}": param_fn(iv)}) return optuna_dict - @staticmethod - def _uniform_float_dist(val): + def _uniform_float_dist(self, val): """Assemble the optuna.distributions.(Log)UniformDistribution object https://optuna.readthedocs.io/en/stable/reference/generated/optuna.distributions.UniformDistribution.html @@ -132,13 +126,7 @@ def _uniform_float_dist(val): optuna.distributions.UniformDistribution or optuna.distributions.LogUniformDistribution """ - try: - low = float(val.bounds[0]) - high = float(val.bounds[1]) - except TypeError: - print( - f"Attempted to cast into type: {val.type} but failed -- check the inputs to RangeHyperParameter" - ) + low, high = self._try_range_cast(val, type_string="RangeHyperParameter") log_scale = val.log_scale return ( optuna.distributions.LogUniformDistribution(low=low, high=high) @@ -146,8 +134,7 @@ def _uniform_float_dist(val): else optuna.distributions.UniformDistribution(low=low, high=high) ) - @staticmethod - def _uniform_int_dist(val): + def _uniform_int_dist(self, val): """Assemble the optuna.distributions.Int(Log)UniformDistribution object https://optuna.readthedocs.io/en/stable/reference/generated/optuna.distributions.IntUniformDistribution.html @@ -162,13 +149,7 @@ def _uniform_int_dist(val): optuna.distributions.IntUniformDistribution or optuna.distributions.IntLogUniformDistribution """ - try: - low = int(val.bounds[0]) - high = int(val.bounds[1]) - except TypeError: - print( - f"Attempted to cast into type: {val.type} but failed -- check the inputs to RangeHyperParameter" - ) + low, high = self._try_range_cast(val, type_string="RangeHyperParameter") log_scale = val.log_scale return ( optuna.distributions.IntLogUniformDistribution(low=low, high=high) @@ -190,12 +171,5 @@ def _categorical_dist(self, val): optuna.distributions.CategoricalDistribution """ - caster = self._get_caster(val) - # Just attempt to cast in a try except - try: - val.choices = [caster(v) for v in val.choices] - except TypeError: - print( - f"Attempted to cast into type: {val.type} but failed -- check the inputs to ChoiceHyperParameter" - ) + val = self._try_choice_cast(val, type_string="ChoiceHyperParameter") return optuna.distributions.CategoricalDistribution(choices=val.choices) diff --git a/spock/addons/tune/tuner.py b/spock/addons/tune/tuner.py index 438e4b97..14cb3269 100644 --- a/spock/addons/tune/tuner.py +++ b/spock/addons/tune/tuner.py @@ -7,7 +7,8 @@ from typing import Union -from spock.addons.tune.config import OptunaTunerConfig +from spock.addons.tune.ax import AxInterface +from spock.addons.tune.config import AxTunerConfig, OptunaTunerConfig from spock.addons.tune.optuna import OptunaInterface from spock.backend.wrappers import Spockspace @@ -25,7 +26,7 @@ class TunerInterface: def __init__( self, - tuner_config: Union[OptunaTunerConfig], + tuner_config: Union[OptunaTunerConfig, AxTunerConfig], tuner_namespace: Spockspace, fixed_namespace: Spockspace, ): @@ -39,20 +40,20 @@ def __init__( """ self._fixed_namespace = fixed_namespace - # Todo: add ax type check here - accept_types = OptunaTunerConfig - if not isinstance(tuner_config, accept_types): - raise TypeError( - f"Passed incorrect tuner_config type of {type(tuner_config)} -- must be of type " - f"{repr(accept_types)}" - ) + accept_types = (OptunaTunerConfig, AxTunerConfig) if isinstance(tuner_config, OptunaTunerConfig): self._lib_interface = OptunaInterface( tuner_config=tuner_config, tuner_namespace=tuner_namespace ) - # # TODO: Add ax class logic - # elif isinstance(tuner_config, (ax.Experiment, ax.SimpleExperiment)): - # pass + elif isinstance(tuner_config, AxTunerConfig): + self._lib_interface = AxInterface( + tuner_config=tuner_config, tuner_namespace=tuner_namespace + ) + else: + raise TypeError( + f"Passed incorrect tuner_config type of {type(tuner_config)} -- must be of type " + f"{repr(accept_types)}" + ) def sample(self): """Public interface to underlying library sepcific sample that returns a single sample/draw from the diff --git a/tests/tune/base_asserts_test.py b/tests/tune/base_asserts_test.py index 5a3c1a3e..63ef9eed 100644 --- a/tests/tune/base_asserts_test.py +++ b/tests/tune/base_asserts_test.py @@ -31,7 +31,7 @@ def test_hp_two(self, arg_builder): class SampleTypes: def test_sampling(self, arg_builder): # Draw 100 random samples and make sure all fall within all of the bounds or sets - for _ in range(100): + for _ in range(25): hp_attrs = arg_builder.sample() assert 10 <= hp_attrs.HPOne.hp_int <= 100 assert isinstance(hp_attrs.HPOne.hp_int, int) is True diff --git a/tests/tune/test_ax.py b/tests/tune/test_ax.py new file mode 100644 index 00000000..09dca316 --- /dev/null +++ b/tests/tune/test_ax.py @@ -0,0 +1,176 @@ +# -*- coding: utf-8 -*- +import datetime +from tests.tune.base_asserts_test import * +from tests.tune.attr_configs_test import * +import pytest +import os +import re +import sys +from spock.builder import ConfigArgBuilder +from spock.addons.tune import AxTunerConfig +from sklearn.datasets import load_iris +from sklearn.linear_model import LogisticRegression +from sklearn.model_selection import train_test_split + + +class TestAxBasic(AllTypes): + @staticmethod + @pytest.fixture + def arg_builder(monkeypatch): + with monkeypatch.context() as m: + m.setattr(sys, 'argv', ['', '--config', + './tests/conf/yaml/test_hp.yaml']) + ax_config = AxTunerConfig(name="Basic Test", minimize=False, objective_name="None", verbose_logging=False) + config = ConfigArgBuilder(HPOne, HPTwo).tuner(ax_config) + return config + + +class TestAxCompose(AllTypes): + @staticmethod + @pytest.fixture + def arg_builder(monkeypatch): + with monkeypatch.context() as m: + m.setattr(sys, 'argv', ['', '--config', + './tests/conf/yaml/test_hp_compose.yaml']) + ax_config = AxTunerConfig(name="Basic Test", minimize=False, objective_name="None", verbose_logging=False) + config = ConfigArgBuilder(HPOne, HPTwo).tuner(ax_config) + return config + + def test_hp_one(self, arg_builder): + assert arg_builder._tune_namespace.HPOne.hp_int.bounds == (20, 200) + assert arg_builder._tune_namespace.HPOne.hp_int.type == 'int' + assert arg_builder._tune_namespace.HPOne.hp_int.log_scale is False + assert arg_builder._tune_namespace.HPOne.hp_int_log.bounds == (10, 100) + assert arg_builder._tune_namespace.HPOne.hp_int_log.type == 'int' + assert arg_builder._tune_namespace.HPOne.hp_int_log.log_scale is True + assert arg_builder._tune_namespace.HPOne.hp_float.bounds == (10.0, 100.0) + assert arg_builder._tune_namespace.HPOne.hp_float.type == 'float' + assert arg_builder._tune_namespace.HPOne.hp_float.log_scale is False + assert arg_builder._tune_namespace.HPOne.hp_float_log.bounds == (10.0, 100.0) + assert arg_builder._tune_namespace.HPOne.hp_float_log.type == 'float' + assert arg_builder._tune_namespace.HPOne.hp_float_log.log_scale is True + + +class TestAxSample(SampleTypes): + @staticmethod + @pytest.fixture + def arg_builder(monkeypatch): + with monkeypatch.context() as m: + m.setattr(sys, 'argv', ['', '--config', + './tests/conf/yaml/test_hp.yaml']) + ax_config = AxTunerConfig(name="Basic Test", minimize=False, objective_name="None", verbose_logging=False) + config = ConfigArgBuilder(HPOne, HPTwo).tuner(ax_config) + return config + + +class TestAxSaveTopLevel: + def test_save_top_level(self, monkeypatch): + with monkeypatch.context() as m: + m.setattr(sys, 'argv', ['', '--config', + './tests/conf/yaml/test_optuna.yaml']) + # Optuna config -- this will internally spawn the study object for the define-and-run style which will be returned + # as part of the call to sample() + ax_config = AxTunerConfig( + name="Iris Logistic Regression Tests", minimize=False, objective_name="None" + ) + now = datetime.datetime.now() + curr_int_time = int(f'{now.year}{now.month}{now.day}{now.hour}{now.second}') + config = ConfigArgBuilder(LogisticRegressionHP).tuner(ax_config).save( + user_specified_path="/tmp", file_name=f'pytest.{curr_int_time}', + ).sample() + # Verify the sample was written out to file + yaml_regex = re.compile(fr'pytest.{curr_int_time}.' + fr'[a-fA-F0-9]{{8}}-[a-fA-F0-9]{{4}}-[a-fA-F0-9]{{4}}-' + fr'[a-fA-F0-9]{{4}}-[a-fA-F0-9]{{12}}.spock.cfg.yaml') + matches = [re.fullmatch(yaml_regex, val) for val in os.listdir('/tmp') + if re.fullmatch(yaml_regex, val) is not None] + fname = f'/tmp/{matches[0].string}' + assert os.path.exists(fname) + with open(fname, 'r') as fin: + print(fin.read()) + # Clean up if assert is good + if os.path.exists(fname): + os.remove(fname) + return config + + +class TestIrisAx: + @staticmethod + @pytest.fixture + def arg_builder(monkeypatch): + with monkeypatch.context() as m: + m.setattr(sys, 'argv', ['', '--config', + './tests/conf/yaml/test_optuna.yaml']) + # Ax config -- this will internally spawn the AxClient service API style which will be returned + # by accessing the tuner_status property on the ConfigArgBuilder object + ax_config = AxTunerConfig( + name="Iris Logistic Regression Tests", minimize=False, objective_name="accuracy" + ) + config = ConfigArgBuilder(LogisticRegressionHP).tuner(ax_config) + return config + + def test_iris(self, arg_builder): + # Load the iris data + X, y = load_iris(return_X_y=True) + # Split the Iris data + X_train, X_valid, y_train, y_valid = train_test_split(X, y) + + # Now we iterate through a bunch of ax samples + for _ in range(10): + # The crux of spock support -- call save w/ the add_tuner_sample flag to write the current draw to file and + # then call sample to return the composed Spockspace of the fixed parameters and the sampled parameters + # Under the hood spock uses the AxClient Ax interface -- thus it handled the underlying call to get the next + # sample and returns the necessary AxClient object in the return dictionary to call 'complete_trial' with the + # associated metrics + now = datetime.datetime.now() + curr_int_time = int(f'{now.year}{now.month}{now.day}{now.hour}{now.second}') + hp_attrs = arg_builder.save( + add_tuner_sample=True, user_specified_path="/tmp", file_name=f'pytest.{curr_int_time}', + ).sample() + # Use the currently sampled parameters in a simple LogisticRegression from sklearn + clf = LogisticRegression( + C=hp_attrs.LogisticRegressionHP.c, + solver=hp_attrs.LogisticRegressionHP.solver, + ) + clf.fit(X_train, y_train) + val_acc = clf.score(X_valid, y_valid) + # Get the status of the tuner -- this dict will contain all the objects needed to update + tuner_status = arg_builder.tuner_status + # Pull the AxClient object and trial index out of the return dictionary and call 'complete_trial' on the + # AxClient object with the correct raw_data that contains the objective name + tuner_status["client"].complete_trial( + trial_index=tuner_status["trial_index"], + raw_data={"accuracy": (val_acc, 0.0)}, + ) + # Always save the current best set of hyper-parameters + arg_builder.save_best(user_specified_path='/tmp', file_name=f'pytest') + # Verify the sample was written out to file + yaml_regex = re.compile(fr'pytest.{curr_int_time}.hp.sample.[0-9]+.' + fr'[a-fA-F0-9]{{8}}-[a-fA-F0-9]{{4}}-[a-fA-F0-9]{{4}}-' + fr'[a-fA-F0-9]{{4}}-[a-fA-F0-9]{{12}}.spock.cfg.yaml') + matches = [re.fullmatch(yaml_regex, val) for val in os.listdir('/tmp') + if re.fullmatch(yaml_regex, val) is not None] + fname = f'/tmp/{matches[0].string}' + assert os.path.exists(fname) + with open(fname, 'r') as fin: + print(fin.read()) + # Clean up if assert is good + if os.path.exists(fname): + os.remove(fname) + + best_config, best_metric = arg_builder.best + print(f'Best HP Config:\n{best_config}') + print(f'Best Metric: {best_metric}') + # Verify the sample was written out to file + yaml_regex = re.compile(fr'pytest.hp.best.' + fr'[a-fA-F0-9]{{8}}-[a-fA-F0-9]{{4}}-[a-fA-F0-9]{{4}}-' + fr'[a-fA-F0-9]{{4}}-[a-fA-F0-9]{{12}}.spock.cfg.yaml') + matches = [re.fullmatch(yaml_regex, val) for val in os.listdir('/tmp') + if re.fullmatch(yaml_regex, val) is not None] + fname = f'/tmp/{matches[0].string}' + assert os.path.exists(fname) + with open(fname, 'r') as fin: + print(fin.read()) + # Clean up if assert is good + if os.path.exists(fname): + os.remove(fname) diff --git a/tests/tune/test_raises.py b/tests/tune/test_raises.py index db660449..e43e93b3 100644 --- a/tests/tune/test_raises.py +++ b/tests/tune/test_raises.py @@ -3,6 +3,7 @@ import pytest import sys from spock.builder import ConfigArgBuilder +from spock.addons.tune import AxTunerConfig, OptunaTunerConfig import optuna @@ -16,21 +17,41 @@ def test_incorrect_tuner_config(self, monkeypatch): config = ConfigArgBuilder(HPOne, HPTwo).tuner(optuna_config) -class TestInvalidCastChoice: +class TestOptunaInvalidCastChoice: def test_invalid_cast_choice(self, monkeypatch): with monkeypatch.context() as m: m.setattr(sys, 'argv', ['', '--config', './tests/conf/yaml/test_hp_cast.yaml']) - optuna_config = optuna.create_study(study_name="Tests", direction='minimize') - with pytest.raises(TypeError): + optuna_config = OptunaTunerConfig(study_name="Basic Tests", direction="maximize") + with pytest.raises(ValueError): config = ConfigArgBuilder(HPOne, HPTwo).tuner(optuna_config) -class TestInvalidCastRange: +class TestOptunaInvalidCastRange: def test_invalid_cast_range(self, monkeypatch): with monkeypatch.context() as m: m.setattr(sys, 'argv', ['', '--config', './tests/conf/yaml/test_hp_cast_bounds.yaml']) - optuna_config = optuna.create_study(study_name="Tests", direction='minimize') + optuna_config = OptunaTunerConfig(study_name="Basic Tests", direction="maximize") with pytest.raises(ValueError): config = ConfigArgBuilder(HPOne, HPTwo).tuner(optuna_config) + + +class TestAxInvalidCastChoice: + def test_invalid_cast_choice(self, monkeypatch): + with monkeypatch.context() as m: + m.setattr(sys, 'argv', ['', '--config', + './tests/conf/yaml/test_hp_cast.yaml']) + ax_config = AxTunerConfig(name="Basic Test", minimize=False, objective_name="None", verbose_logging=False) + with pytest.raises(ValueError): + config = ConfigArgBuilder(HPOne, HPTwo).tuner(ax_config) + + +class TestAxInvalidCastRange: + def test_invalid_cast_range(self, monkeypatch): + with monkeypatch.context() as m: + m.setattr(sys, 'argv', ['', '--config', + './tests/conf/yaml/test_hp_cast_bounds.yaml']) + ax_config = AxTunerConfig(name="Basic Test", minimize=False, objective_name="None", verbose_logging=False) + with pytest.raises(ValueError): + config = ConfigArgBuilder(HPOne, HPTwo).tuner(ax_config)