From 68319020bf92aa2f6e510f13d9eaba10743d974b Mon Sep 17 00:00:00 2001 From: Sait Cakmak Date: Wed, 6 Nov 2024 16:06:17 -0800 Subject: [PATCH] Clean up unused torch model utils (#3027) Summary: Pull Request resolved: https://github.com/facebook/Ax/pull/3027 Removes a couple of unused utilities form ax/models/torch/utils.py, and updates the tests to remove usage of `HeteroskedasticSingleTaskGP`. Reviewed By: esantorella Differential Revision: D65545235 fbshipit-source-id: 42f564861aa9512feb0ebc445325835b42fc30d1 --- ax/models/tests/test_torch_model_utils.py | 22 ++--- ax/models/torch/utils.py | 110 ---------------------- 2 files changed, 6 insertions(+), 126 deletions(-) diff --git a/ax/models/tests/test_torch_model_utils.py b/ax/models/tests/test_torch_model_utils.py index 102cf25cfb3..692f87991e2 100644 --- a/ax/models/tests/test_torch_model_utils.py +++ b/ax/models/tests/test_torch_model_utils.py @@ -6,18 +6,18 @@ # pyre-strict +from unittest import mock + import numpy as np import torch -from ax.exceptions.model import ModelError from ax.models.torch.utils import ( _generate_sobol_points, - is_noiseless, normalize_indices, subset_model, tensor_callable_to_array_callable, ) from ax.utils.common.testutils import TestCase -from botorch.models import HeteroskedasticSingleTaskGP, SingleTaskGP +from botorch.models import SingleTaskGP from botorch.models.deterministic import GenericDeterministicModel from botorch.models.model import ModelList from botorch.models.model_list_gp_regression import ModelListGP @@ -27,17 +27,6 @@ class TorchUtilsTest(TestCase): - def test_is_noiseless(self) -> None: - x = torch.zeros(1, 1) - y = torch.zeros(1, 1) - se = torch.zeros(1, 1) - model = SingleTaskGP(x, y) - self.assertTrue(is_noiseless(model)) - model = HeteroskedasticSingleTaskGP(x, y, se) - self.assertFalse(is_noiseless(model)) - with self.assertRaises(ModelError): - is_noiseless(ModelListGP()) - def test_NormalizeIndices(self) -> None: indices = [0, 2] nlzd_indices = normalize_indices(indices, 3) @@ -181,8 +170,9 @@ def test_with_obj_thresholds_can_subset(self) -> None: def test_unsupported(self) -> None: yvar = torch.ones(1, 2) - model = HeteroskedasticSingleTaskGP(self.x, self.y, yvar) - subset_model_results = subset_model(model, self.obj_weights) + model = SingleTaskGP(train_X=self.x, train_Y=self.y, train_Yvar=yvar) + with mock.patch.object(model, "subset_output", side_effect=NotImplementedError): + subset_model_results = subset_model(model, self.obj_weights) model_sub = subset_model_results.model obj_weights_sub = subset_model_results.objective_weights ocs_sub = subset_model_results.outcome_constraints diff --git a/ax/models/torch/utils.py b/ax/models/torch/utils.py index 3009ea460d4..11f97c08ff8 100644 --- a/ax/models/torch/utils.py +++ b/ax/models/torch/utils.py @@ -15,7 +15,6 @@ import numpy.typing as npt import torch from ax.exceptions.core import UnsupportedError -from ax.exceptions.model import ModelError from ax.models.model_utils import filter_constraints_and_fixed_features, get_observed from ax.models.random.sobol import SobolGenerator from ax.models.types import TConfig @@ -23,7 +22,6 @@ from ax.utils.common.logger import get_logger from botorch.acquisition.acquisition import AcquisitionFunction from botorch.acquisition.analytic import PosteriorMean -from botorch.acquisition.fixed_feature import FixedFeatureAcquisitionFunction from botorch.acquisition.monte_carlo import ( qSimpleRegret, SampleReducingMCAcquisitionFunction, @@ -51,7 +49,6 @@ ) from botorch.acquisition.risk_measures import RiskMeasureMCObjective from botorch.acquisition.utils import get_infeasible_cost -from botorch.models import ModelListGP, SingleTaskGP from botorch.models.model import Model from botorch.posteriors.fully_bayesian import GaussianMixturePosterior from botorch.posteriors.gpytorch import GPyTorchPosterior @@ -66,9 +63,6 @@ logger: Logger = get_logger(__name__) -NOISELESS_MODELS = {SingleTaskGP} - - # Distributions SIMPLEX = "simplex" HYPERSPHERE = "hypersphere" @@ -83,15 +77,6 @@ class SubsetModelData: indices: Tensor -def is_noiseless(model: Model) -> bool: - """Check if a given (single-task) botorch model is noiseless""" - if isinstance(model, ModelListGP): - raise ModelError( - "Checking for noisless models only applies to sub-models of ModelListGP" - ) - return model.__class__ in NOISELESS_MODELS - - def _filter_X_observed( Xs: list[Tensor], objective_weights: Tensor, @@ -508,101 +493,6 @@ def objective(samples: Tensor, X: Tensor | None = None) -> Tensor: return None, transform -def get_out_of_sample_best_point_acqf( - model: Model, - Xs: list[Tensor], - X_observed: Tensor, - objective_weights: Tensor, - mc_samples: int = 512, - fixed_features: dict[int, float] | None = None, - fidelity_features: list[int] | None = None, - target_fidelities: dict[int, float] | None = None, - outcome_constraints: tuple[Tensor, Tensor] | None = None, - seed_inner: int | None = None, - qmc: bool = True, - risk_measure: RiskMeasureMCObjective | None = None, - **kwargs: Any, -) -> tuple[AcquisitionFunction, list[int] | None]: - """Picks an appropriate acquisition function to find the best - out-of-sample (predicted by the given surrogate model) point - and instantiates it. - - NOTE: Typically the appropriate function is the posterior mean, - but can differ to account for fidelities etc. - """ - model = model - - # subset model only to the outcomes we need for the optimization - if kwargs.get(Keys.SUBSET_MODEL, True): - subset_model_results = subset_model( - model=model, - objective_weights=objective_weights, - outcome_constraints=outcome_constraints, - ) - model = subset_model_results.model - objective_weights = subset_model_results.objective_weights - outcome_constraints = subset_model_results.outcome_constraints - - fixed_features = fixed_features or {} - target_fidelities = target_fidelities or {} - - if fidelity_features: - # we need to optimize at the target fidelities - if any(f in fidelity_features for f in fixed_features): - raise RuntimeError("Fixed features cannot also be fidelity features.") - elif set(fidelity_features) != set(target_fidelities): - raise RuntimeError( - "Must provide a target fidelity for every fidelity feature." - ) - # make sure to not modify fixed_features in-place - fixed_features = {**fixed_features, **target_fidelities} - elif target_fidelities: - raise RuntimeError( - "Must specify fidelity_features in fit() when using target fidelities." - ) - - acqf_class, acqf_options = pick_best_out_of_sample_point_acqf_class( - outcome_constraints=outcome_constraints, - mc_samples=mc_samples, - qmc=qmc, - seed_inner=seed_inner, - risk_measure=risk_measure, - ) - objective, posterior_transform = get_botorch_objective_and_transform( - botorch_acqf_class=acqf_class, - model=model, - objective_weights=objective_weights, - outcome_constraints=outcome_constraints, - X_observed=X_observed, - risk_measure=risk_measure, - ) - - if objective is not None: - if not isinstance(objective, MCAcquisitionObjective): - raise UnsupportedError( - f"Unknown objective type: {objective.__class__}" # pragma: nocover - ) - acqf_options = {"objective": objective, **acqf_options} - if posterior_transform is not None: - acqf_options = {"posterior_transform": posterior_transform, **acqf_options} - - acqf = acqf_class(model=model, **acqf_options) # pyre-ignore [45] - - if fixed_features: - acqf = FixedFeatureAcquisitionFunction( - acq_function=acqf, - d=X_observed.size(-1), - columns=list(fixed_features.keys()), - values=list(fixed_features.values()), - ) - non_fixed_idcs = [i for i in range(Xs[0].size(-1)) if i not in fixed_features] - - else: - non_fixed_idcs = None - - return acqf, non_fixed_idcs - - def pick_best_out_of_sample_point_acqf_class( outcome_constraints: tuple[Tensor, Tensor] | None = None, mc_samples: int = 512,