def test_qEI(self):
for double in (True, False):
self._setUp(double=double)
qEI = qExpectedImprovement(self.model_st, best_f=0.0)
candidates, _ = optimize_acqf(
acq_function=qEI,
bounds=self.bounds,
q=3,
num_restarts=10,
raw_samples=20,
options={"maxiter": 5},
)
self.assertTrue(torch.all(-EPS <= candidates))
self.assertTrue(torch.all(candidates <= 1 + EPS))
qEI = qExpectedImprovement(self.model_fn, best_f=0.0)
candidates, _ = optimize_acqf(
acq_function=qEI,
bounds=self.bounds,
q=3,
num_restarts=10,
raw_samples=20,
options={"maxiter": 5},
)
self.assertTrue(torch.all(-EPS <= candidates))
self.assertTrue(torch.all(candidates <= 1 + EPS))
candidates_batch_limit, _ = optimize_acqf(
acq_function=qEI,
bounds=self.bounds,
q=3,
num_restarts=10,
raw_samples=20,
options={"maxiter": 5, "batch_limit": 5},
)
self.assertTrue(torch.all(-EPS <= candidates_batch_limit))
self.assertTrue(torch.all(candidates_batch_limit <= 1 + EPS))
def test_qEI(self, cuda=False):
for double in (True, False):
self._setUp(double=double, cuda=cuda)
qEI = qExpectedImprovement(self.model_st, best_f=0.0)
candidates = joint_optimize(
acq_function=qEI,
bounds=self.bounds,
q=3,
num_restarts=10,
raw_samples=20,
options={"maxiter": 5},
)
self.assertTrue(torch.all(-EPS <= candidates))
self.assertTrue(torch.all(candidates <= 1 + EPS))
qEI = qExpectedImprovement(self.model_fn, best_f=0.0)
candidates = joint_optimize(
acq_function=qEI,
bounds=self.bounds,
q=3,
num_restarts=10,
raw_samples=20,
options={"maxiter": 5},
)
self.assertTrue(torch.all(-EPS <= candidates))
self.assertTrue(torch.all(candidates <= 1 + EPS))
def test_qEI(self, cuda=False):
for double in (True, False):
self._setUp(double=double, cuda=cuda)
qEI = qExpectedImprovement(self.model_st, best_f=0.0)
candidates = joint_optimize(
acq_function=qEI,
bounds=self.bounds,
q=3,
num_restarts=10,
raw_samples=20,
options={"maxiter": 5},
)
self.assertTrue(torch.all(-EPS <= candidates))
self.assertTrue(torch.all(candidates <= 1 + EPS))
qEI = qExpectedImprovement(self.model_fn, best_f=0.0)
candidates = joint_optimize(
acq_function=qEI,
bounds=self.bounds,
q=3,
num_restarts=10,
raw_samples=20,
options={"maxiter": 5},
)
self.assertTrue(torch.all(-EPS <= candidates))
self.assertTrue(torch.all(candidates <= 1 + EPS))
def test_gen_candidates(self, gen_candidates=gen_candidates_scipy, options=None):
options = options or {}
options = {**options, "maxiter": 5}
for double in (True, False):
self._setUp(double=double)
acqfs = [
qExpectedImprovement(self.model, best_f=self.f_best),
qKnowledgeGradient(
self.model, num_fantasies=4, current_value=self.f_best
),
]
for acqf in acqfs:
ics = self.initial_conditions
if isinstance(acqf, qKnowledgeGradient):
ics = ics.repeat(5, 1)
candidates, _ = gen_candidates(
initial_conditions=ics,
acquisition_function=acqf,
lower_bounds=0,
upper_bounds=1,
options=options or {},
)
if isinstance(acqf, qKnowledgeGradient):
candidates = acqf.extract_candidates(candidates)
self.assertTrue(-EPS <= candidates <= 1 + EPS)
def get_candidate(model, acq, full_train_Y, q, bounds, dim):
if acq == 'EI':
if q == 1:
EI = ExpectedImprovement(model, full_train_Y.max().item())
else:
EI = qExpectedImprovement(model, full_train_Y.max().item())
bounds_t = torch.FloatTensor([[bounds[0]] * dim, [bounds[1]] * dim])
candidate, acq_value = optimize_acqf(
EI,
bounds=bounds_t,
q=q,
num_restarts=15,
raw_samples=5000,
)
elif acq == 'TS':
sobol = SobolEngine(dim, scramble=True)
n_candidates = min(5000, max(20000, 2000 * dim))
pert = sobol.draw(n_candidates)
X_cand = (bounds[1] - bounds[0]) * pert + bounds[0]
thompson_sampling = MaxPosteriorSampling(model=model,
replacement=False)
candidate = thompson_sampling(X_cand, num_samples=q)
else:
raise NotImplementedError('Only TS and EI are implemented')
return candidate, EI if acq == 'EI' else None
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_gen_candidates(self, cuda=False, gen_candidates=gen_candidates_scipy):
for double in (True, False):
self._setUp(double=double, cuda=cuda)
qEI = qExpectedImprovement(self.model, best_f=self.f_best)
candidates, _ = gen_candidates(
initial_conditions=self.initial_conditions,
acquisition_function=qEI,
lower_bounds=0,
upper_bounds=1,
)
self.assertTrue(-EPS <= candidates <= 1 + EPS)
def optimize_loop(model, loss, X_train, y_train, bounds, n_samples=10):
best_value = y_train.max()
acq_func = qExpectedImprovement(model, best_f=best_value)
X_new, acq_value = optimize_acqf(acq_func,
bounds=bounds,
q=20,
num_restarts=10,
raw_samples=76)
#X_new = X_new.view((n_samples,-1))
print(X_new)
return X_new
def test_gen_candidates(self, cuda=False, gen_candidates=gen_candidates_scipy):
for double in (True, False):
self._setUp(double=double, cuda=cuda)
qEI = qExpectedImprovement(self.model, best_f=self.f_best)
candidates, _ = gen_candidates(
initial_conditions=self.initial_conditions,
acquisition_function=qEI,
lower_bounds=0,
upper_bounds=1,
)
self.assertTrue(-EPS <= candidates <= 1 + EPS)
def test_gen_candidates_with_none_fixed_features(
self, cuda=False, gen_candidates=gen_candidates_scipy):
for double in (True, False):
self._setUp(double=double, cuda=cuda, expand=True)
qEI = qExpectedImprovement(self.model, best_f=self.f_best)
candidates, _ = gen_candidates(
initial_conditions=self.initial_conditions,
acquisition_function=qEI,
lower_bounds=0,
upper_bounds=1,
fixed_features={1: None},
)
candidates = candidates.squeeze(0)
self.assertTrue(-EPS <= candidates[0] <= 1 + EPS)
self.assertTrue(candidates[1].item() == 1.0)
def test_gen_candidates(self,
gen_candidates=gen_candidates_scipy,
options=None):
options = options or {}
options = {**options, "maxiter": 5}
for double in (True, False):
self._setUp(double=double)
qEI = qExpectedImprovement(self.model, best_f=self.f_best)
candidates, _ = gen_candidates(
initial_conditions=self.initial_conditions,
acquisition_function=qEI,
lower_bounds=0,
upper_bounds=1,
options=options or {},
)
self.assertTrue(-EPS <= candidates <= 1 + EPS)
def test_gen_candidates_with_fixed_features(
self, cuda=False, gen_candidates=gen_candidates_scipy
):
for double in (True, False):
self._setUp(double=double, cuda=cuda, expand=True)
qEI = qExpectedImprovement(self.model, best_f=self.f_best)
candidates, _ = gen_candidates(
initial_conditions=self.initial_conditions,
acquisition_function=qEI,
lower_bounds=0,
upper_bounds=1,
fixed_features={1: 0.25},
)
candidates = candidates.squeeze(0)
self.assertTrue(-EPS <= candidates[0] <= 1 + EPS)
self.assertTrue(candidates[1].item() == 0.25)
def test_random_restart_optimization(self):
for double in (True, False):
self._setUp(double=double)
with gpt_settings.debug(False):
best_f = self.model(self.train_x).mean.max().item()
qEI = qExpectedImprovement(self.model, best_f=best_f)
bounds = torch.tensor([[0.0], [1.0]]).type_as(self.train_x)
batch_ics = torch.rand(2, 1).type_as(self.train_x)
batch_candidates, batch_acq_values = gen_candidates_scipy(
initial_conditions=batch_ics,
acquisition_function=qEI,
lower_bounds=bounds[0],
upper_bounds=bounds[1],
options={"maxiter": 3},
)
candidates = get_best_candidates(batch_candidates=batch_candidates,
batch_values=batch_acq_values)
self.assertTrue(-EPS <= candidates <= 1 + EPS)
def test_gen_candidates_with_fixed_features(
self, gen_candidates=gen_candidates_scipy, options=None):
options = options or {}
options = {**options, "maxiter": 5}
for double in (True, False):
self._setUp(double=double, expand=True)
qEI = qExpectedImprovement(self.model, best_f=self.f_best)
candidates, _ = gen_candidates(
initial_conditions=self.initial_conditions,
acquisition_function=qEI,
lower_bounds=0,
upper_bounds=1,
fixed_features={1: 0.25},
options=options,
)
candidates = candidates.squeeze(0)
self.assertTrue(-EPS <= candidates[0] <= 1 + EPS)
self.assertTrue(candidates[1].item() == 0.25)
def test_random_restart_optimization(self, cuda=False):
for double in (True, False):
self._setUp(double=double, cuda=cuda)
with settings.debug(False):
best_f = self.model(self.train_x).mean.max().item()
qEI = qExpectedImprovement(self.model, best_f=best_f)
bounds = torch.tensor([[0.0], [1.0]]).type_as(self.train_x)
batch_ics = torch.rand(2, 1).type_as(self.train_x)
batch_candidates, batch_acq_values = gen_candidates_scipy(
initial_conditions=batch_ics,
acquisition_function=qEI,
lower_bounds=bounds[0],
upper_bounds=bounds[1],
options={"maxiter": 3},
)
candidates = get_best_candidates(
batch_candidates=batch_candidates, batch_values=batch_acq_values
)
self.assertTrue(-EPS <= candidates <= 1 + EPS)
def test_gen_candidates_scipy_with_fixed_features_inequality_constraints(self):
options = {"maxiter": 5}
for double in (True, False):
self._setUp(double=double, expand=True)
qEI = qExpectedImprovement(self.model, best_f=self.f_best)
candidates, _ = gen_candidates_scipy(
initial_conditions=self.initial_conditions.reshape(1, 1, -1),
acquisition_function=qEI,
inequality_constraints=[
(torch.tensor([0]), torch.tensor([1]), 0),
(torch.tensor([1]), torch.tensor([-1]), -1),
],
fixed_features={1: 0.25},
options=options,
)
# candidates is of dimension 1 x 1 x 2
# so we are squeezing all the singleton dimensions
candidates = candidates.squeeze()
self.assertTrue(-EPS <= candidates[0] <= 1 + EPS)
self.assertTrue(candidates[1].item() == 0.25)
def bayes_opt(x0, y0):
"""
Main Bayesian optimization loop. Begins by initializing model, then for each
iteration, it fits the GP to the data, gets a new point with the acquisition
function, adds it to the dataset, and exits if it's a successful attack
"""
best_observed = []
query_count, success = 0, 0
# call helper function to initialize model
train_x, train_obj, mll, model, best_value, mean, std = initialize_model(
x0, y0, n=args.initial_samples)
if args.standardize_every_iter:
train_obj = (train_obj - train_obj.mean()) / train_obj.std()
best_observed.append(best_value)
query_count += args.initial_samples
# run args.iter rounds of BayesOpt after the initial random batch
for _ in range(args.iter):
# fit the model
fit_gpytorch_model(mll)
# define the qNEI acquisition module using a QMC sampler
if args.q != 1:
qmc_sampler = SobolQMCNormalSampler(num_samples=2000, seed=seed)
qEI = qExpectedImprovement(model=model,
sampler=qmc_sampler,
best_f=best_value)
else:
if args.acqf == 'EI':
qEI = ExpectedImprovement(model=model, best_f=best_value)
elif args.acqf == 'PM':
qEI = PosteriorMean(model)
elif args.acqf == 'POI':
qEI = ProbabilityOfImprovement(model, best_f=best_value)
elif args.acqf == 'UCB':
qEI = UpperConfidenceBound(model, beta=args.beta)
# optimize and get new observation
new_x, new_obj = optimize_acqf_and_get_observation(qEI, x0, y0)
if args.standardize:
new_obj = (new_obj - mean) / std
# update training points
train_x = torch.cat((train_x, new_x))
train_obj = torch.cat((train_obj, new_obj))
if args.standardize_every_iter:
train_obj = (train_obj - train_obj.mean()) / train_obj.std()
# update progress
best_value, best_index = train_obj.max(0)
best_observed.append(best_value.item())
best_candidate = train_x[best_index]
# reinitialize the model so it is ready for fitting on next iteration
torch.cuda.empty_cache()
model.set_train_data(train_x, train_obj, strict=False)
# get objective value of best candidate; if we found an adversary, exit
best_candidate = best_candidate.view(1, -1)
best_candidate = transform(best_candidate, args.dset, args.arch,
args.cos, args.sin).to(device)
best_candidate = proj(best_candidate, args.eps, args.inf_norm,
args.discrete)
with torch.no_grad():
adv_label = torch.argmax(
cnn_model.predict_scores(best_candidate + x0))
if adv_label != y0:
success = 1
if args.inf_norm:
print('Adversarial Label', adv_label.item(), 'Norm:',
best_candidate.abs().max().item())
else:
print('Adversarial Label', adv_label.item(), 'Norm:',
best_candidate.norm().item())
return query_count, success
query_count += args.q
# not successful (ran out of query budget)
return query_count, success
def fit_augmented_objective(
model, augmented_objective: AugmentedLagrangianMCObjective,
x_train: Tensor, y_train: Tensor):
sampler = SobolQMCNormalSampler(num_samples=1500)
print()
print("Optimizing augmented lagrangian on GP surrogate")
print(f"x_1\tx_2\tf_est\tc1_est\tmult1\tmult2")
for i in range(5):
acqfn = qSimpleRegret(model=model,
sampler=sampler,
objective=augmented_objective)
canditate, _ = optimize_acqf(acq_function=acqfn,
bounds=Tensor([[0.0, 0.0], [6.0, 6.0]]),
q=1,
num_restarts=1,
raw_samples=500)
x = canditate.detach()
samples = sampler(model.posterior(x))
augmented_objective.update_mults(samples)
augmented_objective.r = 100.0
x_ = x.numpy()[0]
acqfn_ = acqfn(x).detach().numpy()[0]
pred_ = model.posterior(x).mean.detach().numpy()[0]
mults_ = augmented_objective.mults.detach().numpy()
print(
f"{x_[0]:>6.4f}\t{x_[1]:>6.4f}\t{pred_[0]:>6.4f}\t{pred_[1]:>6.4f}\t{mults_[0][0]:>6.4f}"
)
f_best = augmented_objective(model.posterior(x).mean)
ei = qExpectedImprovement(
model=model,
best_f=f_best,
sampler=sampler,
objective=augmented_objective #linearized_objective
)
for i in range(x_train.shape[0]):
xx = x_train[i, :]
print(i, xx, ei(xx.unsqueeze(-1).T), f_best,
augmented_objective(y_train[i, :].unsqueeze(-1).T),
y_train[i, :])
canditate, _ = optimize_acqf(acq_function=ei,
bounds=Tensor([[0.0, 0.0], [6.0, 6.0]]),
q=1,
num_restarts=10,
raw_samples=500)
x_new = canditate.detach()
x_new_ = canditate.detach().numpy()[0]
f_best_ = f_best.detach().numpy()[0]
f_new_ = augmented_objective(
model.posterior(x_new).mean).detach().numpy()[0]
ei_new_ = ei(x_new).detach().numpy()[0]
print()
print("Optimizing EI on linearized objective")
print(f"x_1\tx_2\tf_best\tf_new_\tei")
print(f"{x_new_[0]:>6.4f}\t", f"{x_new_[1]:>6.4f}\t", f"{f_best_:>6.4f}\t",
f"{f_new_:>6.4f}\t", f"{ei_new_:>6.4f}")
return x_new
def generate_batch(
state,
model, # GP model
X, # Evaluated points on the domain [0, 1]^d
Y, # Function values
batch_size,
n_candidates=None, # Number of candidates for Thompson sampling
num_restarts=10,
raw_samples=512,
acqf="ts", # "ei" or "ts"
deup=False,
turbo=True,
):
dim = X.shape[-1]
assert acqf in ("ts", "ei")
assert X.min() >= 0.0 and X.max() <= 1.0 and torch.all(torch.isfinite(Y))
if n_candidates is None:
n_candidates = min(5000, max(2000, 200 * X.shape[-1]))
# Scale the TR to be proportional to the lengthscales
x_center = X[Y.argmax(), :].clone()
if not deup:
weights = model.covar_module.base_kernel.lengthscale.squeeze().detach()
else:
weights = model.f_predictor.covar_module.base_kernel.lengthscale.squeeze(
).detach()
weights = weights / weights.mean()
weights = weights / torch.prod(weights.pow(1.0 / len(weights)))
tr_lb = torch.clamp(x_center - weights * state.length / 2.0, 0.0, 1.0)
tr_ub = torch.clamp(x_center + weights * state.length / 2.0, 0.0, 1.0)
if not turbo:
tr_lb = torch.zeros(dim)
tr_ub = torch.ones(dim)
if acqf == "ts":
sobol = SobolEngine(dim, scramble=True)
pert = sobol.draw(n_candidates).to(dtype=dtype, device=device)
pert = tr_lb + (tr_ub - tr_lb) * pert
# Create a perturbation mask
prob_perturb = min(20.0 / dim, 1.0)
mask = (torch.rand(n_candidates, dim, dtype=dtype, device=device) <=
prob_perturb)
ind = torch.where(mask.sum(dim=1) == 0)[0]
mask[ind,
torch.randint(0, dim - 1, size=(len(ind), ), device=device)] = 1
# Create candidate points from the perturbations and the mask
X_cand = x_center.expand(n_candidates, dim).clone()
X_cand[mask] = pert[mask]
# Sample on the candidate points
thompson_sampling = MaxPosteriorSampling(model=model,
replacement=False)
X_next = thompson_sampling(X_cand, num_samples=batch_size)
elif acqf == "ei":
if batch_size > 1:
ei = qExpectedImprovement(model, Y.max(), maximize=True)
else:
ei = ExpectedImprovement(model, Y.max(), maximize=True)
try:
X_next, acq_value = optimize_acqf(
ei,
bounds=torch.stack([tr_lb, tr_ub]),
q=batch_size,
num_restarts=num_restarts,
raw_samples=raw_samples,
)
except NotPSDError:
sobol = SobolEngine(dim, scramble=True)
pert = sobol.draw(batch_size).to(dtype=dtype, device=device)
pert = tr_lb + (tr_ub - tr_lb) * pert
X_next = pert
print(
'Warning: NotPSDError, using {} purely random candidates for this step'
.format(batch_size))
return X_next