def test_evaluate_kg(self): # a thorough test using real model and dtype double d = 2 dtype = torch.double bounds = torch.tensor([[0], [1]], device=self.device, dtype=dtype).repeat(1, d) train_X = torch.rand(3, d, device=self.device, dtype=dtype) train_Y = torch.rand(3, 1, device=self.device, dtype=dtype) model = SingleTaskGP(train_X, train_Y) qKG = qKnowledgeGradient( model=model, num_fantasies=2, objective=None, X_pending=torch.rand(2, d, device=self.device, dtype=dtype), current_value=torch.rand(1, device=self.device, dtype=dtype), ) X = torch.rand(4, 3, d, device=self.device, dtype=dtype) options = {"num_inner_restarts": 2, "raw_inner_samples": 3} val = qKG.evaluate(X, bounds=bounds, num_restarts=2, raw_samples=3, options=options) # verify output shape self.assertEqual(val.size(), torch.Size([4])) # verify dtype self.assertEqual(val.dtype, dtype) # test i) no dimension is squeezed out, ii) dtype float, iii) MC objective, # and iv) t_batch_mode_transform dtype = torch.float bounds = torch.tensor([[0], [1]], device=self.device, dtype=dtype) train_X = torch.rand(1, 1, device=self.device, dtype=dtype) train_Y = torch.rand(1, 1, device=self.device, dtype=dtype) model = SingleTaskGP(train_X, train_Y) qKG = qKnowledgeGradient( model=model, num_fantasies=1, objective=GenericMCObjective( objective=lambda Y, X: Y.norm(dim=-1)), ) X = torch.rand(1, 1, device=self.device, dtype=dtype) options = {"num_inner_restarts": 1, "raw_inner_samples": 1} val = qKG.evaluate(X, bounds=bounds, num_restarts=1, raw_samples=1, options=options) # verify output shape self.assertEqual(val.size(), torch.Size([1])) # verify dtype self.assertEqual(val.dtype, dtype)
def test_gen_one_shot_kg_initial_conditions(self): num_fantasies = 8 num_restarts = 4 raw_samples = 16 for dtype in (torch.float, torch.double): mean = torch.zeros(1, 1, device=self.device, dtype=dtype) mm = MockModel(MockPosterior(mean=mean)) mock_kg = qKnowledgeGradient(model=mm, num_fantasies=num_fantasies) bounds = torch.tensor([[0, 0], [1, 1]], device=self.device, dtype=dtype) # test option error with self.assertRaises(ValueError): gen_one_shot_kg_initial_conditions( acq_function=mock_kg, bounds=bounds, q=1, num_restarts=num_restarts, raw_samples=raw_samples, options={"frac_random": 2.0}, ) # test generation logic q = 2 mock_random_ics = torch.rand(num_restarts, q + num_fantasies, 2) mock_fantasy_cands = torch.ones(20, 1, 2) mock_fantasy_vals = torch.randn(20) with ExitStack() as es: mock_gbics = es.enter_context( mock.patch( "botorch.optim.initializers.gen_batch_initial_conditions", return_value=mock_random_ics, ) ) mock_optacqf = es.enter_context( mock.patch( "botorch.optim.optimize.optimize_acqf", return_value=(mock_fantasy_cands, mock_fantasy_vals), ) ) ics = gen_one_shot_kg_initial_conditions( acq_function=mock_kg, bounds=bounds, q=q, num_restarts=num_restarts, raw_samples=raw_samples, ) mock_gbics.assert_called_once() mock_optacqf.assert_called_once() n_value = int((1 - 0.1) * num_fantasies) self.assertTrue( torch.equal( ics[..., :-n_value, :], mock_random_ics[..., :-n_value, :] ) ) self.assertTrue(torch.all(ics[..., -n_value:, :] == 1))
def _instantiate_KG( model: Model, objective: AcquisitionObjective, qmc: bool = True, n_fantasies: int = 64, mc_samples: int = 256, num_trace_observations: int = 0, seed_inner: Optional[int] = None, seed_outer: Optional[int] = None, X_pending: Optional[Tensor] = None, current_value: Optional[Tensor] = None, target_fidelities: Optional[Dict[int, float]] = None, fidelity_weights: Optional[Dict[int, float]] = None, cost_intercept: float = 1.0, ) -> qKnowledgeGradient: r"""Instantiate either a `qKnowledgeGradient` or `qMultiFidelityKnowledgeGradient` acquisition function depending on whether `target_fidelities` is defined. """ sampler_cls = SobolQMCNormalSampler if qmc else IIDNormalSampler fantasy_sampler = sampler_cls(num_samples=n_fantasies, seed=seed_outer) if isinstance(objective, MCAcquisitionObjective): inner_sampler = sampler_cls(num_samples=mc_samples, seed=seed_inner) else: inner_sampler = None if target_fidelities: if fidelity_weights is None: fidelity_weights = {f: 1.0 for f in target_fidelities} if not set(target_fidelities) == set(fidelity_weights): raise RuntimeError( "Must provide the same indices for target_fidelities " f"({set(target_fidelities)}) and fidelity_weights " f" ({set(fidelity_weights)}).") cost_model = AffineFidelityCostModel(fidelity_weights=fidelity_weights, fixed_cost=cost_intercept) cost_aware_utility = InverseCostWeightedUtility(cost_model=cost_model) def project(X: Tensor) -> Tensor: return project_to_target_fidelity( X=X, target_fidelities=target_fidelities) def expand(X: Tensor) -> Tensor: return expand_trace_observations( X=X, fidelity_dims=sorted(target_fidelities), # pyre-ignore: [6] num_trace_obs=num_trace_observations, ) return qMultiFidelityKnowledgeGradient( model=model, num_fantasies=n_fantasies, sampler=fantasy_sampler, objective=objective, inner_sampler=inner_sampler, X_pending=X_pending, current_value=current_value, cost_aware_utility=cost_aware_utility, project=project, expand=expand, ) return qKnowledgeGradient( model=model, num_fantasies=n_fantasies, sampler=fantasy_sampler, objective=objective, inner_sampler=inner_sampler, X_pending=X_pending, current_value=current_value, )
def test_initialize_q_knowledge_gradient(self): for dtype in (torch.float, torch.double): mean = torch.zeros(1, 1, device=self.device, dtype=dtype) mm = MockModel(MockPosterior(mean=mean)) # test error when neither specifying neither sampler nor num_fantasies with self.assertRaises(ValueError): qKnowledgeGradient(model=mm, num_fantasies=None) # test error when sampler and num_fantasies arg are inconsistent sampler = IIDNormalSampler(num_samples=16) with self.assertRaises(ValueError): qKnowledgeGradient(model=mm, num_fantasies=32, sampler=sampler) # test default construction qKG = qKnowledgeGradient(model=mm, num_fantasies=32) self.assertEqual(qKG.num_fantasies, 32) self.assertIsInstance(qKG.sampler, SobolQMCNormalSampler) self.assertEqual(qKG.sampler.sample_shape, torch.Size([32])) self.assertIsNone(qKG.objective) self.assertIsNone(qKG.inner_sampler) self.assertIsNone(qKG.X_pending) self.assertIsNone(qKG.current_value) self.assertEqual(qKG.get_augmented_q_batch_size(q=3), 32 + 3) # test custom construction obj = GenericMCObjective(lambda Y, X: Y.mean(dim=-1)) sampler = IIDNormalSampler(num_samples=16) X_pending = torch.zeros(2, 2, device=self.device, dtype=dtype) qKG = qKnowledgeGradient( model=mm, num_fantasies=16, sampler=sampler, objective=obj, X_pending=X_pending, ) self.assertEqual(qKG.num_fantasies, 16) self.assertEqual(qKG.sampler, sampler) self.assertEqual(qKG.sampler.sample_shape, torch.Size([16])) self.assertEqual(qKG.objective, obj) self.assertIsInstance(qKG.inner_sampler, SobolQMCNormalSampler) self.assertEqual(qKG.inner_sampler.sample_shape, torch.Size([128])) self.assertTrue(torch.equal(qKG.X_pending, X_pending)) self.assertIsNone(qKG.current_value) self.assertEqual(qKG.get_augmented_q_batch_size(q=3), 16 + 3) # test assignment of num_fantasies from sampler if not provided qKG = qKnowledgeGradient(model=mm, num_fantasies=None, sampler=sampler) self.assertEqual(qKG.sampler.sample_shape, torch.Size([16])) # test custom construction with inner sampler and current value inner_sampler = SobolQMCNormalSampler(num_samples=256) current_value = torch.zeros(1, device=self.device, dtype=dtype) qKG = qKnowledgeGradient( model=mm, num_fantasies=8, objective=obj, inner_sampler=inner_sampler, current_value=current_value, ) self.assertEqual(qKG.num_fantasies, 8) self.assertEqual(qKG.sampler.sample_shape, torch.Size([8])) self.assertEqual(qKG.objective, obj) self.assertIsInstance(qKG.inner_sampler, SobolQMCNormalSampler) self.assertEqual(qKG.inner_sampler, inner_sampler) self.assertIsNone(qKG.X_pending) self.assertTrue(torch.equal(qKG.current_value, current_value)) self.assertEqual(qKG.get_augmented_q_batch_size(q=3), 8 + 3) # test construction with non-MC objective (ScalarizedObjective) qKG_s = qKnowledgeGradient( model=mm, num_fantasies=16, sampler=sampler, objective=ScalarizedObjective(weights=torch.rand(2)), ) self.assertIsNone(qKG_s.inner_sampler) self.assertIsInstance(qKG_s.objective, ScalarizedObjective) # test error if no objective and multi-output model mean2 = torch.zeros(1, 2, device=self.device, dtype=dtype) mm2 = MockModel(MockPosterior(mean=mean2)) with self.assertRaises(UnsupportedError): qKnowledgeGradient(model=mm2)
def test_evaluate_q_knowledge_gradient(self): for dtype in (torch.float, torch.double): # basic test n_f = 4 mean = torch.rand(n_f, 1, 1, device=self.device, dtype=dtype) variance = torch.rand(n_f, 1, 1, device=self.device, dtype=dtype) mfm = MockModel(MockPosterior(mean=mean, variance=variance)) with mock.patch.object(MockModel, "fantasize", return_value=mfm) as patch_f: with mock.patch(NO, new_callable=mock.PropertyMock) as mock_num_outputs: mock_num_outputs.return_value = 1 mm = MockModel(None) qKG = qKnowledgeGradient(model=mm, num_fantasies=n_f) X = torch.rand(n_f + 1, 1, device=self.device, dtype=dtype) val = qKG(X) patch_f.assert_called_once() cargs, ckwargs = patch_f.call_args self.assertEqual(ckwargs["X"].shape, torch.Size([1, 1, 1])) self.assertTrue(torch.allclose(val, mean.mean(), atol=1e-4)) self.assertTrue(torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :])) # batched evaluation b = 2 mean = torch.rand(n_f, b, 1, device=self.device, dtype=dtype) variance = torch.rand(n_f, b, 1, device=self.device, dtype=dtype) mfm = MockModel(MockPosterior(mean=mean, variance=variance)) X = torch.rand(b, n_f + 1, 1, device=self.device, dtype=dtype) with mock.patch.object(MockModel, "fantasize", return_value=mfm) as patch_f: with mock.patch(NO, new_callable=mock.PropertyMock) as mock_num_outputs: mock_num_outputs.return_value = 1 mm = MockModel(None) qKG = qKnowledgeGradient(model=mm, num_fantasies=n_f) val = qKG(X) patch_f.assert_called_once() cargs, ckwargs = patch_f.call_args self.assertEqual(ckwargs["X"].shape, torch.Size([b, 1, 1])) self.assertTrue( torch.allclose(val, mean.mean(dim=0).squeeze(-1), atol=1e-4) ) self.assertTrue(torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :])) # pending points and current value X_pending = torch.rand(2, 1, device=self.device, dtype=dtype) mean = torch.rand(n_f, 1, 1, device=self.device, dtype=dtype) variance = torch.rand(n_f, 1, 1, device=self.device, dtype=dtype) mfm = MockModel(MockPosterior(mean=mean, variance=variance)) current_value = torch.rand(1, device=self.device, dtype=dtype) X = torch.rand(n_f + 1, 1, device=self.device, dtype=dtype) with mock.patch.object(MockModel, "fantasize", return_value=mfm) as patch_f: with mock.patch(NO, new_callable=mock.PropertyMock) as mock_num_outputs: mock_num_outputs.return_value = 1 mm = MockModel(None) qKG = qKnowledgeGradient( model=mm, num_fantasies=n_f, X_pending=X_pending, current_value=current_value, ) val = qKG(X) patch_f.assert_called_once() cargs, ckwargs = patch_f.call_args self.assertEqual(ckwargs["X"].shape, torch.Size([1, 3, 1])) self.assertTrue(torch.allclose(val, mean.mean() - current_value, atol=1e-4)) self.assertTrue(torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :])) # test objective (inner MC sampling) objective = GenericMCObjective(objective=lambda Y, X: Y.norm(dim=-1)) samples = torch.randn(3, 1, 1, device=self.device, dtype=dtype) mfm = MockModel(MockPosterior(samples=samples)) X = torch.rand(n_f + 1, 1, device=self.device, dtype=dtype) with mock.patch.object(MockModel, "fantasize", return_value=mfm) as patch_f: with mock.patch(NO, new_callable=mock.PropertyMock) as mock_num_outputs: mock_num_outputs.return_value = 1 mm = MockModel(None) qKG = qKnowledgeGradient( model=mm, num_fantasies=n_f, objective=objective ) val = qKG(X) patch_f.assert_called_once() cargs, ckwargs = patch_f.call_args self.assertEqual(ckwargs["X"].shape, torch.Size([1, 1, 1])) self.assertTrue(torch.allclose(val, objective(samples).mean(), atol=1e-4)) self.assertTrue(torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :])) # test non-MC objective (ScalarizedObjective) weights = torch.rand(2, device=self.device, dtype=dtype) objective = ScalarizedObjective(weights=weights) mean = torch.tensor([1.0, 0.5], device=self.device, dtype=dtype).expand( n_f, 1, 2 ) cov = torch.tensor( [[1.0, 0.1], [0.1, 0.5]], device=self.device, dtype=dtype ).expand(n_f, 2, 2) posterior = GPyTorchPosterior(MultitaskMultivariateNormal(mean, cov)) mfm = MockModel(posterior) with mock.patch.object(MockModel, "fantasize", return_value=mfm) as patch_f: with mock.patch(NO, new_callable=mock.PropertyMock) as mock_num_outputs: mock_num_outputs.return_value = 2 mm = MockModel(None) qKG = qKnowledgeGradient( model=mm, num_fantasies=n_f, objective=objective ) val = qKG(X) patch_f.assert_called_once() cargs, ckwargs = patch_f.call_args self.assertEqual(ckwargs["X"].shape, torch.Size([1, 1, 1])) val_expected = (mean * weights).sum(-1).mean(0) self.assertTrue(torch.allclose(val, val_expected))
def main( benchmark_name, dataset_name, dimensions, method_name, num_runs, run_start, num_iterations, acquisition_name, # acquisition_optimizer_name, gamma, num_random_init, mc_samples, batch_size, num_fantasies, num_restarts, raw_samples, noise_variance_init, # use_ard, # use_input_warping, standardize_targets, input_dir, output_dir): # TODO(LT): Turn into options # device = "cpu" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = torch.double benchmark = make_benchmark(benchmark_name, dimensions=dimensions, dataset_name=dataset_name, input_dir=input_dir) name = make_name(benchmark_name, dimensions=dimensions, dataset_name=dataset_name) output_path = Path(output_dir).joinpath(name, method_name) output_path.mkdir(parents=True, exist_ok=True) options = dict(gamma=gamma, num_random_init=num_random_init, acquisition_name=acquisition_name, mc_samples=mc_samples, batch_size=batch_size, num_restarts=num_restarts, raw_samples=raw_samples, num_fantasies=num_fantasies, noise_variance_init=noise_variance_init, standardize_targets=standardize_targets) with output_path.joinpath("options.yaml").open('w') as f: yaml.dump(options, f) config_space = DenseConfigurationSpace(benchmark.get_config_space()) bounds = create_bounds(config_space.get_bounds(), device=device, dtype=dtype) input_dim = config_space.get_dimensions() def func(tensor, *args, **kwargs): """ Wrapper that receives and returns torch.Tensor """ config = dict_from_tensor(tensor, cs=config_space) # turn into maximization problem res = -benchmark.evaluate(config).value return torch.tensor(res, device=device, dtype=dtype) for run_id in trange(run_start, num_runs, unit="run"): run_begin_t = batch_end_t_adj = batch_end_t = datetime.now() frames = [] features = [] targets = [] noise_variance = torch.tensor(noise_variance_init, device=device, dtype=dtype) state_dict = None with trange(num_iterations) as iterations: for batch in iterations: if len(targets) < num_random_init: # click.echo(f"Completed {i}/{num_random_init} initial runs. " # "Suggesting random candidate...") # TODO(LT): support random seed X_batch = torch.rand(size=(batch_size, input_dim), device=device, dtype=dtype) else: # construct dataset X = torch.vstack(features) y = torch.hstack(targets).unsqueeze(axis=-1) y = standardize(y) if standardize_targets else y # construct model # model = FixedNoiseGP(X, standardize(y), noise_variance.expand_as(y), model = FixedNoiseGP(X, y, noise_variance.expand_as(y), input_transform=None).to(X) mll = ExactMarginalLogLikelihood(model.likelihood, model) if state_dict is not None: model.load_state_dict(state_dict) # update model fit_gpytorch_model(mll) # construct acquisition function tau = torch.quantile(y, q=1 - gamma) iterations.set_postfix(tau=tau.item()) if acquisition_name == "q-KG": assert num_fantasies is not None and num_fantasies > 0 acq = qKnowledgeGradient(model, num_fantasies=num_fantasies) elif acquisition_name == "q-EI": assert mc_samples is not None and mc_samples > 0 qmc_sampler = SobolQMCNormalSampler( num_samples=mc_samples) acq = qExpectedImprovement(model=model, best_f=tau, sampler=qmc_sampler) # optimize acquisition function X_batch, b = optimize_acqf(acq_function=acq, bounds=bounds, q=batch_size, num_restarts=num_restarts, raw_samples=raw_samples, options=dict(batch_limit=5, maxiter=200)) state_dict = model.state_dict() # begin batch evaluation batch_begin_t = datetime.now() decision_duration = batch_begin_t - batch_end_t batch_begin_t_adj = batch_end_t_adj + decision_duration eval_end_times = [] # TODO(LT): Deliberately not doing broadcasting for now since # batch sizes are so small anyway. Can revisit later if there # is a compelling reason to do it. rows = [] for j, x_next in enumerate(X_batch): # eval begin time eval_begin_t = datetime.now() # evaluate blackbox objective y_next = func(x_next) # eval end time eval_end_t = datetime.now() # eval duration eval_duration = eval_end_t - eval_begin_t # adjusted eval end time is the duration added to the # time at which batch eval was started eval_end_t_adj = batch_begin_t_adj + eval_duration eval_end_times.append(eval_end_t_adj) elapsed = eval_end_t_adj - run_begin_t # update dataset features.append(x_next) targets.append(y_next) row = dict_from_tensor(x_next, cs=config_space) row["loss"] = -y_next.item() row["cost_eval"] = eval_duration.total_seconds() row["finished"] = elapsed.total_seconds() rows.append(row) batch_end_t = datetime.now() batch_end_t_adj = max(eval_end_times) frame = pd.DataFrame(data=rows) \ .assign(batch=batch, cost_decision=decision_duration.total_seconds()) frames.append(frame) data = pd.concat(frames, axis="index", ignore_index=True) data.to_csv(output_path.joinpath(f"{run_id:03d}.csv")) return 0
def test_evaluate_q_knowledge_gradient(self): for dtype in (torch.float, torch.double): # basic test n_f = 4 mean = torch.rand(n_f, 1, device=self.device, dtype=dtype) variance = torch.rand(n_f, 1, device=self.device, dtype=dtype) mfm = MockModel(MockPosterior(mean=mean, variance=variance)) with mock.patch.object(MockModel, "fantasize", return_value=mfm) as patch_f: mm = MockModel(None) qKG = qKnowledgeGradient(model=mm, num_fantasies=n_f) X = torch.rand(n_f + 1, 1, device=self.device, dtype=dtype) val = qKG(X) patch_f.assert_called_once() cargs, ckwargs = patch_f.call_args self.assertEqual(ckwargs["X"].shape, torch.Size([1, 1])) self.assertTrue(torch.allclose(val, mean.mean(), atol=1e-4)) self.assertTrue( torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :])) # batched evaluation b = 2 mean = torch.rand(n_f, b, 1, device=self.device, dtype=dtype) variance = torch.rand(n_f, b, 1, device=self.device, dtype=dtype) mfm = MockModel(MockPosterior(mean=mean, variance=variance)) X = torch.rand(b, n_f + 1, 1, device=self.device, dtype=dtype) with mock.patch.object(MockModel, "fantasize", return_value=mfm) as patch_f: mm = MockModel(None) qKG = qKnowledgeGradient(model=mm, num_fantasies=n_f) val = qKG(X) patch_f.assert_called_once() cargs, ckwargs = patch_f.call_args self.assertEqual(ckwargs["X"].shape, torch.Size([b, 1, 1])) self.assertTrue( torch.allclose(val, mean.mean(dim=0).squeeze(-1), atol=1e-4)) self.assertTrue( torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :])) # pending points and current value mean = torch.rand(n_f, 1, device=self.device, dtype=dtype) variance = torch.rand(n_f, 1, device=self.device, dtype=dtype) X_pending = torch.rand(2, 1, device=self.device, dtype=dtype) mfm = MockModel(MockPosterior(mean=mean, variance=variance)) current_value = torch.rand(1, device=self.device, dtype=dtype) X = torch.rand(n_f + 1, 1, device=self.device, dtype=dtype) with mock.patch.object(MockModel, "fantasize", return_value=mfm) as patch_f: mm = MockModel(None) qKG = qKnowledgeGradient( model=mm, num_fantasies=n_f, X_pending=X_pending, current_value=current_value, ) val = qKG(X) patch_f.assert_called_once() cargs, ckwargs = patch_f.call_args self.assertEqual(ckwargs["X"].shape, torch.Size([3, 1])) self.assertTrue( torch.allclose(val, mean.mean() - current_value, atol=1e-4)) self.assertTrue( torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :])) # test objective (inner MC sampling) objective = GenericMCObjective(objective=lambda Y: Y.norm(dim=-1)) samples = torch.randn(3, 1, 1, device=self.device, dtype=dtype) mfm = MockModel(MockPosterior(samples=samples)) X = torch.rand(n_f + 1, 1, device=self.device, dtype=dtype) with mock.patch.object(MockModel, "fantasize", return_value=mfm) as patch_f: mm = MockModel(None) qKG = qKnowledgeGradient(model=mm, num_fantasies=n_f, objective=objective) val = qKG(X) patch_f.assert_called_once() cargs, ckwargs = patch_f.call_args self.assertEqual(ckwargs["X"].shape, torch.Size([1, 1])) self.assertTrue( torch.allclose(val, objective(samples).mean(), atol=1e-4)) self.assertTrue( torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :]))