def get_NEI(
model: Model,
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
X_observed: Optional[Tensor] = None,
X_pending: Optional[Tensor] = None,
**kwargs: Any,
) -> AcquisitionFunction:
r"""Instantiates a qNoisyExpectedImprovement acquisition function.
Args:
objective_weights: The objective is to maximize a weighted sum of
the columns of f(x). These are the weights.
outcome_constraints: A tuple of (A, b). For k outcome constraints
and m outputs at f(x), A is (k x m) and b is (k x 1) such that
A f(x) <= b. (Not used by single task models)
X_observed: A tensor containing points observed for all objective
outcomes and outcomes that appear in the outcome constraints (if
there are any).
X_pending: A tensor containing points whose evaluation is pending (i.e.
that have been submitted for evaluation) present for all objective
outcomes and outcomes that appear in the outcome constraints (if
there are any).
mc_samples: The number of MC samples to use (default: 512).
qmc: If True, use qMC instead of MC (default: True).
prune_baseline: If True, prune the baseline points for NEI (default: True).
Returns:
qNoisyExpectedImprovement: The instantiated acquisition function.
"""
if X_observed is None:
raise ValueError("There are no feasible observed points.")
# Parse random_scalarization params
objective_weights = _extract_random_scalarization_settings(
objective_weights, outcome_constraints, **kwargs)
# construct Objective module
if outcome_constraints is None:
objective = LinearMCObjective(weights=objective_weights)
else:
obj_tf = get_objective_weights_transform(objective_weights)
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
X_observed = torch.as_tensor(X_observed)
inf_cost = get_infeasible_cost(X=X_observed,
model=model,
objective=obj_tf)
objective = ConstrainedMCObjective(objective=obj_tf,
constraints=con_tfs or [],
infeasible_cost=inf_cost)
return get_acquisition_function(
acquisition_function_name="qNEI",
model=model,
objective=objective,
X_observed=X_observed,
X_pending=X_pending,
prune_baseline=kwargs.get("prune_baseline", True),
mc_samples=kwargs.get("mc_samples", 512),
qmc=kwargs.get("qmc", True),
seed=torch.randint(1, 10000, (1, )).item(),
)
def get_PosteriorMean(
model: Model,
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
X_observed: Optional[Tensor] = None,
X_pending: Optional[Tensor] = None,
**kwargs: Any,
) -> AcquisitionFunction:
r"""Instantiates a PosteriorMean acquisition function.
Note: If no OutcomeConstraints given, return an analytic acquisition
function. This requires {optimizer_kwargs: {joint_optimization: True}} or an
optimizer that does not assume pending point support.
Args:
objective_weights: The objective is to maximize a weighted sum of
the columns of f(x). These are the weights.
outcome_constraints: A tuple of (A, b). For k outcome constraints
and m outputs at f(x), A is (k x m) and b is (k x 1) such that
A f(x) <= b. (Not used by single task models)
X_observed: A tensor containing points observed for all objective
outcomes and outcomes that appear in the outcome constraints (if
there are any).
X_pending: A tensor containing points whose evaluation is pending (i.e.
that have been submitted for evaluation) present for all objective
outcomes and outcomes that appear in the outcome constraints (if
there are any).
Returns:
PosteriorMean: The instantiated acquisition function.
"""
if X_observed is None:
raise ValueError("There are no feasible observed points.")
# construct Objective module
if kwargs.get("chebyshev_scalarization", False):
obj_tf = get_chebyshev_scalarization(
weights=objective_weights,
Y=torch.stack(kwargs.get("Ys")).transpose(0, 1).squeeze(-1),
)
else:
obj_tf = get_objective_weights_transform(objective_weights)
def obj_fn(samples: Tensor, X: Optional[Tensor] = None) -> Tensor:
return obj_tf(samples)
if outcome_constraints is None:
objective = GenericMCObjective(objective=obj_fn)
else:
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
inf_cost = get_infeasible_cost(X=X_observed,
model=model,
objective=obj_fn)
objective = ConstrainedMCObjective(objective=obj_fn,
constraints=con_tfs or [],
infeasible_cost=inf_cost)
# Use qSimpleRegret, not analytic posterior, to handle arbitrary objective fns.
acq_func = qSimpleRegret(model, objective=objective)
return acq_func
def test_BotorchModel(self, dtype=torch.float, cuda=False):
Xs1, Ys1, Yvars1, bounds, task_features, feature_names = _get_torch_test_data(
dtype=dtype, cuda=cuda, constant_noise=True)
Xs2, Ys2, Yvars2, _, _, _ = _get_torch_test_data(dtype=dtype,
cuda=cuda,
constant_noise=True)
model = BotorchModel()
# Test ModelListGP
# make training data different for each output
Xs2_diff = [Xs2[0] + 0.1]
with mock.patch(FIT_MODEL_MO_PATH) as _mock_fit_model:
model.fit(
Xs=Xs1 + Xs2_diff,
Ys=Ys1 + Ys2,
Yvars=Yvars1 + Yvars2,
bounds=bounds,
task_features=task_features,
feature_names=feature_names,
fidelity_features=[],
)
_mock_fit_model.assert_called_once()
# Check attributes
self.assertTrue(torch.equal(model.Xs[0], Xs1[0]))
self.assertTrue(torch.equal(model.Xs[1], Xs2_diff[0]))
self.assertEqual(model.dtype, Xs1[0].dtype)
self.assertEqual(model.device, Xs1[0].device)
self.assertIsInstance(model.model, ModelListGP)
# Check fitting
model_list = model.model.models
self.assertTrue(torch.equal(model_list[0].train_inputs[0], Xs1[0]))
self.assertTrue(torch.equal(model_list[1].train_inputs[0],
Xs2_diff[0]))
self.assertTrue(
torch.equal(model_list[0].train_targets, Ys1[0].view(-1)))
self.assertTrue(
torch.equal(model_list[1].train_targets, Ys2[0].view(-1)))
self.assertIsInstance(model_list[0].likelihood,
_GaussianLikelihoodBase)
self.assertIsInstance(model_list[1].likelihood,
_GaussianLikelihoodBase)
# Test batched multi-output FixedNoiseGP
with mock.patch(FIT_MODEL_MO_PATH) as _mock_fit_model:
model.fit(
Xs=Xs1 + Xs2,
Ys=Ys1 + Ys2,
Yvars=Yvars1 + Yvars2,
bounds=bounds,
task_features=task_features,
feature_names=feature_names,
fidelity_features=[],
)
_mock_fit_model.assert_called_once()
# Check attributes
self.assertTrue(torch.equal(model.Xs[0], Xs1[0]))
self.assertTrue(torch.equal(model.Xs[1], Xs2[0]))
self.assertEqual(model.dtype, Xs1[0].dtype)
self.assertEqual(model.device, Xs1[0].device)
self.assertIsInstance(model.model, FixedNoiseGP)
# Check fitting
# train inputs should be `o x n x 1`
self.assertTrue(
torch.equal(
model.model.train_inputs[0],
Xs1[0].unsqueeze(0).expand(torch.Size([2]) + Xs1[0].shape),
))
# train targets should be `o x n`
self.assertTrue(
torch.equal(model.model.train_targets,
torch.cat(Ys1 + Ys2, dim=-1).permute(1, 0)))
self.assertIsInstance(model.model.likelihood, _GaussianLikelihoodBase)
# Check infeasible cost can be computed on the model
device = torch.device("cuda") if cuda else torch.device("cpu")
objective_weights = torch.tensor([1.0, 0.0],
dtype=dtype,
device=device)
objective_transform = get_objective_weights_transform(
objective_weights)
infeasible_cost = torch.tensor(
get_infeasible_cost(X=Xs1[0],
model=model.model,
objective=objective_transform))
expected_infeasible_cost = -1 * torch.min(
objective_transform(
model.model.posterior(Xs1[0]).mean -
6 * model.model.posterior(Xs1[0]).variance.sqrt()).min(),
torch.tensor(0.0, dtype=dtype, device=device),
)
self.assertTrue(
torch.abs(infeasible_cost - expected_infeasible_cost) < 1e-5)
# Check prediction
X = torch.tensor([[6.0, 7.0, 8.0]], dtype=dtype, device=device)
f_mean, f_cov = model.predict(X)
self.assertTrue(f_mean.shape == torch.Size([1, 2]))
self.assertTrue(f_cov.shape == torch.Size([1, 2, 2]))
# Check generation
objective_weights = torch.tensor([1.0, 0.0],
dtype=dtype,
device=device)
outcome_constraints = (
torch.tensor([[0.0, 1.0]], dtype=dtype, device=device),
torch.tensor([[5.0]], dtype=dtype, device=device),
)
linear_constraints = (torch.tensor([[0.0, 1.0,
1.0]]), torch.tensor([[100.0]]))
fixed_features = None
pending_observations = [
torch.tensor([[1.0, 3.0, 4.0]], dtype=dtype, device=device),
torch.tensor([[2.0, 6.0, 8.0]], dtype=dtype, device=device),
]
n = 3
X_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], dtype=dtype, device=device)
model_gen_options = {}
# test sequential optimize
with mock.patch("ax.models.torch.botorch_defaults.sequential_optimize",
return_value=X_dummy) as mock_optimize_acqf:
Xgen, wgen = model.gen(
n=n,
bounds=bounds,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
linear_constraints=linear_constraints,
fixed_features=fixed_features,
pending_observations=pending_observations,
model_gen_options=model_gen_options,
rounding_func=dummy_func,
)
# note: gen() always returns CPU tensors
self.assertTrue(torch.equal(Xgen, X_dummy.cpu()))
self.assertTrue(torch.equal(wgen, torch.ones(n, dtype=dtype)))
# test joint optimize
with mock.patch("ax.models.torch.botorch_defaults.joint_optimize",
return_value=X_dummy) as mock_optimize_acqf:
Xgen, wgen = model.gen(
n=n,
bounds=bounds,
objective_weights=objective_weights,
outcome_constraints=None,
linear_constraints=None,
fixed_features=fixed_features,
pending_observations=pending_observations,
model_gen_options={
"optimizer_kwargs": {
"joint_optimization": True
}
},
)
# note: gen() always returns CPU tensors
self.assertTrue(torch.equal(Xgen, X_dummy.cpu()))
self.assertTrue(torch.equal(wgen, torch.ones(n, dtype=dtype)))
mock_optimize_acqf.assert_called_once()
# test get_rounding_func
dummy_rounding = get_rounding_func(rounding_func=dummy_func)
X_temp = torch.rand(1, 2, 3, 4)
self.assertTrue(torch.equal(X_temp, dummy_rounding(X_temp)))
# Check best point selection
xbest = model.best_point(bounds=bounds,
objective_weights=objective_weights)
xbest = model.best_point(
bounds=bounds,
objective_weights=objective_weights,
fixed_features={0: 100.0},
)
self.assertIsNone(xbest)
# Test cross-validation
mean, variance = model.cross_validate(
Xs_train=Xs1 + Xs2,
Ys_train=Ys1 + Ys2,
Yvars_train=Yvars1 + Yvars2,
X_test=torch.tensor([[1.2, 3.2, 4.2], [2.4, 5.2, 3.2]],
dtype=dtype,
device=device),
)
self.assertTrue(mean.shape == torch.Size([2, 2]))
self.assertTrue(variance.shape == torch.Size([2, 2, 2]))
# Test cross-validation with refit_on_cv
model.refit_on_cv = True
mean, variance = model.cross_validate(
Xs_train=Xs1 + Xs2,
Ys_train=Ys1 + Ys2,
Yvars_train=Yvars1 + Yvars2,
X_test=torch.tensor([[1.2, 3.2, 4.2], [2.4, 5.2, 3.2]],
dtype=dtype,
device=device),
)
self.assertTrue(mean.shape == torch.Size([2, 2]))
self.assertTrue(variance.shape == torch.Size([2, 2, 2]))
# Test update
model.refit_on_update = False
model.update(Xs=Xs2 + Xs2, Ys=Ys2 + Ys2, Yvars=Yvars2 + Yvars2)
# Test feature_importances
importances = model.feature_importances()
self.assertEqual(importances.shape, torch.Size([2, 1, 3]))
# When calling update directly, the data is completely overwritten.
self.assertTrue(torch.equal(model.Xs[0], Xs2[0]))
self.assertTrue(torch.equal(model.Xs[1], Xs2[0]))
self.assertTrue(torch.equal(model.Ys[0], Ys2[0]))
self.assertTrue(torch.equal(model.Yvars[0], Yvars2[0]))
model.refit_on_update = True
with mock.patch(FIT_MODEL_MO_PATH) as _mock_fit_model:
model.update(Xs=Xs2 + Xs2, Ys=Ys2 + Ys2, Yvars=Yvars2 + Yvars2)
# test unfit model CV, update, and feature_importances
unfit_model = BotorchModel()
with self.assertRaises(RuntimeError):
unfit_model.cross_validate(
Xs_train=Xs1 + Xs2,
Ys_train=Ys1 + Ys2,
Yvars_train=Yvars1 + Yvars2,
X_test=Xs1[0],
)
with self.assertRaises(RuntimeError):
unfit_model.update(Xs=Xs1 + Xs2,
Ys=Ys1 + Ys2,
Yvars=Yvars1 + Yvars2)
with self.assertRaises(RuntimeError):
unfit_model.feature_importances()
# Test loading state dict
tkwargs = {"device": device, "dtype": dtype}
true_state_dict = {
"mean_module.constant": [3.5004],
"covar_module.raw_outputscale":
2.2438,
"covar_module.base_kernel.raw_lengthscale":
[[-0.9274, -0.9274, -0.9274]],
"covar_module.base_kernel.lengthscale_prior.concentration":
3.0,
"covar_module.base_kernel.lengthscale_prior.rate":
6.0,
"covar_module.outputscale_prior.concentration":
2.0,
"covar_module.outputscale_prior.rate":
0.15,
}
true_state_dict = {
key: torch.tensor(val, **tkwargs)
for key, val in true_state_dict.items()
}
model = get_and_fit_model(
Xs=Xs1,
Ys=Ys1,
Yvars=Yvars1,
task_features=[],
fidelity_features=[],
state_dict=true_state_dict,
refit_model=False,
)
for k, v in chain(model.named_parameters(), model.named_buffers()):
self.assertTrue(torch.equal(true_state_dict[k], v))
# Test for some change in model parameters & buffer for refit_model=True
true_state_dict["mean_module.constant"] += 0.1
true_state_dict["covar_module.raw_outputscale"] += 0.1
true_state_dict["covar_module.base_kernel.raw_lengthscale"] += 0.1
true_state_dict = {
key: torch.tensor(val, **tkwargs)
for key, val in true_state_dict.items()
}
model = get_and_fit_model(
Xs=Xs1,
Ys=Ys1,
Yvars=Yvars1,
task_features=[],
fidelity_features=[],
state_dict=true_state_dict,
refit_model=True,
)
self.assertTrue(
any(not torch.equal(true_state_dict[k], v) for k, v in chain(
model.named_parameters(), model.named_buffers())))
def get_NEI(
model: Model,
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
X_observed: Optional[Tensor] = None,
X_pending: Optional[Tensor] = None,
**kwargs: Any,
) -> AcquisitionFunction:
r"""Instantiates a qNoisyExpectedImprovement acquisition function.
Args:
model: The underlying model which the acqusition function uses
to estimate acquisition values of candidates.
objective_weights: The objective is to maximize a weighted sum of
the columns of f(x). These are the weights.
outcome_constraints: A tuple of (A, b). For k outcome constraints
and m outputs at f(x), A is (k x m) and b is (k x 1) such that
A f(x) <= b. (Not used by single task models)
X_observed: A tensor containing points observed for all objective
outcomes and outcomes that appear in the outcome constraints (if
there are any).
X_pending: A tensor containing points whose evaluation is pending (i.e.
that have been submitted for evaluation) present for all objective
outcomes and outcomes that appear in the outcome constraints (if
there are any).
mc_samples: The number of MC samples to use (default: 512).
qmc: If True, use qMC instead of MC (default: True).
prune_baseline: If True, prune the baseline points for NEI (default: True).
chebyshev_scalarization: Use augmented Chebyshev scalarization.
Returns:
qNoisyExpectedImprovement: The instantiated acquisition function.
"""
if X_observed is None:
raise ValueError("There are no feasible observed points.")
# construct Objective module
if kwargs.get("chebyshev_scalarization", False):
if "Ys" not in kwargs:
raise ValueError("Chebyshev Scalarization requires Ys argument")
Y_tensor = torch.cat(kwargs.get("Ys"), dim=-1)
obj_tf = get_chebyshev_scalarization(weights=objective_weights,
Y=Y_tensor)
else:
obj_tf = get_objective_weights_transform(objective_weights)
if outcome_constraints is None:
objective = GenericMCObjective(objective=obj_tf)
else:
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
inf_cost = get_infeasible_cost(X=X_observed,
model=model,
objective=obj_tf)
objective = ConstrainedMCObjective(objective=obj_tf,
constraints=con_tfs or [],
infeasible_cost=inf_cost)
return get_acquisition_function(
acquisition_function_name="qNEI",
model=model,
objective=objective,
X_observed=X_observed,
X_pending=X_pending,
prune_baseline=kwargs.get("prune_baseline", True),
mc_samples=kwargs.get("mc_samples", 512),
qmc=kwargs.get("qmc", True),
# pyre-fixme[6]: Expected `Optional[int]` for 9th param but got
# `Union[float, int]`.
seed=torch.randint(1, 10000, (1, )).item(),
)
def test_NoMCSamples(self):
Y = torch.ones(2, 4, 2)
objective_transform = get_objective_weights_transform(
torch.tensor([1.0, 1.0]))
Y_transformed = objective_transform(Y)
self.assertTrue(torch.equal(torch.sum(Y, dim=-1), Y_transformed))
def test_IncompatibleNumberOfWeights(self):
Y = torch.ones(5, 2, 4, 1)
objective_transform = get_objective_weights_transform(
torch.tensor([1.0, 2.0]))
with self.assertRaises(RuntimeError):
objective_transform(Y)
def test_OneWeightBroadcasting(self):
Y = torch.ones(5, 2, 4, 1)
objective_transform = get_objective_weights_transform(
torch.tensor([0.5]))
Y_transformed = objective_transform(Y)
self.assertTrue(torch.equal(0.5 * Y.sum(dim=-1), Y_transformed))
def test_NoWeights(self):
Y = torch.ones(5, 2, 4, 1)
objective_transform = get_objective_weights_transform(None)
Y_transformed = objective_transform(Y)
self.assertTrue(torch.equal(Y.squeeze(-1), Y_transformed))
def _get_acquisition_func(
model: Model,
acquisition_function_name: str,
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
X_observed: Optional[Tensor] = None,
X_pending: Optional[Tensor] = None,
mc_objective: Type[GenericMCObjective] = GenericMCObjective,
constrained_mc_objective: Optional[
Type[ConstrainedMCObjective]
] = ConstrainedMCObjective,
mc_objective_kwargs: Optional[Dict] = None,
**kwargs: Any,
) -> AcquisitionFunction:
r"""Instantiates a acquisition function.
Args:
model: The underlying model which the acqusition function uses
to estimate acquisition values of candidates.
acquisition_function_name: Name of the acquisition function.
objective_weights: The objective is to maximize a weighted sum of
the columns of f(x). These are the weights.
outcome_constraints: A tuple of (A, b). For k outcome constraints
and m outputs at f(x), A is (k x m) and b is (k x 1) such that
A f(x) <= b. (Not used by single task models)
X_observed: A tensor containing points observed for all objective
outcomes and outcomes that appear in the outcome constraints (if
there are any).
X_pending: A tensor containing points whose evaluation is pending (i.e.
that have been submitted for evaluation) present for all objective
outcomes and outcomes that appear in the outcome constraints (if
there are any).
mc_objective: GenericMCObjective class, used for constructing a
MC-objective. If constructing a penalized MC-objective, pass in
PenalizedMCObjective together with mc_objective_kwargs .
constrained_mc_objective: ConstrainedMCObjective class, used when
applying constraints on the outcomes.
mc_objective_kwargs: kwargs for constructing MC-objective.
For GenericMCObjective, leave it as None. For PenalizedMCObjective,
it needs to be specified in the format of kwargs.
mc_samples: The number of MC samples to use (default: 512).
qmc: If True, use qMC instead of MC (default: True).
prune_baseline: If True, prune the baseline points for NEI (default: True).
chebyshev_scalarization: Use augmented Chebyshev scalarization.
Returns:
The instantiated acquisition function.
"""
if X_observed is None:
raise ValueError("There are no feasible observed points.")
# construct Objective module
if kwargs.get("chebyshev_scalarization", False):
with torch.no_grad():
Y = model.posterior(X_observed).mean
obj_tf = get_chebyshev_scalarization(weights=objective_weights, Y=Y)
else:
obj_tf = get_objective_weights_transform(objective_weights)
def objective(samples: Tensor, X: Optional[Tensor] = None) -> Tensor:
return obj_tf(samples)
if outcome_constraints is None:
mc_objective_kwargs = {} if mc_objective_kwargs is None else mc_objective_kwargs
objective = mc_objective(objective=objective, **mc_objective_kwargs)
else:
if constrained_mc_objective is None:
raise ValueError(
"constrained_mc_objective cannot be set to None "
"when applying outcome constraints."
)
if issubclass(mc_objective, PenalizedMCObjective):
raise RuntimeError(
"Outcome constraints are not supported for PenalizedMCObjective."
)
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
inf_cost = get_infeasible_cost(X=X_observed, model=model, objective=objective)
objective = constrained_mc_objective(
objective=objective, constraints=con_tfs or [], infeasible_cost=inf_cost
)
return get_acquisition_function(
acquisition_function_name=acquisition_function_name,
model=model,
objective=objective,
X_observed=X_observed,
X_pending=X_pending,
prune_baseline=kwargs.get("prune_baseline", True),
mc_samples=kwargs.get("mc_samples", 512),
qmc=kwargs.get("qmc", True),
# pyre-fixme[6]: Expected `Optional[int]` for 9th param but got
# `Union[float, int]`.
seed=torch.randint(1, 10000, (1,)).item(),
marginalize_dim=kwargs.get("marginalize_dim"),
)
def fit_albo_objective(self) -> AlboMCObjective:
r"""Inner loop of Augmented Lagrangian algorithm
Args:
model: A BoTorch model, fitted to observed data
objective_callable: A callable transformation from model outputs to objective
constraints_callable_list: A callable transformation from model outputs to constraints,
with negative values imply feasibility
sampler: An MCSampler instance for monte-carlo acquisition
Returns:
albo_objective: augmented objective with fitted Lagrangian multipliers
trace: optimization trace
"""
objective_callable = get_objective_weights_transform(
self.objective_weights)
constraints_callable_list = get_outcome_constraint_transforms(
self.outcome_constraints)
penalty_rate = self.init_penalty_rate
num_mults = self.outcome_constraints[0].shape[0]
if self.init_mults is not None:
assert num_mults == self.init_mults.shape[-1]
mults = self.init_mults
else:
mults = torch.Tensor(
[self._default_init_mult for _ in range(num_mults)])
x_trace = torch.zeros_like(self.bounds[0].unsqueeze(0))
mults_trace = mults.unsqueeze(0)
output_means = torch.zeros((1, self.model.num_outputs), dtype=float)
output_variances = torch.zeros((1, self.model.num_outputs),
dtype=float)
for i in range(self.num_iter):
self._execute_callbacks("on_iter_start", locals())
# 1. Optimize the augmented objective with fixed multipliers to find the next point for multipliers update
albo_objective = self.albo_objective_constructor(
objective=objective_callable,
constraints=constraints_callable_list,
penalty_rate=penalty_rate,
lagrange_mults=mults)
# Using predictive mean for inner loop optimization
acq_function = qSimpleRegret(model=self.model,
objective=albo_objective,
sampler=self.sampler)
x, val = optimize_acqf(acq_function=acq_function,
bounds=self.bounds,
q=1,
num_restarts=self.num_restarts,
raw_samples=self.raw_samples)
# 2. Compute update of lagrange multipliers at the optimal point of augmented objective
posterior = self.model.posterior(x.unsqueeze(0))
samples = self.sampler(posterior)
mults_next, mults_stds_next = albo_objective.get_mults_update(
samples)
# 3. Possibly apply heuristics here, i.e clamp mults before update and increase penalty rate
mults = mults_next # currently not using any heuristics
penalty_rate = self.init_penalty_rate # seems to work just fine with a constant penalty rate
# 4. Write trace of inner-loop optimization for debugging
x_trace = torch.cat([x_trace, x], dim=0)
mults_trace = torch.cat([mults_trace, mults.unsqueeze(0)], dim=0)
output_means = torch.cat(
[output_means,
posterior.mean.detach().squeeze(dim=0)], dim=0)
output_variances = torch.cat(
[output_variances,
posterior.variance.detach().squeeze(dim=0)],
dim=0)
# 5. Check stopping condition for inner loop (not implemented)
self._execute_callbacks("on_iter_end", locals())
continue
# Construct final objective
albo_objective = self.albo_objective_constructor(
objective=objective_callable,
constraints=constraints_callable_list,
penalty_rate=penalty_rate,
lagrange_mults=mults)
trace = {
'x': x_trace,
'mults': mults_trace,
'output': {
'mean': output_means,
'variance': output_variances
}
}
return albo_objective, trace
def test_FullyBayesianBotorchModel(self,
dtype=torch.float,
cuda=False):
Xs1, Ys1, Yvars1, bounds, tfs, fns, mns = get_torch_test_data(
dtype=dtype, cuda=cuda, constant_noise=True)
Xs2, Ys2, Yvars2, _, _, _, _ = get_torch_test_data(
dtype=dtype, cuda=cuda, constant_noise=True)
Yvars_inferred_noise = [
torch.full_like(Yvars1[0], float("nan")),
torch.full_like(Yvars2[0], float("nan")),
]
# make input different for each output
Xs2_diff = [Xs2[0] + 0.1]
Xs = Xs1 + Xs2_diff
Ys = Ys1 + Ys2
for inferred_noise, use_input_warping, use_saas in product(
(True, False), repeat=3):
Yvars = Yvars_inferred_noise if inferred_noise else Yvars1 + Yvars2
model = self.model_cls(
use_input_warping=use_input_warping,
thinning=1,
num_samples=4,
use_saas=use_saas,
disable_progbar=True,
max_tree_depth=1,
)
if use_input_warping:
self.assertTrue(model.use_input_warping)
# Test ModelListGP
# make training data different for each output
tkwargs = {"dtype": dtype, "device": Xs1[0].device}
dummy_samples_list = _get_dummy_mcmc_samples(num_samples=4,
num_outputs=2,
**tkwargs)
for dummy_samples in dummy_samples_list:
if use_input_warping:
dummy_samples["c0"] = (
torch.rand(4, 1, Xs1[0].shape[-1], **tkwargs) * 0.5
+ 0.1)
dummy_samples["c1"] = (
torch.rand(4, 1, Xs1[0].shape[-1], **tkwargs) * 0.5
+ 0.1)
if inferred_noise:
dummy_samples["noise"] = torch.rand(
4, 1,
**tkwargs).clamp_min(MIN_INFERRED_NOISE_LEVEL)
with mock.patch(
RUN_INFERENCE_PATH,
side_effect=dummy_samples_list,
) as _mock_fit_model:
model.fit(
Xs=Xs,
Ys=Ys,
Yvars=Yvars,
search_space_digest=SearchSpaceDigest(
feature_names=fns,
bounds=bounds,
task_features=tfs,
),
metric_names=mns,
)
self.assertEqual(_mock_fit_model.call_count, 2)
for i, call in enumerate(_mock_fit_model.call_args_list):
_, ckwargs = call
X = Xs[i]
Y = Ys[i]
Yvar = Yvars[i]
self.assertIs(ckwargs["pyro_model"], pyro_model)
self.assertTrue(torch.equal(ckwargs["X"], X))
self.assertTrue(torch.equal(ckwargs["Y"], Y))
if inferred_noise:
self.assertTrue(torch.isnan(ckwargs["Yvar"]).all())
else:
self.assertTrue(torch.equal(ckwargs["Yvar"], Yvar))
self.assertEqual(ckwargs["num_samples"], 4)
self.assertEqual(ckwargs["warmup_steps"], 1024)
self.assertEqual(ckwargs["max_tree_depth"], 1)
self.assertTrue(ckwargs["disable_progbar"])
self.assertEqual(ckwargs["use_input_warping"],
use_input_warping)
self.assertEqual(ckwargs["use_saas"], use_saas)
# Check attributes
self.assertTrue(torch.equal(model.Xs[i], Xs[i]))
self.assertEqual(model.dtype, Xs[i].dtype)
self.assertEqual(model.device, Xs[i].device)
self.assertIsInstance(model.model, ModelListGP)
# Check fitting
# Note each model in the model list is a batched model, where
# the batch dim corresponds to the MCMC samples
model_list = model.model.models
# Put model in `eval` mode to transform the train inputs.
m = model_list[i].eval()
# check mcmc samples
dummy_samples = dummy_samples_list[i]
expected_train_inputs = Xs[i].expand(4, *Xs[i].shape)
if use_input_warping:
# train inputs should be warped inputs
expected_train_inputs = m.input_transform(
expected_train_inputs)
self.assertTrue(
torch.equal(
m.train_inputs[0],
expected_train_inputs,
))
self.assertTrue(
torch.equal(
m.train_targets,
Ys[i].view(1, -1).expand(4, Ys[i].numel()),
))
expected_noise = (dummy_samples["noise"].view(
m.likelihood.noise.shape) if inferred_noise else
Yvars[i].view(1, -1).expand(
4, Yvars[i].numel()))
self.assertTrue(
torch.allclose(
m.likelihood.noise.detach(),
expected_noise,
))
self.assertIsInstance(m.likelihood,
_GaussianLikelihoodBase)
self.assertTrue(
torch.allclose(
m.covar_module.base_kernel.lengthscale.detach(
),
dummy_samples["lengthscale"].view(
m.covar_module.base_kernel.lengthscale.
shape),
))
self.assertTrue(
torch.allclose(
m.covar_module.outputscale.detach(),
dummy_samples["outputscale"].view(
m.covar_module.outputscale.shape),
))
self.assertTrue(
torch.allclose(
m.mean_module.constant.detach(),
dummy_samples["mean"].view(
m.mean_module.constant.shape),
))
if use_input_warping:
self.assertTrue(hasattr(m, "input_transform"))
self.assertIsInstance(m.input_transform, Warp)
self.assertTrue(
torch.equal(
m.input_transform.concentration0,
dummy_samples_list[i]["c0"],
))
self.assertTrue(
torch.equal(
m.input_transform.concentration1,
dummy_samples_list[i]["c1"],
))
else:
self.assertFalse(hasattr(m, "input_transform"))
# test that multi-task is not implemented
(
Xs_mt,
Ys_mt,
Yvars_mt,
bounds_mt,
tfs_mt,
fns_mt,
mns_mt,
) = get_torch_test_data(dtype=dtype,
cuda=cuda,
constant_noise=True,
task_features=[2])
with mock.patch(
RUN_INFERENCE_PATH,
side_effect=dummy_samples_list,
) as _mock_fit_model, self.assertRaises(NotImplementedError):
model.fit(
Xs=Xs_mt,
Ys=Ys_mt,
Yvars=Yvars_mt,
search_space_digest=SearchSpaceDigest(
feature_names=fns_mt,
bounds=bounds_mt,
task_features=tfs_mt,
),
metric_names=mns_mt,
)
with mock.patch(
RUN_INFERENCE_PATH,
side_effect=dummy_samples_list,
) as _mock_fit_model, self.assertRaises(NotImplementedError):
model.fit(
Xs=Xs1 + Xs2,
Ys=Ys1 + Ys2,
Yvars=Yvars1 + Yvars2,
search_space_digest=SearchSpaceDigest(
feature_names=fns,
bounds=bounds,
fidelity_features=[0],
),
metric_names=mns,
)
# fit model with same inputs (otherwise X_observed will be None)
model = self.model_cls(
use_input_warping=use_input_warping,
thinning=1,
num_samples=4,
use_saas=use_saas,
disable_progbar=True,
max_tree_depth=1,
)
Yvars = Yvars1 + Yvars2
dummy_samples_list = _get_dummy_mcmc_samples(num_samples=4,
num_outputs=2,
**tkwargs)
with mock.patch(
RUN_INFERENCE_PATH,
side_effect=dummy_samples_list,
) as _mock_fit_model:
model.fit(
Xs=Xs1 + Xs2,
Ys=Ys1 + Ys2,
Yvars=Yvars,
search_space_digest=SearchSpaceDigest(
feature_names=fns,
bounds=bounds,
task_features=tfs,
),
metric_names=mns,
)
# Check infeasible cost can be computed on the model
device = torch.device("cuda") if cuda else torch.device("cpu")
objective_weights = torch.tensor([1.0, 0.0],
dtype=dtype,
device=device)
objective_transform = get_objective_weights_transform(
objective_weights)
infeasible_cost = torch.tensor(
get_infeasible_cost(X=Xs1[0],
model=model.model,
objective=objective_transform))
expected_infeasible_cost = -1 * torch.min(
objective_transform(
model.model.posterior(Xs1[0]).mean - 6 *
model.model.posterior(Xs1[0]).variance.sqrt()).min(),
torch.tensor(0.0, dtype=dtype, device=device),
)
self.assertTrue(
torch.abs(infeasible_cost -
expected_infeasible_cost) < 1e-5)
# Check prediction
X = torch.tensor([[6.0, 7.0, 8.0]], **tkwargs)
f_mean, f_cov = model.predict(X)
self.assertTrue(f_mean.shape == torch.Size([1, 2]))
self.assertTrue(f_cov.shape == torch.Size([1, 2, 2]))
# Check generation
objective_weights = torch.tensor(
[1.0, 0.0] if self.model_cls is FullyBayesianBotorchModel
else [1.0, 1.0],
**tkwargs,
)
outcome_constraints = (
torch.tensor([[0.0, 1.0]], **tkwargs),
torch.tensor([[5.0]], **tkwargs),
)
gen_kwargs = ({
"objective_thresholds": torch.zeros(2, **tkwargs)
} if self.model_cls is FullyBayesianMOOBotorchModel else {})
linear_constraints = (
torch.tensor([[0.0, 1.0, 1.0]]),
torch.tensor([[100.0]]),
)
fixed_features = None
pending_observations = [
torch.tensor([[1.0, 3.0, 4.0]], **tkwargs),
torch.tensor([[2.0, 6.0, 8.0]], **tkwargs),
]
n = 3
X_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], **tkwargs)
acqfv_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], **tkwargs)
model_gen_options = {
Keys.OPTIMIZER_KWARGS: {
"maxiter": 1
},
Keys.ACQF_KWARGS: {
"mc_samples": 3
},
}
# test sequential optimize with constraints
with mock.patch(
"ax.models.torch.botorch_defaults.optimize_acqf",
return_value=(X_dummy, acqfv_dummy),
) as _:
Xgen, wgen, gen_metadata, cand_metadata = model.gen(
n=n,
bounds=bounds,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
linear_constraints=linear_constraints,
fixed_features=fixed_features,
pending_observations=pending_observations,
model_gen_options=model_gen_options,
rounding_func=dummy_func,
**gen_kwargs,
)
# note: gen() always returns CPU tensors
self.assertTrue(torch.equal(Xgen, X_dummy.cpu()))
self.assertTrue(
torch.equal(wgen, torch.ones(n, dtype=dtype)))
# actually test optimization for 1 step without constraints
with mock.patch(
"ax.models.torch.botorch_defaults.optimize_acqf",
wraps=optimize_acqf,
return_value=(X_dummy, acqfv_dummy),
) as _:
Xgen, wgen, gen_metadata, cand_metadata = model.gen(
n=n,
bounds=bounds,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
fixed_features=fixed_features,
pending_observations=pending_observations,
model_gen_options=model_gen_options,
**gen_kwargs,
)
# note: gen() always returns CPU tensors
self.assertTrue(torch.equal(Xgen, X_dummy.cpu()))
self.assertTrue(
torch.equal(wgen, torch.ones(n, dtype=dtype)))
# Check best point selection
xbest = model.best_point(bounds=bounds,
objective_weights=objective_weights)
xbest = model.best_point(
bounds=bounds,
objective_weights=objective_weights,
fixed_features={0: 100.0},
)
self.assertIsNone(xbest)
# Test cross-validation
mean, variance = model.cross_validate(
Xs_train=Xs1 + Xs2,
Ys_train=Ys,
Yvars_train=Yvars,
X_test=torch.tensor([[1.2, 3.2, 4.2], [2.4, 5.2, 3.2]],
dtype=dtype,
device=device),
)
self.assertTrue(mean.shape == torch.Size([2, 2]))
self.assertTrue(variance.shape == torch.Size([2, 2, 2]))
# Test cross-validation with refit_on_cv
model.refit_on_cv = True
with mock.patch(
RUN_INFERENCE_PATH,
side_effect=dummy_samples_list,
) as _mock_fit_model:
mean, variance = model.cross_validate(
Xs_train=Xs1 + Xs2,
Ys_train=Ys,
Yvars_train=Yvars,
X_test=torch.tensor(
[[1.2, 3.2, 4.2], [2.4, 5.2, 3.2]],
dtype=dtype,
device=device,
),
)
self.assertTrue(mean.shape == torch.Size([2, 2]))
self.assertTrue(variance.shape == torch.Size([2, 2, 2]))
# Test update
model.refit_on_update = False
model.update(Xs=Xs2 + Xs2, Ys=Ys2 + Ys2, Yvars=Yvars2 + Yvars2)
# Test feature_importances
importances = model.feature_importances()
self.assertEqual(importances.shape, torch.Size([2, 1, 3]))
# When calling update directly, the data is completely overwritten.
self.assertTrue(torch.equal(model.Xs[0], Xs2[0]))
self.assertTrue(torch.equal(model.Xs[1], Xs2[0]))
self.assertTrue(torch.equal(model.Ys[0], Ys2[0]))
self.assertTrue(torch.equal(model.Yvars[0], Yvars2[0]))
model.refit_on_update = True
with mock.patch(
RUN_INFERENCE_PATH,
side_effect=dummy_samples_list) as _mock_fit_model:
model.update(Xs=Xs2 + Xs2,
Ys=Ys2 + Ys2,
Yvars=Yvars2 + Yvars2)
# test unfit model CV, update, and feature_importances
unfit_model = self.model_cls()
with self.assertRaises(RuntimeError):
unfit_model.cross_validate(
Xs_train=Xs1 + Xs2,
Ys_train=Ys1 + Ys2,
Yvars_train=Yvars1 + Yvars2,
X_test=Xs1[0],
)
with self.assertRaises(RuntimeError):
unfit_model.update(Xs=Xs1 + Xs2,
Ys=Ys1 + Ys2,
Yvars=Yvars1 + Yvars2)
with self.assertRaises(RuntimeError):
unfit_model.feature_importances()
