Source code for ax.benchmark.runners.surrogate

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# pyre-strict

import warnings
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

import torch
from ax.benchmark.runners.base import BenchmarkRunner
from ax.core.arm import Arm
from ax.core.base_trial import BaseTrial, TrialStatus
from ax.core.observation import ObservationFeatures
from ax.core.parameter import RangeParameter
from ax.core.search_space import SearchSpace
from ax.core.types import TParameterization
from ax.modelbridge.transforms.int_to_float import IntToFloat
from ax.modelbridge.transforms.log import Log
from ax.models.torch.botorch_modular.surrogate import Surrogate
from ax.utils.common.serialization import TClassDecoderRegistry, TDecoderRegistry
from ax.utils.common.typeutils import not_none
from botorch.utils.datasets import SupervisedDataset
from torch import Tensor


[docs]class SurrogateRunner(BenchmarkRunner): def __init__( self, name: str, surrogate: Surrogate, datasets: List[SupervisedDataset], search_space: SearchSpace, outcome_names: List[str], noise_stds: Union[float, Dict[str, float]] = 0.0, ) -> None: """Runner for surrogate benchmark problems. Args: name: The name of the runner. surrogate: The modular BoTorch model `Surrogate` to use for generating observations. search_space: The search space of the problem (used for parameter transforms). datasets: The data sets used to fit the surrogate model. outcome_names: The names of the outcomes of the Surrogate. noise_stds: Noise standard deviations to add to the surrogate output(s). If a single float is provided, noise with that standard deviation is added to all outputs. Alternatively, a dictionary mapping outcome names to noise standard deviations can be provided to specify different noise levels for different outputs. """ self.name = name self.surrogate = surrogate self._outcome_names = outcome_names self.datasets = datasets self.search_space = search_space self.noise_stds = noise_stds self.statuses: Dict[int, TrialStatus] = {} # If there are log scale parameters, these need to be transformed. if any( isinstance(p, RangeParameter) and p.log_scale for p in search_space.parameters.values() ): int_to_float_tf = IntToFloat(search_space=search_space) log_tf = Log( search_space=int_to_float_tf.transform_search_space( search_space.clone() ) ) self.transforms: Optional[Tuple[IntToFloat, Log]] = ( int_to_float_tf, log_tf, ) else: self.transforms = None @property def outcome_names(self) -> List[str]: return self._outcome_names def _get_transformed_parameters( self, parameters: TParameterization ) -> TParameterization: if self.transforms is None: return parameters obs_ft = ObservationFeatures(parameters=parameters) for t in not_none(self.transforms): obs_ft = t.transform_observation_features([obs_ft])[0] return obs_ft.parameters
[docs] def get_noise_stds(self) -> Union[None, float, Dict[str, float]]: return self.noise_stds
[docs] def get_Y_true(self, arm: Arm) -> Tensor: X = torch.tensor( [list(self._get_transformed_parameters(arm.parameters).values())], device=self.surrogate.device, dtype=self.surrogate.dtype, ) # We're ignoring the uncertainty predictions of the surrogate model here and # use the mean predictions as the outcomes (before potentially adding noise) means, _ = self.surrogate.predict(X=X) return means.squeeze(0)
[docs] def run(self, trial: BaseTrial) -> Dict[str, Any]: """Run the trial by evaluating its parameterization(s) on the surrogate model. Note: This also sets the status of the trial to COMPLETED. Args: trial: The trial to evaluate. Returns: A dictionary with the following keys: - outcome_names: The names of the metrics being evaluated. - Ys: A dict mapping arm names to lists of corresponding outcomes, where the order of the outcomes is the same as in `outcome_names`. - Ystds: A dict mapping arm names to lists of corresponding outcome noise standard deviations (possibly nan if the noise level is unobserved), where the order of the outcomes is the same as in `outcome_names`. - Ys_true: A dict mapping arm names to lists of corresponding ground truth outcomes, where the order of the outcomes is the same as in `outcome_names`. """ self.statuses[trial.index] = TrialStatus.COMPLETED run_metadata = super().run(trial=trial) run_metadata["outcome_names"] = self.outcome_names return run_metadata
[docs] def poll_trial_status( self, trials: Iterable[BaseTrial] ) -> Dict[TrialStatus, Set[int]]: return {TrialStatus.COMPLETED: {t.index for t in trials}}
[docs] @classmethod # pyre-fixme[2]: Parameter annotation cannot be `Any`. def serialize_init_args(cls, obj: Any) -> Dict[str, Any]: """Serialize the properties needed to initialize the runner. Used for storage. WARNING: Because of issues with consistently saving and loading BoTorch and GPyTorch modules the SurrogateRunner cannot be serialized at this time. At load time the runner will be replaced with a SyntheticRunner. """ warnings.warn( "Because of issues with consistently saving and loading BoTorch and " f"GPyTorch modules, {cls.__name__} cannot be serialized at this time. " "At load time the runner will be replaced with a SyntheticRunner.", stacklevel=3, ) return {}
[docs] @classmethod def deserialize_init_args( cls, args: Dict[str, Any], decoder_registry: Optional[TDecoderRegistry] = None, class_decoder_registry: Optional[TClassDecoderRegistry] = None, ) -> Dict[str, Any]: return {}