def prepare_acquisition_function(args, model_obj, train_x, train_y, bounds, step):
if args.num_steps > 500:
sampler = IIDNormalSampler(num_samples=256)
else:
sampler = SobolQMCNormalSampler(num_samples=256)
if args.acqf == "ei":
acqf = qExpectedImprovement(
model=model_obj, best_f=train_y.max(), sampler=sampler,
)
elif args.acqf == "ucb":
acqf = qUpperConfidenceBound(model=model_obj, beta=0.9 ** step)
elif args.acqf == "nei":
acqf = qNoisyExpectedImprovement(
model=model_obj, X_baseline=train_x, sampler=sampler
)
elif args.acqf == "kg":
acqf = qKnowledgeGradient(
model=model_obj,
sampler=sampler,
num_fantasies=None,
current_value=train_y.max(),
)
elif args.acqf == "mves":
candidate_set = torch.rand(10000, bounds.size(0), device=bounds.device)
candidate_set = bounds[..., 0] + (bounds[..., 1] - bounds[..., 0]) * candidate_set
acqf = qMaxValueEntropy(
model=model_obj, candidate_set=candidate_set, train_inputs=train_x,
)
return acqf
def test_batch_decoupled_sampler(self):
train_n = 6
dim = 2
transform = None
num_fantasies = 3
num_fant_X = 4
model = SingleTaskGP(
torch.rand(train_n, dim),
torch.rand(train_n, 1),
outcome_transform=transform,
)
fantasy_X = torch.rand(num_fant_X, 1, dim)
fantasy_model = model.fantasize(fantasy_X,
IIDNormalSampler(num_fantasies))
sample_shape = [5]
num_basis = 16
sampler = decoupled_sampler(model=fantasy_model,
sample_shape=sample_shape,
num_basis=num_basis)
num_test = 8
test_X = torch.rand(num_test, dim)
ds_samples = sampler(test_X)
self.assertEqual(
list(ds_samples.shape),
sample_shape + [num_fantasies, num_fant_X, num_test, 1],
)
def __init__(
self,
pref_model: Model,
sampler: Optional[MCSampler] = None,
):
r"""Learned preference objective constructed from a preference model.
Args:
pref_model: A BoTorch model, which models the latent preference/utility
function. Given an input tensor of size
`sample_size x batch_shape x N x d`, its `posterior` method should
return a `Posterior` object with single outcome representing the
utility values of the input.
sampler: Sampler for the preference model to account for uncertainty in
preferece when calculating the objective; it's not the one used
in MC acquisition functions. If None,
it uses `IIDNormalSampler(num_samples=1)`.
"""
super().__init__()
self.pref_model = pref_model
if isinstance(pref_model, DeterministicModel):
assert sampler is None
self.sampler = None
else:
if sampler is None:
self.sampler = IIDNormalSampler(num_samples=1)
else:
self.sampler = sampler
def test_likelihood_and_fantasize(self):
self.assertIsInstance(self.model.likelihood, GaussianLikelihood)
self.assertTrue(self.model._is_custom_likelihood, True)
test_X = torch.randn(5, 1, device=self.device)
with self.assertRaises(NotImplementedError):
self.model.fantasize(test_X, sampler=IIDNormalSampler(num_samples=32))
def test_fantasize(self):
manual_seed(0)
test_x = rand(2, 5, 1, device=self.device)
sampler = IIDNormalSampler(num_samples=32).to(self.device)
_ = self.model.posterior(test_x)
fantasy_model = self.model.fantasize(test_x, sampler=sampler)
self.assertIsInstance(fantasy_model, HigherOrderGP)
self.assertEqual(fantasy_model.train_inputs[0].shape[:2], Size((32, 2)))
def __call__(self,
model: Model,
objective_weights: Tensor,
outcome_constraints: Tuple[Tensor, Tensor],
X_observed: Optional[Tensor] = None,
X_pending: Optional[Tensor] = None,
mc_samples: int = 512,
qmc: bool = True,
seed: Optional[Tensor] = None,
**kwargs: Any) -> AcquisitionFunction:
r"""Creates an acquisition function.
The callable interface is compatible with `ax.modelbridge.factory.get_botorch` factory function.
See for example `ax.models.torch.botorch_defaults.get_NEI`.
Args:
model: A fitted GPyTorch model
objective_weights: Transform parameters from model output to raw objective
outcome_constraints: Transform parameters from model output to raw constaints
X_observed: Tensor of evaluated points
X_pending: Tensor of points whose evaluation is pending
mc_samples: The number of MC samples to use in the inner-loop optimization
qmc: If True, use qMC instead of MC
seed: Optional seed for MCSampler
Returns:
An instance of the acquisition function
"""
self.model = model
if qmc:
self.sampler = SobolQMCNormalSampler(num_samples=mc_samples,
seed=seed)
else:
self.sampler = IIDNormalSampler(num_samples=mc_samples, seed=seed)
# Store objective and constraints
self.objective_weights = objective_weights
self.outcome_constraints = outcome_constraints
# Run inner loop of Augmented Lagrangian algorithm for fitting of Lagrange Multipliers
# and store the fitted albo objective and the trace of inner loop optimization
# (for debugging and visualization).
self.albo_objective, self.trace_inner = self.fit_albo_objective()
# Create an acquisition function over the fitted AL objective
return get_acquisition_function(
acquisition_function_name=self.acquisition_function_name,
model=model,
objective=self.albo_objective,
X_observed=X_observed,
X_pending=X_pending,
mc_samples=mc_samples,
seed=seed,
**kwargs)
def test_InverseCostWeightedUtility(self):
for batch_shape in ([], [2]):
for dtype in (torch.float, torch.double):
# the event shape is `batch_shape x q x t`
mean = 1 + torch.rand(
*batch_shape, 2, 1, device=self.device, dtype=dtype)
mm = MockModel(MockPosterior(mean=mean))
X = torch.randn(*batch_shape,
3,
2,
device=self.device,
dtype=dtype)
deltas = torch.rand(4,
*batch_shape,
device=self.device,
dtype=dtype)
# test that sampler is required if use_mean=False
icwu = InverseCostWeightedUtility(mm, use_mean=False)
with self.assertRaises(RuntimeError):
icwu(X, deltas)
# check warning for negative cost
mm = MockModel(MockPosterior(mean=mean.clamp_max(-1e-6)))
icwu = InverseCostWeightedUtility(mm)
with warnings.catch_warnings(
record=True) as ws, settings.debug(True):
icwu(X, deltas)
self.assertTrue(
any(
issubclass(w.category, CostAwareWarning)
for w in ws))
# basic test
mm = MockModel(MockPosterior(mean=mean))
icwu = InverseCostWeightedUtility(mm)
ratios = icwu(X, deltas)
self.assertTrue(
torch.equal(ratios, deltas / mean.squeeze(-1).sum(dim=-1)))
# sampling test
samples = 1 + torch.rand( # event shape is q x m
*batch_shape,
3,
1,
device=self.device,
dtype=dtype)
mm = MockModel(MockPosterior(samples=samples))
icwu = InverseCostWeightedUtility(mm, use_mean=False)
ratios = icwu(X, deltas, sampler=IIDNormalSampler(4))
self.assertTrue(
torch.equal(ratios,
deltas / samples.squeeze(-1).sum(dim=-1)))
def test_fantasize(self):
for dtype in [torch.float, torch.double]:
torch.random.manual_seed(0)
test_x = torch.rand(2, 5, 1, device=self.device, dtype=dtype)
sampler = IIDNormalSampler(num_samples=32)
self.model.to(dtype)
if dtype == torch.double:
# need to clear float caches
self.model.train()
self.model.eval()
_ = self.model.posterior(test_x)
fantasy_model = self.model.fantasize(test_x, sampler=sampler)
self.assertIsInstance(fantasy_model, HigherOrderGP)
self.assertEqual(fantasy_model.train_inputs[0].shape[:2],
torch.Size((32, 2)))
def test_posterior(self):
# test the posterior works
sample_shaping = [5, 3, 5]
for post_collection in self.post_list:
model, test_x, posterior = post_collection
self.assertIsInstance(posterior, GPyTorchPosterior)
correct_shape = [2] if test_x.shape[0] == 2 else []
[correct_shape.append(s) for s in sample_shaping]
# test providing no base samples
samples_0 = posterior.rsample()
self.assertEqual(samples_0.shape, Size((1, *correct_shape)))
# test that providing all base samples produces non-random results
if test_x.shape[0] == 2:
base_samples = randn(8,
2, (5 + 10 + 10) * 3 * 5,
device=self.device)
else:
base_samples = randn(8, (5 + 10 + 10) * 3 * 5,
device=self.device)
samples_1 = posterior.rsample(base_samples=base_samples,
sample_shape=Size((8, )))
samples_2 = posterior.rsample(base_samples=base_samples,
sample_shape=Size((8, )))
self.assertTrue(allclose(samples_1, samples_2))
# test that botorch.sampler picks up the correct shapes
sampler = IIDNormalSampler(num_samples=5)
samples_det_shape = sampler(posterior).shape
self.assertEqual(samples_det_shape, Size([5, *correct_shape]))
# test that providing only some base samples is okay
base_samples = randn(8, prod(correct_shape), device=self.device)
samples_3 = posterior.rsample(base_samples=base_samples,
sample_shape=Size((8, )))
self.assertEqual(samples_3.shape, Size([8, *correct_shape]))
# test that providing the wrong number base samples is okay
base_samples = randn(8, 50 * 2 * 3 * 5, device=self.device)
samples_4 = posterior.rsample(base_samples=base_samples,
sample_shape=Size((8, )))
self.assertEqual(samples_4.shape, Size([8, *correct_shape]))
# test that providing the wrong shapes of base samples fails
base_samples = randn(8, 5 * 2 * 3 * 5, device=self.device)
with self.assertRaises(RuntimeError):
samples_4 = posterior.rsample(base_samples=base_samples,
sample_shape=Size((4, )))
# finally we check the quality of the variances and the samples
# test that the posterior variances are the same as the evaluation variance
posterior_variance = posterior.variance
model.eval()
eval_mode_variance = model(test_x).variance.reshape_as(
posterior_variance)
self.assertLess(
(posterior_variance - eval_mode_variance).norm() /
eval_mode_variance.norm(),
4e-2,
)
# and finally test that sampling with no base samples is okay
samples_3 = posterior.rsample(sample_shape=Size((5000, )))
sampled_variance = (samples_3.std(dim=0)**2).view(-1)
posterior_variance = posterior_variance.view(-1)
self.assertLess(
(posterior_variance - sampled_variance).norm() /
posterior_variance.norm(),
5e-2,
)
def test_HigherOrderGPPosterior(self):
sample_shaping = torch.Size([5, 3, 5])
for post_collection in self.post_list:
model, test_x, posterior = post_collection
self.assertIsInstance(posterior, HigherOrderGPPosterior)
batch_shape = test_x.shape[:-2]
# expected_event_shape = batch_shape + sample_shaping[-2:]
expected_event_shape = batch_shape + sample_shaping
self.assertEqual(posterior.event_shape, expected_event_shape)
# test providing no base samples
samples_0 = posterior.rsample()
self.assertEqual(samples_0.shape,
torch.Size((1, *expected_event_shape)))
# test that providing all base samples produces non-torch.random results
if len(batch_shape) > 0:
base_sample_shape = (8, 2, (5 + 10 + 10) * 3 * 5)
else:
base_sample_shape = (8, (5 + 10 + 10) * 3 * 5)
base_samples = torch.randn(*base_sample_shape, device=self.device)
samples_1 = posterior.rsample(base_samples=base_samples,
sample_shape=torch.Size((8, )))
samples_2 = posterior.rsample(base_samples=base_samples,
sample_shape=torch.Size((8, )))
self.assertTrue(torch.allclose(samples_1, samples_2))
# test that botorch.sampler picks up the correct shapes
sampler = IIDNormalSampler(num_samples=5)
samples_det_shape = sampler(posterior).shape
self.assertEqual(samples_det_shape,
torch.Size([5, *expected_event_shape]))
# test that providing only some base samples is okay
base_samples = torch.randn(8,
np.prod(expected_event_shape),
device=self.device)
samples_3 = posterior.rsample(base_samples=base_samples,
sample_shape=torch.Size((8, )))
self.assertEqual(samples_3.shape,
torch.Size([8, *expected_event_shape]))
# test that providing the wrong number base samples is okay
base_samples = torch.randn(8, 50 * 2 * 3 * 5, device=self.device)
samples_4 = posterior.rsample(base_samples=base_samples,
sample_shape=torch.Size((8, )))
self.assertEqual(samples_4.shape,
torch.Size([8, *expected_event_shape]))
# test that providing the wrong shapes of base samples fails
base_samples = torch.randn(8, 5 * 2 * 3 * 5, device=self.device)
with self.assertRaises(RuntimeError):
samples_4 = posterior.rsample(base_samples=base_samples,
sample_shape=torch.Size((4, )))
# finally we check the quality of the variances and the samples
# test that the posterior variances are the same as the evaluation variance
posterior_variance = posterior.variance
model.eval()
eval_mode_variance = model(test_x).variance.reshape_as(
posterior_variance)
self.assertLess(
(posterior_variance - eval_mode_variance).norm() /
eval_mode_variance.norm(),
4e-2,
)
# and finally test that sampling with no base samples is okay
samples_3 = posterior.rsample(sample_shape=torch.Size((5000, )))
sampled_variance = samples_3.var(dim=0).view(-1)
posterior_variance = posterior_variance.view(-1)
self.assertLess(
(posterior_variance - sampled_variance).norm() /
posterior_variance.norm(),
5e-2,
)
def test_gen_value_function_initial_conditions(self):
num_fantasies = 2
num_solutions = 3
num_restarts = 4
raw_samples = 5
n_train = 6
dim = 2
dtype = torch.float
# run a thorough test with dtype float
train_X = torch.rand(n_train, dim, device=self.device, dtype=dtype)
train_Y = torch.rand(n_train, 1, device=self.device, dtype=dtype)
model = SingleTaskGP(train_X, train_Y)
fant_X = torch.rand(num_solutions,
1,
dim,
device=self.device,
dtype=dtype)
fantasy_model = model.fantasize(fant_X,
IIDNormalSampler(num_fantasies))
bounds = torch.tensor([[0, 0], [1, 1]],
device=self.device,
dtype=dtype)
value_function = PosteriorMean(fantasy_model)
# test option error
with self.assertRaises(ValueError):
gen_value_function_initial_conditions(
acq_function=value_function,
bounds=bounds,
num_restarts=num_restarts,
raw_samples=raw_samples,
current_model=model,
options={"frac_random": 2.0},
)
# test output shape
ics = gen_value_function_initial_conditions(
acq_function=value_function,
bounds=bounds,
num_restarts=num_restarts,
raw_samples=raw_samples,
current_model=model,
)
self.assertEqual(
ics.shape,
torch.Size([num_restarts, num_fantasies, num_solutions, 1, dim]))
# test bounds
self.assertTrue(torch.all(ics >= bounds[0]))
self.assertTrue(torch.all(ics <= bounds[1]))
# test dtype
self.assertEqual(dtype, ics.dtype)
# minimal test cases for when all raw samples are random, with dtype double
dtype = torch.double
n_train = 2
dim = 1
num_solutions = 1
train_X = torch.rand(n_train, dim, device=self.device, dtype=dtype)
train_Y = torch.rand(n_train, 1, device=self.device, dtype=dtype)
model = SingleTaskGP(train_X, train_Y)
fant_X = torch.rand(1, 1, dim, device=self.device, dtype=dtype)
fantasy_model = model.fantasize(fant_X,
IIDNormalSampler(num_fantasies))
bounds = torch.tensor([[0], [1]], device=self.device, dtype=dtype)
value_function = PosteriorMean(fantasy_model)
ics = gen_value_function_initial_conditions(
acq_function=value_function,
bounds=bounds,
num_restarts=1,
raw_samples=1,
current_model=model,
options={"frac_random": 0.99},
)
self.assertEqual(ics.shape,
torch.Size([1, num_fantasies, num_solutions, 1, dim]))
# test bounds
self.assertTrue(torch.all(ics >= bounds[0]))
self.assertTrue(torch.all(ics <= bounds[1]))
# test dtype
self.assertEqual(dtype, ics.dtype)
def test_MultitaskGPPosterior(self):
sample_shaping = torch.Size([5, 3])
for post_collection in self.post_list:
model, test_x, posterior = post_collection
self.assertIsInstance(posterior, MultitaskGPPosterior)
batch_shape = test_x.shape[:-2]
expected_event_shape = batch_shape + sample_shaping
self.assertEqual(posterior.event_shape, expected_event_shape)
# test providing no base samples
samples_0 = posterior.rsample()
self.assertEqual(samples_0.shape,
torch.Size((1, *expected_event_shape)))
# test that providing all base samples produces non-torch.random results
scale = 2 if posterior.observation_noise else 1
base_sample_shaping = torch.Size([
2 * model.train_targets.numel() +
scale * sample_shaping[0] * sample_shaping[1]
])
expected_base_sample_shape = batch_shape + base_sample_shaping
self.assertEqual(
posterior.base_sample_shape,
expected_base_sample_shape,
)
base_samples = torch.randn(8,
*expected_base_sample_shape,
device=self.device)
samples_1 = posterior.rsample(base_samples=base_samples,
sample_shape=torch.Size((8, )))
samples_2 = posterior.rsample(base_samples=base_samples,
sample_shape=torch.Size((8, )))
self.assertTrue(torch.allclose(samples_1, samples_2))
# test that botorch.sampler picks up the correct shapes
sampler = IIDNormalSampler(num_samples=5)
samples_det_shape = sampler(posterior).shape
self.assertEqual(samples_det_shape,
torch.Size([5, *expected_event_shape]))
# test that providing only some base samples is okay
base_samples = torch.randn(8,
np.prod(expected_event_shape),
device=self.device)
samples_3 = posterior.rsample(base_samples=base_samples,
sample_shape=torch.Size((8, )))
self.assertEqual(samples_3.shape,
torch.Size([8, *expected_event_shape]))
# test that providing the wrong number base samples is okay
base_samples = torch.randn(8, 50 * 2 * 3, device=self.device)
samples_4 = posterior.rsample(base_samples=base_samples,
sample_shape=torch.Size((8, )))
self.assertEqual(samples_4.shape,
torch.Size([8, *expected_event_shape]))
# test that providing the wrong shapes of base samples fails
base_samples = torch.randn(8, 5 * 2 * 3, device=self.device)
with self.assertRaises(RuntimeError):
samples_4 = posterior.rsample(base_samples=base_samples,
sample_shape=torch.Size((4, )))
# finally we check the quality of the variances and the samples
# test that the posterior variances are the same as the evaluation variance
posterior_variance = posterior.variance
posterior_mean = posterior.mean.view(-1)
model.eval()
if not posterior.observation_noise:
eval_mode_variance = model(test_x).variance
else:
eval_mode_variance = model.likelihood(model(test_x)).variance
eval_mode_variance = eval_mode_variance.reshape_as(
posterior_variance)
self.assertLess(
(posterior_variance - eval_mode_variance).norm() /
eval_mode_variance.norm(),
4e-2,
)
# and finally test that sampling with no base samples is okay
samples_3 = posterior.rsample(sample_shape=torch.Size((10000, )))
sampled_variance = samples_3.var(dim=0).view(-1)
sampled_mean = samples_3.mean(dim=0).view(-1)
posterior_variance = posterior_variance.view(-1)
# slightly higher tolerance here because of the potential for low norms
self.assertLess(
(posterior_mean - sampled_mean).norm() / posterior_mean.norm(),
1e-1,
)
self.assertLess(
(posterior_variance - sampled_variance).norm() /
posterior_variance.norm(),
5e-2,
)
