def test_task_feature(self, get_model_mock):
x = [torch.zeros(2, 2)]
y = [torch.zeros(2, 1)]
yvars = [torch.ones(2, 1)]
get_and_fit_model(
Xs=x,
Ys=y,
Yvars=yvars,
task_features=[1],
fidelity_features=[],
metric_names=["L2NormMetric"],
state_dict=None,
refit_model=False,
)
# Check that task feature was correctly passed to _get_model
self.assertEqual(get_model_mock.mock_calls[0][2]["task_feature"], 1)
# check error on multiple task features
with self.assertRaises(NotImplementedError):
get_and_fit_model(
Xs=x,
Ys=y,
Yvars=yvars,
task_features=[0, 1],
fidelity_features=[],
metric_names=["L2NormMetric"],
state_dict=None,
refit_model=False,
)
# check error on multiple fidelity features
with self.assertRaises(NotImplementedError):
get_and_fit_model(
Xs=x,
Ys=y,
Yvars=yvars,
task_features=[],
fidelity_features=[-1, -2],
metric_names=["L2NormMetric"],
state_dict=None,
refit_model=False,
)
# check error on botch task and fidelity feature
with self.assertRaises(NotImplementedError):
get_and_fit_model(
Xs=x,
Ys=y,
Yvars=yvars,
task_features=[1],
fidelity_features=[-1],
metric_names=["L2NormMetric"],
state_dict=None,
refit_model=False,
)
def test_task_feature(self, gp_mock, get_model_mock):
x = [torch.zeros(2, 2)]
y = [torch.zeros(2, 1)]
yvars = [torch.ones(2, 1)]
get_and_fit_model(Xs=x,
Ys=y,
Yvars=yvars,
task_features=[1],
state_dict=[])
# Check that task feature was correctly passed to _get_model
self.assertEqual(get_model_mock.mock_calls[0][2]["task_feature"], 1)
with self.assertRaises(ValueError):
get_and_fit_model(Xs=x,
Ys=y,
Yvars=yvars,
task_features=[0, 1],
state_dict=[])
def test_task_feature(self, gp_mock, get_model_mock):
x = [torch.zeros(2, 2)]
y = [torch.zeros(2, 1)]
yvars = [torch.ones(2, 1)]
get_and_fit_model(
Xs=x,
Ys=y,
Yvars=yvars,
task_features=[1],
fidelity_features=[],
state_dict=[],
refit_model=False,
)
# Check that task feature was correctly passed to _get_model
self.assertEqual(get_model_mock.mock_calls[0][2]["task_feature"], 1)
with self.assertRaises(ValueError):
get_and_fit_model(
Xs=x,
Ys=y,
Yvars=yvars,
task_features=[0, 1],
fidelity_features=[],
state_dict=[],
refit_model=False,
)
with self.assertRaises(NotImplementedError):
get_and_fit_model(
Xs=x,
Ys=y,
Yvars=yvars,
task_features=[0, 1],
fidelity_features=[],
state_dict=[],
fidelity_model_id=0,
refit_model=False,
)
with self.assertRaises(UnsupportedError):
get_and_fit_model(
Xs=x,
Ys=y,
Yvars=yvars,
task_features=[],
fidelity_features=[-1, -2],
state_dict=[],
fidelity_model_id=0,
refit_model=False,
)
def get_and_fit_model(
self,
Xs: List[Tensor],
Ys: List[Tensor],
Yvars: List[Tensor],
state_dicts: Optional[List[MutableMapping[str, Tensor]]] = None,
) -> GPyTorchModel:
return get_and_fit_model(
Xs=Xs,
Ys=Ys,
Yvars=Yvars,
task_features=[],
fidelity_features=[],
metric_names=[],
state_dict=None,
)
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())))
