def test_prune_baseline(self): no = "botorch.utils.testing.MockModel.num_outputs" prune = "botorch.acquisition.monte_carlo.prune_inferior_points" for dtype in (torch.float, torch.double): X_baseline = torch.zeros(1, 1, device=self.device, dtype=dtype) X_pruned = torch.rand(1, 1, device=self.device, dtype=dtype) with mock.patch( no, new_callable=mock.PropertyMock) as mock_num_outputs: mock_num_outputs.return_value = 1 mm = MockModel(mock.Mock()) with mock.patch(prune, return_value=X_pruned) as mock_prune: acqf = qNoisyExpectedImprovement( model=mm, X_baseline=X_baseline, prune_baseline=True, cache_root=False, ) mock_prune.assert_called_once() self.assertTrue(torch.equal(acqf.X_baseline, X_pruned)) with mock.patch(prune, return_value=X_pruned) as mock_prune: acqf = qNoisyExpectedImprovement( model=mm, X_baseline=X_baseline, prune_baseline=True, marginalize_dim=-3, cache_root=False, ) _, kwargs = mock_prune.call_args self.assertEqual(kwargs["marginalize_dim"], -3)
def test_q_noisy_expected_improvement(self, cuda=False): device = torch.device("cuda") if cuda else torch.device("cpu") for dtype in (torch.float, torch.double): # the event shape is `b x q x t` = 1 x 2 x 1 samples_noisy = torch.tensor([1.0, 0.0], device=device, dtype=dtype) samples_noisy = samples_noisy.view(1, 2, 1) # X_baseline is `q' x d` = 1 x 1 X_baseline = torch.zeros(1, 1, device=device, dtype=dtype) mm_noisy = MockModel(MockPosterior(samples=samples_noisy)) # X is `q x d` = 1 x 1 X = torch.zeros(1, 1, device=device, dtype=dtype) # basic test sampler = IIDNormalSampler(num_samples=2) acqf = qNoisyExpectedImprovement( model=mm_noisy, X_baseline=X_baseline, sampler=sampler ) res = acqf(X) self.assertEqual(res.item(), 1.0) # basic test, no resample sampler = IIDNormalSampler(num_samples=2, seed=12345) acqf = qNoisyExpectedImprovement( model=mm_noisy, X_baseline=X_baseline, sampler=sampler ) res = acqf(X) self.assertEqual(res.item(), 1.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # basic test, qmc, no resample sampler = SobolQMCNormalSampler(num_samples=2) acqf = qNoisyExpectedImprovement( model=mm_noisy, X_baseline=X_baseline, sampler=sampler ) res = acqf(X) self.assertEqual(res.item(), 1.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # basic test, qmc, resample sampler = SobolQMCNormalSampler(num_samples=2, resample=True, seed=12345) acqf = qNoisyExpectedImprovement( model=mm_noisy, X_baseline=X_baseline, sampler=sampler ) res = acqf(X) self.assertEqual(res.item(), 1.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertFalse(torch.equal(acqf.sampler.base_samples, bs))
def test_prune_baseline(self): prune = "botorch.acquisition.monte_carlo.prune_inferior_points" for dtype in (torch.float, torch.double): X_baseline = torch.zeros(1, 1, device=self.device, dtype=dtype) X_pruned = torch.rand(1, 1, device=self.device, dtype=dtype) mm = MockModel(mock.Mock()) with mock.patch(prune, return_value=X_pruned) as mock_prune: acqf = qNoisyExpectedImprovement( model=mm, X_baseline=X_baseline, prune_baseline=True ) mock_prune.assert_called_once() self.assertTrue(torch.equal(acqf.X_baseline, X_pruned))
def test_acquisition_functions(self): tkwargs = {"device": self.device, "dtype": torch.double} train_X, train_Y, train_Yvar, model = self._get_data_and_model( infer_noise=True, **tkwargs ) fit_fully_bayesian_model_nuts( model, warmup_steps=8, num_samples=5, thinning=2, disable_progbar=True ) sampler = IIDNormalSampler(num_samples=2) acquisition_functions = [ ExpectedImprovement(model=model, best_f=train_Y.max()), ProbabilityOfImprovement(model=model, best_f=train_Y.max()), PosteriorMean(model=model), UpperConfidenceBound(model=model, beta=4), qExpectedImprovement(model=model, best_f=train_Y.max(), sampler=sampler), qNoisyExpectedImprovement(model=model, X_baseline=train_X, sampler=sampler), qProbabilityOfImprovement( model=model, best_f=train_Y.max(), sampler=sampler ), qSimpleRegret(model=model, sampler=sampler), qUpperConfidenceBound(model=model, beta=4, sampler=sampler), qNoisyExpectedHypervolumeImprovement( model=ModelListGP(model, model), X_baseline=train_X, ref_point=torch.zeros(2, **tkwargs), sampler=sampler, ), qExpectedHypervolumeImprovement( model=ModelListGP(model, model), ref_point=torch.zeros(2, **tkwargs), sampler=sampler, partitioning=NondominatedPartitioning( ref_point=torch.zeros(2, **tkwargs), Y=train_Y.repeat([1, 2]) ), ), ] for acqf in acquisition_functions: for batch_shape in [[5], [6, 5, 2]]: test_X = torch.rand(*batch_shape, 1, 4, **tkwargs) self.assertEqual(acqf(test_X).shape, torch.Size(batch_shape))
def main(argv): dataset = 1 try: opts, args = getopt.getopt(argv, "hd:", ["dataset="]) except getopt.GetoptError: print('random parallel with input dataset') sys.exit(2) for opt, arg in opts: if opt == '-h': print('random parallel with input dataset') sys.exit() elif opt in ("-d", "--dataset"): dataset = int(arg) # average over multiple trials for trial in range(1, N_TRIALS + 1): print(f"\nTrial {trial:>2} of {N_TRIALS} ", end="") best_observed_ei, best_observed_nei = [], [] # call helper functions to generate initial training data and initialize model train_x_ei, train_obj_ei, best_observed_value_ei, current_best_config = generate_initial_data( dataset) train_x_nei, train_obj_nei = train_x_ei, train_obj_ei best_observed_value_nei = best_observed_value_ei mll_nei, model_nei = initialize_model(train_x_nei, train_obj_nei) best_observed_nei.append(best_observed_value_nei) # run N_BATCH rounds of BayesOpt after the initial random batch for iteration in range(1, N_BATCH + 1): # fit the models fit_gpytorch_model(mll_nei) # define the qEI and qNEI acquisition modules using a QMC sampler qmc_sampler = SobolQMCNormalSampler(num_samples=MC_SAMPLES) # for best_f, we use the best observed noisy values as an approximation qNEI = qNoisyExpectedImprovement( model=model_nei, X_baseline=train_x_nei, sampler=qmc_sampler, ) # optimize and get new observation new_x_nei, new_obj_nei = optimize_acqf_and_get_observation( qNEI, dataset) # update training points train_x_nei = torch.cat([train_x_nei, new_x_nei]) train_obj_nei = torch.cat([train_obj_nei, new_obj_nei]) # update progress best_value_nei = train_obj_nei.max().item() best_observed_nei.append(best_value_nei) # reinitialize the models so they are ready for fitting on next iteration # use the current state dict to speed up fitting mll_nei, model_nei = initialize_model( train_x_nei, train_obj_nei, model_nei.state_dict(), ) # return the best configuration best_tensor_nei, indices_nei = torch.max(train_obj_nei, 0) train_best_x_nei = train_x_nei[indices_nei].cpu().numpy() from botorch.acquisition import PosteriorMean argmax_pmean_nei, max_pmean_nei = optimize_acqf( acq_function=PosteriorMean(model_nei), bounds=bounds, q=1, num_restarts=20, raw_samples=2048, ) csv_file_name = '/home/junjie/modes/botorch/' + folder_name + '/modes-i/hp-ngp-qnei-dataset-' + str( dataset) + '-trail' + str(trial) + '.csv' with open(csv_file_name, 'w') as csvFile: writer = csv.writer(csvFile) writer.writerow([ str(argmax_pmean_nei.cpu().numpy()), str(max_pmean_nei.cpu().numpy()) ]) # nei prediction writer.writerow( [str(train_best_x_nei), str(best_tensor_nei.cpu().numpy())]) # nei observation csvFile.close()
def get_acquisition_function( acquisition_function_name: str, model: Model, objective: MCAcquisitionObjective, X_observed: Tensor, X_pending: Optional[Tensor] = None, constraints: Optional[List[Callable[[Tensor], Tensor]]] = None, mc_samples: int = 500, qmc: bool = True, seed: Optional[int] = None, **kwargs, ) -> monte_carlo.MCAcquisitionFunction: r"""Convenience function for initializing botorch acquisition functions. Args: acquisition_function_name: Name of the acquisition function. model: A fitted model. objective: A MCAcquisitionObjective. X_observed: A `m1 x d`-dim Tensor of `m1` design points that have already been observed. X_pending: A `m2 x d`-dim Tensor of `m2` design points whose evaluation is pending. constraints: A list of callables, each mapping a Tensor of dimension `sample_shape x batch-shape x q x m` to a Tensor of dimension `sample_shape x batch-shape x q`, where negative values imply feasibility. Used when constraint_transforms are not passed as part of the objective. mc_samples: The number of samples to use for (q)MC evaluation of the acquisition function. qmc: If True, use quasi-Monte-Carlo sampling (instead of iid). seed: If provided, perform deterministic optimization (i.e. the function to optimize is fixed and not stochastic). Returns: The requested acquisition function. Example: >>> model = SingleTaskGP(train_X, train_Y) >>> obj = LinearMCObjective(weights=torch.tensor([1.0, 2.0])) >>> acqf = get_acquisition_function("qEI", model, obj, train_X) """ # initialize the sampler if qmc: sampler = SobolQMCNormalSampler(num_samples=mc_samples, seed=seed) else: sampler = IIDNormalSampler(num_samples=mc_samples, seed=seed) # instantiate and return the requested acquisition function if acquisition_function_name == "qEI": best_f = objective(model.posterior(X_observed).mean).max().item() return monte_carlo.qExpectedImprovement( model=model, best_f=best_f, sampler=sampler, objective=objective, X_pending=X_pending, ) elif acquisition_function_name == "qPI": best_f = objective(model.posterior(X_observed).mean).max().item() return monte_carlo.qProbabilityOfImprovement( model=model, best_f=best_f, sampler=sampler, objective=objective, X_pending=X_pending, tau=kwargs.get("tau", 1e-3), ) elif acquisition_function_name == "qNEI": return monte_carlo.qNoisyExpectedImprovement( model=model, X_baseline=X_observed, sampler=sampler, objective=objective, X_pending=X_pending, prune_baseline=kwargs.get("prune_baseline", False), ) elif acquisition_function_name == "qSR": return monte_carlo.qSimpleRegret(model=model, sampler=sampler, objective=objective, X_pending=X_pending) elif acquisition_function_name == "qUCB": if "beta" not in kwargs: raise ValueError("`beta` must be specified in kwargs for qUCB.") return monte_carlo.qUpperConfidenceBound( model=model, beta=kwargs["beta"], sampler=sampler, objective=objective, X_pending=X_pending, ) elif acquisition_function_name == "qEHVI": # pyre-fixme [16]: `Model` has no attribute `train_targets` try: ref_point = kwargs["ref_point"] except KeyError: raise ValueError( "`ref_point` must be specified in kwargs for qEHVI") try: Y = kwargs["Y"] except KeyError: raise ValueError("`Y` must be specified in kwargs for qEHVI") # get feasible points if constraints is not None: feas = torch.stack([c(Y) <= 0 for c in constraints], dim=-1).all(dim=-1) Y = Y[feas] obj = objective(Y) partitioning = NondominatedPartitioning( ref_point=torch.as_tensor(ref_point, dtype=Y.dtype, device=Y.device), Y=obj, alpha=kwargs.get("alpha", 0.0), ) return moo_monte_carlo.qExpectedHypervolumeImprovement( model=model, ref_point=ref_point, partitioning=partitioning, sampler=sampler, objective=objective, constraints=constraints, X_pending=X_pending, ) raise NotImplementedError( f"Unknown acquisition function {acquisition_function_name}")
def test_q_noisy_expected_improvement_batch(self): for dtype in (torch.float, torch.double): # the event shape is `b x q x t` = 2 x 3 x 1 samples_noisy = torch.zeros(2, 3, 1, device=self.device, dtype=dtype) samples_noisy[0, 0, 0] = 1.0 mm_noisy = MockModel(MockPosterior(samples=samples_noisy)) # X is `q x d` = 1 x 1 X = torch.zeros(1, 1, 1, device=self.device, dtype=dtype) X_baseline = torch.zeros(1, 1, device=self.device, dtype=dtype) # test batch mode sampler = IIDNormalSampler(num_samples=2) acqf = qNoisyExpectedImprovement( model=mm_noisy, X_baseline=X_baseline, sampler=sampler ) res = acqf(X) self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) # test batch mode, no resample sampler = IIDNormalSampler(num_samples=2, seed=12345) acqf = qNoisyExpectedImprovement( model=mm_noisy, X_baseline=X_baseline, sampler=sampler ) res = acqf(X) # 1-dim batch self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) res = acqf(X.expand(2, 1, 1)) # 2-dim batch self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) # the base samples should have the batch dim collapsed self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1])) bs = acqf.sampler.base_samples.clone() acqf(X.expand(2, 1, 1)) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # test batch mode, qmc, no resample sampler = SobolQMCNormalSampler(num_samples=2) acqf = qNoisyExpectedImprovement( model=mm_noisy, X_baseline=X_baseline, sampler=sampler ) res = acqf(X) self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # test X_pending w/ batch mode, qmc, resample sampler = SobolQMCNormalSampler(num_samples=2, resample=True, seed=12345) acqf = qNoisyExpectedImprovement( model=mm_noisy, X_baseline=X_baseline, sampler=sampler ) res = acqf(X) # 1-dim batch self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertFalse(torch.equal(acqf.sampler.base_samples, bs)) res = acqf(X.expand(2, 1, 1)) # 2-dim batch self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) # the base samples should have the batch dim collapsed self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1])) bs = acqf.sampler.base_samples.clone() acqf(X.expand(2, 1, 1)) self.assertFalse(torch.equal(acqf.sampler.base_samples, bs))
def test_q_noisy_expected_improvement(self): for dtype in (torch.float, torch.double): # the event shape is `b x q x t` = 1 x 2 x 1 samples_noisy = torch.tensor([1.0, 0.0], device=self.device, dtype=dtype) samples_noisy = samples_noisy.view(1, 2, 1) # X_baseline is `q' x d` = 1 x 1 X_baseline = torch.zeros(1, 1, device=self.device, dtype=dtype) mm_noisy = MockModel(MockPosterior(samples=samples_noisy)) # X is `q x d` = 1 x 1 X = torch.zeros(1, 1, device=self.device, dtype=dtype) # basic test sampler = IIDNormalSampler(num_samples=2) acqf = qNoisyExpectedImprovement( model=mm_noisy, X_baseline=X_baseline, sampler=sampler ) res = acqf(X) self.assertEqual(res.item(), 1.0) # basic test, no resample sampler = IIDNormalSampler(num_samples=2, seed=12345) acqf = qNoisyExpectedImprovement( model=mm_noisy, X_baseline=X_baseline, sampler=sampler ) res = acqf(X) self.assertEqual(res.item(), 1.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # basic test, qmc, no resample sampler = SobolQMCNormalSampler(num_samples=2) acqf = qNoisyExpectedImprovement( model=mm_noisy, X_baseline=X_baseline, sampler=sampler ) res = acqf(X) self.assertEqual(res.item(), 1.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # basic test, qmc, resample sampler = SobolQMCNormalSampler(num_samples=2, resample=True, seed=12345) acqf = qNoisyExpectedImprovement( model=mm_noisy, X_baseline=X_baseline, sampler=sampler ) res = acqf(X) self.assertEqual(res.item(), 1.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertFalse(torch.equal(acqf.sampler.base_samples, bs)) # basic test for X_pending and warning sampler = SobolQMCNormalSampler(num_samples=2) samples_noisy_pending = torch.tensor( [1.0, 0.0, 0.0], device=self.device, dtype=dtype ) samples_noisy_pending = samples_noisy_pending.view(1, 3, 1) mm_noisy_pending = MockModel(MockPosterior(samples=samples_noisy_pending)) acqf = qNoisyExpectedImprovement( model=mm_noisy_pending, X_baseline=X_baseline, sampler=sampler ) acqf.set_X_pending() self.assertIsNone(acqf.X_pending) acqf.set_X_pending(None) self.assertIsNone(acqf.X_pending) acqf.set_X_pending(X) self.assertEqual(acqf.X_pending, X) res = acqf(X) X2 = torch.zeros( 1, 1, 1, device=self.device, dtype=dtype, requires_grad=True ) with warnings.catch_warnings(record=True) as ws, settings.debug(True): acqf.set_X_pending(X2) self.assertEqual(acqf.X_pending, X2) self.assertEqual(len(ws), 1) self.assertTrue(issubclass(ws[-1].category, BotorchWarning))
def optimize(): verbose = True best_observed_all_ei, best_observed_all_nei, best_random_all = [], [], [] train_x_all_ei, train_x_all_nei, train_x_all_random = [], [], [] train_y_all_ei, train_y_all_nei, train_y_all_random = [], [], [] # statistics over multiple trials for trial in range(1, N_TRIALS + 1): print('\nTrial {} of {}'.format(trial, N_TRIALS)) best_observed_ei, best_observed_nei = [], [] best_random = [] # generate initial training data and initialize model print('\nGenerating {} random samples'.format(N_INITIAL_SAMPLES)) train_x_ei, train_y_ei, best_y_ei, mean_y, std_y = generate_initial_data( n_samples=N_INITIAL_SAMPLES) denormalize = lambda x: -(x * std_y + mean_y) mll_ei, model_ei = initialize_model(train_x_ei, train_y_ei) train_x_nei, train_y_nei, best_y_nei = train_x_ei, train_y_ei, best_y_ei mll_nei, model_nei = initialize_model(train_x_nei, train_y_nei) train_x_random, train_y_random, best_y_random = train_x_ei, train_y_ei, best_y_ei best_observed_ei.append(denormalize(best_y_ei)) best_observed_nei.append(denormalize(best_y_nei)) best_random.append(denormalize(best_y_random)) # run N_BATCH rounds of BayesOpt after the initial random batch for iteration in range(1, N_BATCH + 1): print('\nBatch {} of {}\n'.format(iteration, N_BATCH)) t0 = time.time() # fit the models fit_gpytorch_model(mll_ei) fit_gpytorch_model(mll_nei) # update acquisition functions qEI = qExpectedImprovement( model=model_ei, best_f=train_y_ei.max(), sampler=qmc_sampler, ) qNEI = qNoisyExpectedImprovement( model=model_nei, X_baseline=train_x_nei, sampler=qmc_sampler, ) # optimize acquisition function and evaluate new sample new_x_ei, new_y_ei = optimize_acqf_and_get_observation( qEI, mean_y=mean_y, std_y=std_y) print('EI: time to traverse is {:.4f}s'.format( -(new_y_ei.numpy().ravel()[0] * std_y + mean_y))) new_x_nei, new_y_nei = optimize_acqf_and_get_observation( qNEI, mean_y=mean_y, std_y=std_y) print('NEI: time to traverse is {:.4f}s'.format( -(new_y_nei.numpy().ravel()[0] * std_y + mean_y))) new_x_random, new_y_random = sample_random_observations( mean_y=mean_y, std_y=std_y) print('Random: time to traverse is {:.4f}s'.format( -(new_y_random.numpy().ravel()[0] * std_y + mean_y))) # update training points train_x_ei = torch.cat([train_x_ei, new_x_ei]) train_y_ei = torch.cat([train_y_ei, new_y_ei]) train_x_nei = torch.cat([train_x_nei, new_x_nei]) train_y_nei = torch.cat([train_y_nei, new_y_nei]) train_x_random = torch.cat([train_x_random, new_x_random]) train_y_random = torch.cat([train_y_random, new_y_random]) # update progress best_value_ei = denormalize(train_y_ei.max().item()) best_value_nei = denormalize(train_y_nei.max().item()) best_value_random = denormalize(train_y_random.max().item()) best_observed_ei.append(best_value_ei) best_observed_nei.append(best_value_nei) best_random.append(best_value_random) # reinitialize the models so they are ready for fitting on next iteration # use the current state dict to speed up fitting mll_ei, model_ei = initialize_model( train_x_ei, train_y_ei, model_ei.state_dict(), ) mll_nei, model_nei = initialize_model( train_x_nei, train_y_nei, model_nei.state_dict(), ) t1 = time.time() if verbose: print( 'best lap time (random, qEI, qNEI) = {:.2f}, {:.2f}, {:.2f}, time to compute = {:.2f}s' .format(best_value_random, best_value_ei, best_value_nei, t1 - t0)) else: print(".") best_observed_all_ei.append(best_observed_ei) best_observed_all_nei.append(best_observed_nei) best_random_all.append(best_random) train_x_all_ei.append(train_x_ei.cpu().numpy()) train_x_all_nei.append(train_x_nei.cpu().numpy()) train_x_all_random.append(train_x_random.cpu().numpy()) train_y_all_ei.append(denormalize(train_y_ei.cpu().numpy())) train_y_all_nei.append(denormalize(train_y_nei.cpu().numpy())) train_y_all_random.append(denormalize(train_y_random.cpu().numpy())) iters = np.arange(N_BATCH + 1) * BATCH_SIZE y_ei = np.asarray(best_observed_all_ei) y_nei = np.asarray(best_observed_all_nei) y_rnd = np.asarray(best_random_all) savestr = time.strftime('%Y%m%d%H%M%S') ##################################################################### # save results if SAVE_RESULTS: np.savez( 'results/{}_raceline_data-{}.npz'.format('UCB', savestr), y_ei=y_ei, y_nei=y_nei, y_rnd=y_rnd, iters=iters, train_x_all_ei=np.asarray(train_x_all_ei), train_x_all_nei=np.asarray(train_x_all_nei), train_x_all_random=np.asarray(train_x_all_random), train_y_all_ei=np.asarray(train_y_all_ei), train_y_all_nei=np.asarray(train_y_all_nei), train_y_all_random=np.asarray(train_y_all_random), SEED=SEED, ) ##################################################################### # plot results if PLOT_RESULTS: def ci(y): return 1.96 * y.std(axis=0) / np.sqrt(N_TRIALS) plt.figure() plt.gca().set_prop_cycle(None) plt.plot(iters, y_rnd.mean(axis=0), linewidth=1.5) plt.plot(iters, y_ei.mean(axis=0), linewidth=1.5) plt.plot(iters, y_nei.mean(axis=0), linewidth=1.5) plt.gca().set_prop_cycle(None) plt.fill_between(iters, y_rnd.mean(axis=0) - ci(y_rnd), y_rnd.mean(axis=0) + ci(y_rnd), label='random', alpha=0.2) plt.fill_between(iters, y_ei.mean(axis=0) - ci(y_ei), y_ei.mean(axis=0) + ci(y_ei), label='qEI', alpha=0.2) plt.fill_between(iters, y_nei.mean(axis=0) - ci(y_nei), y_nei.mean(axis=0) + ci(y_nei), label='qNEI', alpha=0.2) plt.xlabel('number of observations (beyond initial points)') plt.ylabel('best lap times') plt.grid(True) plt.legend(loc=0) plt.savefig('results/{}_laptimes-{}.png'.format('UCB', savestr), dpi=600) plt.show()
def test_q_noisy_expected_improvement(self, cuda=False): device = torch.device("cuda") if cuda else torch.device("cpu") for dtype in (torch.float, torch.double): # the event shape is `b x q x t` = 1 x 2 x 1 samples_noisy = torch.tensor([1.0, 0.0], device=device, dtype=dtype) samples_noisy = samples_noisy.view(1, 2, 1) # X_baseline is `q' x d` = 1 x 1 X_baseline = torch.zeros(1, 1, device=device, dtype=dtype) mm_noisy = MockModel(MockPosterior(samples=samples_noisy)) # X is `q x d` = 1 x 1 X = torch.zeros(1, 1, device=device, dtype=dtype) # basic test sampler = IIDNormalSampler(num_samples=2) acqf = qNoisyExpectedImprovement(model=mm_noisy, X_baseline=X_baseline, sampler=sampler) res = acqf(X) self.assertEqual(res.item(), 1.0) # basic test, no resample sampler = IIDNormalSampler(num_samples=2, seed=12345) acqf = qNoisyExpectedImprovement(model=mm_noisy, X_baseline=X_baseline, sampler=sampler) res = acqf(X) self.assertEqual(res.item(), 1.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # basic test, qmc, no resample sampler = SobolQMCNormalSampler(num_samples=2) acqf = qNoisyExpectedImprovement(model=mm_noisy, X_baseline=X_baseline, sampler=sampler) res = acqf(X) self.assertEqual(res.item(), 1.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # basic test, qmc, resample sampler = SobolQMCNormalSampler(num_samples=2, resample=True, seed=12345) acqf = qNoisyExpectedImprovement(model=mm_noisy, X_baseline=X_baseline, sampler=sampler) res = acqf(X) self.assertEqual(res.item(), 1.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertFalse(torch.equal(acqf.sampler.base_samples, bs))
def run_RGPE(test_task: int, objective, bounds, base_model_list): input_dim = bounds.shape[1] best_rgpe_all = [] best_argmax_rgpe_all = [] # Average over multiple trials for trial in range(N_TRIALS): print(f"Trial {trial + 1} of {N_TRIALS}") best_BMs = [] best_rgpe = [] # Initial random observations raw_x = draw_sobol_samples(bounds=bounds, n=RANDOM_INITIALIZATION_SIZE, q=1, seed=trial).squeeze(1) train_x = normalize(raw_x, bounds=bounds) train_y_noiseless = objective(raw_x, shift=test_task) train_y = train_y_noiseless + noise_std * torch.randn_like( train_y_noiseless) train_yvar = torch.full_like(train_y, noise_std**2) # keep track of the best observed point at each iteration best_value = train_y.max().item() best_rgpe.append(best_value) # Run N_BATCH rounds of BayesOpt after the initial random batch for iteration in range(N_BATCH): target_model = get_fitted_model(train_x, train_y, train_yvar) model_list = base_model_list + [target_model] rank_weights = compute_rank_weights( train_x, train_y, base_model_list, target_model, NUM_POSTERIOR_SAMPLES, ) # create model and acquisition function rgpe_model = RGPE(model_list, rank_weights) sampler_qnei = SobolQMCNormalSampler(num_samples=MC_SAMPLES) qNEI = qNoisyExpectedImprovement( model=rgpe_model, X_baseline=train_x, sampler=sampler_qnei, ) # optimize candidate, _ = optimize_acqf( acq_function=qNEI, bounds=bounds, q=Q_BATCH_SIZE, num_restarts=N_RESTARTS, raw_samples=N_RESTART_CANDIDATES, ) # fetch the new values new_x = candidate.detach() new_y_noiseless = objective(unnormalize(new_x, bounds=bounds), shift=test_task) new_y = new_y_noiseless + noise_std * torch.randn_like( new_y_noiseless) new_yvar = torch.full_like(new_y, noise_std**2) # update training points train_x = torch.cat((train_x, new_x)) train_y = torch.cat((train_y, new_y)) train_yvar = torch.cat((train_yvar, new_yvar)) # get the new best observed value best_value = train_y.max().item() best_idx = torch.argmax(train_y).item() best_candidate = train_x[best_idx].view(1, -1) _, best_BM = objective(unnormalize(best_candidate, bounds=bounds), shift=test_task, include_BMs=True) best_rgpe.append(best_value) best_BMs.append(best_BM) best_rgpe_all.append(best_rgpe) best_argmax_rgpe_all.append(best_BMs) BM_winner_idx = np.argmax(np.array(best_rgpe_all)[:, -1], axis=0) BM_winner = np.reshape(np.array(best_argmax_rgpe_all[BM_winner_idx][-1]), (2, input_dim)) return BM_winner
def test_sample_points_around_best(self): tkwargs = {"device": self.device} _bounds = torch.ones(2, 2) _bounds[1] = 2 for dtype in (torch.float, torch.double): tkwargs["dtype"] = dtype bounds = _bounds.to(**tkwargs) X_train = 1 + torch.rand(20, 2, **tkwargs) model = MockModel( MockPosterior(mean=(2 * X_train + 1).sum(dim=-1, keepdim=True)) ) # test NEI with X_baseline acqf = qNoisyExpectedImprovement(model, X_baseline=X_train) with mock.patch( "botorch.optim.initializers.sample_perturbed_subset_dims" ) as mock_subset_dims: X_rnd = sample_points_around_best( acq_function=acqf, n_discrete_points=4, sigma=1e-3, bounds=bounds, ) mock_subset_dims.assert_not_called() self.assertTrue(X_rnd.shape, torch.Size([4, 2])) self.assertTrue((X_rnd >= 1).all()) self.assertTrue((X_rnd <= 2).all()) # test model that returns a batched mean model = MockModel( MockPosterior( mean=(2 * X_train + 1).sum(dim=-1, keepdim=True).unsqueeze(0) ) ) acqf = qNoisyExpectedImprovement(model, X_baseline=X_train) X_rnd = sample_points_around_best( acq_function=acqf, n_discrete_points=4, sigma=1e-3, bounds=bounds, ) self.assertTrue(X_rnd.shape, torch.Size([4, 2])) self.assertTrue((X_rnd >= 1).all()) self.assertTrue((X_rnd <= 2).all()) # test EI without X_baseline acqf = qExpectedImprovement(model, best_f=0.0) with warnings.catch_warnings(record=True) as w, settings.debug(True): X_rnd = sample_points_around_best( acq_function=acqf, n_discrete_points=4, sigma=1e-3, bounds=bounds, ) self.assertEqual(len(w), 1) self.assertTrue(issubclass(w[-1].category, BotorchWarning)) self.assertIsNone(X_rnd) # set train inputs model.train_inputs = (X_train,) X_rnd = sample_points_around_best( acq_function=acqf, n_discrete_points=4, sigma=1e-3, bounds=bounds, ) self.assertTrue(X_rnd.shape, torch.Size([4, 2])) self.assertTrue((X_rnd >= 1).all()) self.assertTrue((X_rnd <= 2).all()) # test an acquisition function that has objective=None # and maximize=False pm = PosteriorMean(model, maximize=False) self.assertIsNone(pm.objective) self.assertFalse(pm.maximize) X_rnd = sample_points_around_best( acq_function=pm, n_discrete_points=4, sigma=0, bounds=bounds, best_pct=1e-8, # ensures that we only use best value ) idx = (-model.posterior(X_train).mean).argmax() self.assertTrue((X_rnd == X_train[idx : idx + 1]).all(dim=-1).all()) # test acquisition function that has no model ff = FixedFeatureAcquisitionFunction(pm, d=2, columns=[0], values=[0]) # set X_baseline for testing purposes ff.X_baseline = X_train with warnings.catch_warnings(record=True) as w, settings.debug(True): X_rnd = sample_points_around_best( acq_function=ff, n_discrete_points=4, sigma=1e-3, bounds=bounds, ) self.assertEqual(len(w), 1) self.assertTrue(issubclass(w[-1].category, BotorchWarning)) self.assertIsNone(X_rnd) # test constraints with NEHVI constraints = [lambda Y: Y[..., 0]] ref_point = torch.zeros(2, **tkwargs) # test cases when there are and are not any feasible points for any_feas in (True, False): Y_train = torch.stack( [ torch.linspace(-0.5, 0.5, X_train.shape[0], **tkwargs) if any_feas else torch.ones(X_train.shape[0], **tkwargs), X_train.sum(dim=-1), ], dim=-1, ) moo_model = MockModel(MockPosterior(mean=Y_train, samples=Y_train)) acqf = qNoisyExpectedHypervolumeImprovement( moo_model, ref_point=ref_point, X_baseline=X_train, constraints=constraints, cache_root=False, ) X_rnd = sample_points_around_best( acq_function=acqf, n_discrete_points=4, sigma=0.0, bounds=bounds, ) self.assertTrue(X_rnd.shape, torch.Size([4, 2])) # this should be true since sigma=0 # and we should only be returning feasible points violation = constraints[0](Y_train) neg_violation = -violation.clamp_min(0.0) feas = neg_violation == 0 eq_mask = (X_train.unsqueeze(1) == X_rnd.unsqueeze(0)).all(dim=-1) if feas.any(): # determine # create n_train x n_rnd tensor of booleans eq_mask = (X_train.unsqueeze(1) == X_rnd.unsqueeze(0)).all(dim=-1) # check that all X_rnd correspond to feasible points self.assertEqual(eq_mask[feas].sum(), 4) else: idcs = torch.topk(neg_violation, k=2).indices self.assertEqual(eq_mask[idcs].sum(), 4) self.assertTrue((X_rnd >= 1).all()) self.assertTrue((X_rnd <= 2).all()) # test that subset_dims is called if d>=21 X_train = 1 + torch.rand(20, 21, **tkwargs) model = MockModel( MockPosterior(mean=(2 * X_train + 1).sum(dim=-1, keepdim=True)) ) bounds = torch.ones(2, 21, **tkwargs) bounds[1] = 2 # test NEI with X_baseline acqf = qNoisyExpectedImprovement(model, X_baseline=X_train) with mock.patch( "botorch.optim.initializers.sample_perturbed_subset_dims", wraps=sample_perturbed_subset_dims, ) as mock_subset_dims: X_rnd = sample_points_around_best( acq_function=acqf, n_discrete_points=5, sigma=1e-3, bounds=bounds ) self.assertTrue(X_rnd.shape, torch.Size([5, 2])) self.assertTrue((X_rnd >= 1).all()) self.assertTrue((X_rnd <= 2).all()) mock_subset_dims.assert_called_once() # test tiny prob_perturb to make sure we perturb at least one dimension X_rnd = sample_points_around_best( acq_function=acqf, n_discrete_points=5, sigma=1e-3, bounds=bounds, prob_perturb=1e-8, ) self.assertTrue( ((X_rnd.unsqueeze(0) == X_train.unsqueeze(1)).all(dim=-1)).sum() == 0 )
def test_get_X_baseline(self): tkwargs = {"device": self.device} for dtype in (torch.float, torch.double): tkwargs["dtype"] = dtype X_train = torch.rand(20, 2, **tkwargs) model = MockModel( MockPosterior(mean=(2 * X_train + 1).sum(dim=-1, keepdim=True))) # test NEI with X_baseline acqf = qNoisyExpectedImprovement(model, X_baseline=X_train[:2]) X = get_X_baseline(acq_function=acqf) self.assertTrue(torch.equal(X, acqf.X_baseline)) # test EI without X_baseline acqf = qExpectedImprovement(model, best_f=0.0) with warnings.catch_warnings( record=True) as w, settings.debug(True): X_rnd = get_X_baseline(acq_function=acqf, ) self.assertEqual(len(w), 1) self.assertTrue(issubclass(w[-1].category, BotorchWarning)) self.assertIsNone(X_rnd) # set train inputs model.train_inputs = (X_train, ) X = get_X_baseline(acq_function=acqf, ) self.assertTrue(torch.equal(X, X_train)) # test that we fail back to train_inputs if X_baseline is an empty tensor acqf.register_buffer("X_baseline", X_train[:0]) X = get_X_baseline(acq_function=acqf, ) self.assertTrue(torch.equal(X, X_train)) # test acquisitipon function without X_baseline or model acqf = FixedFeatureAcquisitionFunction(acqf, d=2, columns=[0], values=[0]) with warnings.catch_warnings( record=True) as w, settings.debug(True): X_rnd = get_X_baseline(acq_function=acqf, ) self.assertEqual(len(w), 1) self.assertTrue(issubclass(w[-1].category, BotorchWarning)) self.assertIsNone(X_rnd) Y_train = 2 * X_train[:2] + 1 moo_model = MockModel(MockPosterior(mean=Y_train, samples=Y_train)) ref_point = torch.zeros(2, **tkwargs) # test NEHVI with X_baseline acqf = qNoisyExpectedHypervolumeImprovement( moo_model, ref_point=ref_point, X_baseline=X_train[:2], cache_root=False, ) X = get_X_baseline(acq_function=acqf, ) self.assertTrue(torch.equal(X, acqf.X_baseline)) # test qEHVI without train_inputs acqf = qExpectedHypervolumeImprovement( moo_model, ref_point=ref_point, partitioning=FastNondominatedPartitioning( ref_point=ref_point, Y=Y_train, ), ) # test extracting train_inputs from model list GP model_list = ModelListGP( SingleTaskGP(X_train, Y_train[:, :1]), SingleTaskGP(X_train, Y_train[:, 1:]), ) acqf = qExpectedHypervolumeImprovement( model_list, ref_point=ref_point, partitioning=FastNondominatedPartitioning( ref_point=ref_point, Y=Y_train, ), ) X = get_X_baseline(acq_function=acqf, ) self.assertTrue(torch.equal(X, X_train)) # test MESMO for which we need to use # `acqf.mo_model` batched_mo_model = SingleTaskGP(X_train, Y_train) acqf = qMultiObjectiveMaxValueEntropy( batched_mo_model, sample_pareto_frontiers=lambda model: torch.rand( 10, 2, **tkwargs), ) X = get_X_baseline(acq_function=acqf, ) self.assertTrue(torch.equal(X, X_train)) # test that if there is an input transform that is applied # to the train_inputs when the model is in eval mode, we # extract the untransformed train_inputs model = SingleTaskGP(X_train, Y_train[:, :1], input_transform=Warp(indices=[0, 1])) model.eval() self.assertFalse(torch.equal(model.train_inputs[0], X_train)) acqf = qExpectedImprovement(model, best_f=0.0) X = get_X_baseline(acq_function=acqf, ) self.assertTrue(torch.equal(X, X_train))
def test_q_noisy_expected_improvement_batch(self, cuda=False): device = torch.device("cuda") if cuda else torch.device("cpu") for dtype in (torch.float, torch.double): # the event shape is `b x q x t` = 2 x 3 x 1 samples_noisy = torch.zeros(2, 3, 1, device=device, dtype=dtype) samples_noisy[0, 0, 0] = 1.0 mm_noisy = MockModel(MockPosterior(samples=samples_noisy)) # X is `q x d` = 1 x 1 X = torch.zeros(1, 1, 1, device=device, dtype=dtype) X_baseline = torch.zeros(1, 1, device=device, dtype=dtype) # test batch mode sampler = IIDNormalSampler(num_samples=2) acqf = qNoisyExpectedImprovement( model=mm_noisy, X_baseline=X_baseline, sampler=sampler ) res = acqf(X) self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) # test batch mode, no resample sampler = IIDNormalSampler(num_samples=2, seed=12345) acqf = qNoisyExpectedImprovement( model=mm_noisy, X_baseline=X_baseline, sampler=sampler ) res = acqf(X) # 1-dim batch self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) res = acqf(X.expand(2, 1, 1)) # 2-dim batch self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) # the base samples should have the batch dim collapsed self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1])) bs = acqf.sampler.base_samples.clone() acqf(X.expand(2, 1, 1)) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # test batch mode, qmc, no resample sampler = SobolQMCNormalSampler(num_samples=2) acqf = qNoisyExpectedImprovement( model=mm_noisy, X_baseline=X_baseline, sampler=sampler ) res = acqf(X) self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertTrue(torch.equal(acqf.sampler.base_samples, bs)) # test X_pending w/ batch mode, qmc, resample sampler = SobolQMCNormalSampler(num_samples=2, resample=True, seed=12345) acqf = qNoisyExpectedImprovement( model=mm_noisy, X_baseline=X_baseline, sampler=sampler ) res = acqf(X) # 1-dim batch self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1])) bs = acqf.sampler.base_samples.clone() acqf(X) self.assertFalse(torch.equal(acqf.sampler.base_samples, bs)) res = acqf(X.expand(2, 1, 1)) # 2-dim batch self.assertEqual(res[0].item(), 1.0) self.assertEqual(res[1].item(), 0.0) # the base samples should have the batch dim collapsed self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1])) bs = acqf.sampler.base_samples.clone() acqf(X.expand(2, 1, 1)) self.assertFalse(torch.equal(acqf.sampler.base_samples, bs))
fit_gpytorch_model(mll_nei) # define the qEI and qNEI acquisition modules using a QMC sampler qmc_sampler = SobolQMCNormalSampler(num_samples=MC_SAMPLES) # for best_f, we use the best observed noisy values as an approximation qEI = qExpectedImprovement( model=model_ei, best_f=(train_obj_ei * (train_con_ei <= 0).to(train_obj_ei)).max(), sampler=qmc_sampler, objective=constrained_obj, ) qNEI = qNoisyExpectedImprovement( model=model_nei, X_baseline=train_x_nei, sampler=qmc_sampler, objective=constrained_obj, ) # optimize and get new observation new_x_ei, new_obj_ei, new_con_ei = optimize_acqf_and_get_observation( qEI) new_x_nei, new_obj_nei, new_con_nei = optimize_acqf_and_get_observation( qNEI) # update training points train_x_ei = torch.cat([train_x_ei, new_x_ei]) train_obj_ei = torch.cat([train_obj_ei, new_obj_ei]) train_con_ei = torch.cat([train_con_ei, new_con_ei]) train_x_nei = torch.cat([train_x_nei, new_x_nei])
def get_acquisition_function( acquisition_function_name: str, model: Model, objective: MCAcquisitionObjective, X_observed: Tensor, X_pending: Optional[Tensor] = None, mc_samples: int = 500, qmc: bool = True, seed: Optional[int] = None, **kwargs, ) -> monte_carlo.MCAcquisitionFunction: r"""Convenience function for initializing botorch acquisition functions. Args: acquisition_function_name: Name of the acquisition function. model: A fitted model. objective: A MCAcquisitionObjective. X_observed: A `m1 x d`-dim Tensor of `m1` design points that have already been observed. X_pending: A `m2 x d`-dim Tensor of `m2` design points whose evaluation is pending. mc_samples: The number of samples to use for (q)MC evaluation of the acquisition function. qmc: If True, use quasi-Monte-Carlo sampling (instead of iid). seed: If provided, perform deterministic optimization (i.e. the function to optimize is fixed and not stochastic). Returns: The requested acquisition function. Example: >>> model = SingleTaskGP(train_X, train_Y) >>> obj = LinearMCObjective(weights=torch.tensor([1.0, 2.0])) >>> acqf = get_acquisition_function("qEI", model, obj, train_X) """ # initialize the sampler if qmc: sampler = SobolQMCNormalSampler(num_samples=mc_samples, seed=seed) else: sampler = IIDNormalSampler(num_samples=mc_samples, seed=seed) # instantiate and return the requested acquisition function if acquisition_function_name == "qEI": best_f = objective(model.posterior(X_observed).mean).max().item() return monte_carlo.qExpectedImprovement( model=model, best_f=best_f, sampler=sampler, objective=objective, X_pending=X_pending, ) elif acquisition_function_name == "qPI": best_f = objective(model.posterior(X_observed).mean).max().item() return monte_carlo.qProbabilityOfImprovement( model=model, best_f=best_f, sampler=sampler, objective=objective, X_pending=X_pending, tau=kwargs.get("tau", 1e-3), ) elif acquisition_function_name == "qNEI": return monte_carlo.qNoisyExpectedImprovement( model=model, X_baseline=X_observed, sampler=sampler, objective=objective, X_pending=X_pending, prune_baseline=kwargs.get("prune_baseline", False), ) elif acquisition_function_name == "qSR": return monte_carlo.qSimpleRegret(model=model, sampler=sampler, objective=objective, X_pending=X_pending) elif acquisition_function_name == "qUCB": if "beta" not in kwargs: raise ValueError("`beta` must be specified in kwargs for qUCB.") return monte_carlo.qUpperConfidenceBound( model=model, beta=kwargs["beta"], sampler=sampler, objective=objective, X_pending=X_pending, ) raise NotImplementedError( f"Unknown acquisition function {acquisition_function_name}")
def test_cache_root(self): sample_cached_path = ( "botorch.acquisition.cached_cholesky.sample_cached_cholesky") raw_state_dict = { "likelihood.noise_covar.raw_noise": torch.tensor([[0.0895], [0.2594]], dtype=torch.float64), "mean_module.constant": torch.tensor([[-0.4545], [-0.1285]], dtype=torch.float64), "covar_module.raw_outputscale": torch.tensor([1.4876, 1.4897], dtype=torch.float64), "covar_module.base_kernel.raw_lengthscale": torch.tensor([[[-0.7202, -0.2868]], [[-0.8794, -1.2877]]], dtype=torch.float64), } # test batched models (e.g. for MCMC) for train_batch_shape, m, dtype in product( (torch.Size([]), torch.Size([3])), (1, 2), (torch.float, torch.double)): state_dict = deepcopy(raw_state_dict) for k, v in state_dict.items(): if m == 1: v = v[0] if len(train_batch_shape) > 0: v = v.unsqueeze(0).expand(*train_batch_shape, *v.shape) state_dict[k] = v tkwargs = {"device": self.device, "dtype": dtype} if m == 2: objective = GenericMCObjective(lambda Y, X: Y.sum(dim=-1)) else: objective = None for k, v in state_dict.items(): state_dict[k] = v.to(**tkwargs) all_close_kwargs = ({ "atol": 1e-1, "rtol": 0.0, } if dtype == torch.float else { "atol": 1e-4, "rtol": 0.0 }) torch.manual_seed(1234) train_X = torch.rand(*train_batch_shape, 3, 2, **tkwargs) train_Y = ( torch.sin(train_X * 2 * pi) + torch.randn(*train_batch_shape, 3, 2, **tkwargs))[..., :m] train_Y = standardize(train_Y) model = SingleTaskGP( train_X, train_Y, ) if len(train_batch_shape) > 0: X_baseline = train_X[0] else: X_baseline = train_X model.load_state_dict(state_dict, strict=False) # test sampler with collapse_batch_dims=False sampler = IIDNormalSampler(5, seed=0, collapse_batch_dims=False) with self.assertRaises(UnsupportedError): qNoisyExpectedImprovement( model=model, X_baseline=X_baseline, sampler=sampler, objective=objective, prune_baseline=False, cache_root=True, ) sampler = IIDNormalSampler(5, seed=0) torch.manual_seed(0) acqf = qNoisyExpectedImprovement( model=model, X_baseline=X_baseline, sampler=sampler, objective=objective, prune_baseline=False, cache_root=True, ) orig_base_samples = acqf.base_sampler.base_samples.detach().clone() sampler2 = IIDNormalSampler(5, seed=0) sampler2.base_samples = orig_base_samples torch.manual_seed(0) acqf_no_cache = qNoisyExpectedImprovement( model=model, X_baseline=X_baseline, sampler=sampler2, objective=objective, prune_baseline=False, cache_root=False, ) for q, batch_shape in product( (1, 3), (torch.Size([]), torch.Size([3]), torch.Size([4, 3]))): test_X = (0.3 + 0.05 * torch.randn(*batch_shape, q, 2, **tkwargs) ).requires_grad_(True) with mock.patch( sample_cached_path, wraps=sample_cached_cholesky) as mock_sample_cached: torch.manual_seed(0) val = acqf(test_X) mock_sample_cached.assert_called_once() val.sum().backward() base_samples = acqf.sampler.base_samples.detach().clone() X_grad = test_X.grad.clone() test_X2 = test_X.detach().clone().requires_grad_(True) acqf_no_cache.sampler.base_samples = base_samples with mock.patch( sample_cached_path, wraps=sample_cached_cholesky) as mock_sample_cached: torch.manual_seed(0) val2 = acqf_no_cache(test_X2) mock_sample_cached.assert_not_called() self.assertTrue(torch.allclose(val, val2, **all_close_kwargs)) val2.sum().backward() self.assertTrue( torch.allclose(X_grad, test_X2.grad, **all_close_kwargs)) # test we fall back to standard sampling for # ill-conditioned covariances acqf._baseline_L = torch.zeros_like(acqf._baseline_L) with warnings.catch_warnings( record=True) as ws, settings.debug(True): with torch.no_grad(): acqf(test_X) self.assertEqual(len(ws), 1) self.assertTrue(issubclass(ws[-1].category, BotorchWarning))