import os from copy import deepcopy from bayesmark.abstract_optimizer import AbstractOptimizer from bayesmark.experiment import experiment_main from bayesmark.space import JointSpace import torch import logging from sklearn.preprocessing import MinMaxScaler, StandardScaler import numpy as np from sklearn.compose import ColumnTransformer from scipy.special import logit, expit # from scipy.stats import qmc from torch.quasirandom import SobolEngine from .scripts import acquisition_functions, tune_input_warping class PFNOptimizer(AbstractOptimizer): # Used for determining the version number of package used # primary_import = "" def __init__(self, api_config, pfn_file, minimize=True, acqf_optimizer_name="lbfgs", sobol_sampler=False, device="cpu:0", fit_encoder_from_step=None, verbose=False, rand_bool=False, sample_only_valid=False, round_suggests_to=4, min_initial_design=0, max_initial_design=None, rand_sugg_after_x_steps_of_stagnation=None, fixed_initial_guess=None,minmax_encode_y=False,**acqf_kwargs): """Build wrapper class to use optimizer in benchmark. Parameters ---------- api_config : dict-like of dict-like Configuration of the optimization variables. See API description. """ assert not 'fit_encoder' in acqf_kwargs AbstractOptimizer.__init__(self, api_config) # Do whatever other setup is needed # ... self.space_x = JointSpace(api_config) self.bounds = self.space_x.get_bounds() self.device = device self.model = torch.load(pfn_file) if pfn_file.startswith('/') else torch.load(os.path.dirname(__file__) + '/' + pfn_file) self.X = [] self.y = [] self.api_config = {key: value for key, value in sorted(api_config.items())} self.hp_names = list(self.api_config.keys()) # self.model.encoder.num_features = 18 self.epsilon = 1e-8 self.minimize = minimize self.sobol_sampler = sobol_sampler self.create_scaler() self.sobol = SobolEngine(len(self.max_values), scramble=True) self.acqf_optimizer_name = acqf_optimizer_name self.acqf_kwargs = acqf_kwargs self.fit_encoder_from_step = fit_encoder_from_step assert not (rand_bool and sample_only_valid) self.rand_bool = rand_bool self.sample_only_valid = sample_only_valid self.verbose = verbose self.round_suggests_to = round_suggests_to self.min_initial_design = min_initial_design self.max_initial_design = max_initial_design self.fixed_initial_guess = fixed_initial_guess self.minmax_encode_y = minmax_encode_y self.rand_sugg_after_x_steps_of_stagnation = rand_sugg_after_x_steps_of_stagnation self.model.eval() print(api_config) def create_scaler(self): list_of_scalers = [] self.min_values = [] self.max_values = [] self.spaces = [] self.types = [] for i, feature in enumerate(self.api_config): # list_of_scalers.append((feature, MinMaxScaler(feature_range),i)) self.spaces.append(self.api_config[feature].get("space", "bool")) self.types.append(self.api_config[feature]["type"]) if self.types[-1] == "bool": feature_range = [0, 1] else: feature_range = list(self.api_config[feature]["range"]) feature_range[0] = self.transform_feature(feature_range[0], -1) feature_range[1] = self.transform_feature(feature_range[1], -1) self.min_values.append(feature_range[0]) self.max_values.append(feature_range[1]) self.column_scaler = ColumnTransformer(list_of_scalers) self.max_values: np.array = np.array(self.max_values) self.min_values: np.array = np.array(self.min_values) def transform_feature_inverse(self, x, feature_index): if self.spaces[feature_index] == "log": x = np.exp(x) elif self.spaces[feature_index] == "logit": x = expit(x) if self.types[feature_index] == "int": if self.rand_bool: x = int(x) + int(np.random.rand() < (x-int(x))) else: x = int(np.round(x)) elif self.types[feature_index] == "bool": if self.rand_bool: x = np.random.rand() < x else: x = bool(np.round(x)) return x def transform_feature(self, x, feature_index): if np.isinf(x) or np.isnan(x): return 0 if self.spaces[feature_index] == "log": x = np.log(x) elif self.spaces[feature_index] == "logit": x = logit(x) elif self.types[feature_index] == "bool": x = int(x) return x def random_suggest(self): self.rand_prev = True if self.sobol_sampler: # sampler = qmc.Sobol(d=len(self.max_values), scramble=False) # temp_guess = sampler.random_base2(m=len(self.max_values)) temp_guess = self.sobol.draw(1).numpy()[0] temp_guess = temp_guess * (self.max_values - self.min_values) + self.min_values x_guess = {} for j, feature in enumerate(self.api_config): x = self.transform_feature_inverse(temp_guess[j], j) x_guess[feature] = x x_guess = [x_guess] else: x_guess = {} for i, feature in enumerate(self.api_config): temp_guess = np.random.uniform(self.min_values[i], self.max_values[i], 1)[0] temp_guess = self.transform_feature_inverse(temp_guess, i) x_guess[feature] = temp_guess x_guess = [x_guess] return x_guess def transform_back(self, x_guess): if self.round_suggests_to is not None: x_guess = np.round(x_guess, self.round_suggests_to) # make sure very similar values are actually the same x_guess = x_guess * (self.max_values - self.min_values) + self.min_values x_guess = x_guess.tolist() return self.transform_inverse(x_guess) def min_max_encode(self, temp_X): # this, combined with transform is the inverse of transform_back temp_X = (temp_X - self.min_values) / (self.max_values - self.min_values) temp_X = torch.tensor(temp_X).to(torch.float32) temp_X = torch.clamp(temp_X, min=0., max=1.) return temp_X @torch.no_grad() def suggest(self, n_suggestions=1): """Get suggestion from the optimizer. Parameters ---------- n_suggestions : int Desired number of parallel suggestions in the output Returns ------- next_guess : list of dict List of `n_suggestions` suggestions to evaluate the objective function. Each suggestion is a dictionary where each key corresponds to a parameter being optimized. """ assert n_suggestions == 1, "Only one suggestion at a time is supported" # Do whatever is needed to get the parallel guesses # ... # scaler = MinMaxScaler() # scaler.fit(self.X) try: num_initial_design = max(len(self.bounds), self.min_initial_design) if self.max_initial_design is not None: num_initial_design = min(num_initial_design, self.max_initial_design) if len(self.X) < num_initial_design: if len(self.X) == 0 and self.fixed_initial_guess is not None: x_guess = [self.transform_back(np.array([self.fixed_initial_guess for _ in range(len(self.bounds))]))] else: x_guess = self.random_suggest() return x_guess else: temp_X = np.array(self.X) temp_X = self.min_max_encode(temp_X) if self.minmax_encode_y: temp_y = MinMaxScaler().fit_transform(np.array(self.y).reshape(-1, 1)).reshape(-1) else: temp_y = np.array(self.y) temp_y = torch.tensor(temp_y).to(torch.float32) if self.rand_sugg_after_x_steps_of_stagnation is not None \ and len(self.y) > self.rand_sugg_after_x_steps_of_stagnation\ and not self.rand_prev: if temp_y[:-self.rand_sugg_after_x_steps_of_stagnation].max() == temp_y.max(): print(f"Random suggestion after >= {self.rand_sugg_after_x_steps_of_stagnation} steps of stagnation") x_guess = self.random_suggest() return x_guess if self.verbose: from matplotlib import pyplot as plt print(f"{temp_X=}, {temp_y=}") if temp_X.shape[1] == 2: from scipy.stats import rankdata plt.title('Observations, red -> blue.') plt.scatter(temp_X[:,0], temp_X[:,1], cmap='RdBu', c=rankdata(temp_y)) plt.show() temp_X = temp_X.to(self.device) temp_y = temp_y.to(self.device) if self.fit_encoder_from_step and self.fit_encoder_from_step <= len(self.X): with torch.enable_grad(): w = tune_input_warping.fit_input_warping(self.model, temp_X, temp_y) temp_X_warped = w(temp_X).detach() else: temp_X_warped = temp_X with torch.enable_grad(): if self.acqf_optimizer_name == "lbfgs": def rand_sample_func(n): pre_samples = torch.rand(n, temp_X_warped.shape[1], device='cpu') back_transformed_samples = [self.transform_back(sample) for sample in pre_samples] samples = np.array([self.transform(deepcopy(bt_sample)) for bt_sample in back_transformed_samples]) samples = self.min_max_encode(samples) return samples.to(self.device) if self.sample_only_valid: rand_sample_func = rand_sample_func # dims with bool or int are not continuous, thus no gradient opt is applied dims_wo_gradient_opt = [i for i, t in enumerate(self.types) if t != "real"] else: rand_sample_func = None dims_wo_gradient_opt = [] x_guess, x_options, eis, x_rs, x_rs_eis = acquisition_functions.optimize_acq_w_lbfgs( self.model, temp_X_warped, temp_y, device=self.device, verbose=self.verbose, rand_sample_func=rand_sample_func, dims_wo_gradient_opt=dims_wo_gradient_opt, **{'apply_power_transform':True,**self.acqf_kwargs} ) elif self.acqf_optimizer_name == 'adam': x_guess = acquisition_functions.optimize_acq(self.model, temp_X_warped, temp_y, apply_power_transform=True, device=self.device, **self.acqf_kwargs ).detach().cpu().numpy() else: raise ValueError("Optimizer not recognized, set `acqf_optimizer_name` to 'lbfgs' or 'adam'") back_transformed_x_options = [self.transform_back(x) for x in x_options] opt_X = np.array([self.transform(deepcopy(transformed_x_options)) for transformed_x_options in back_transformed_x_options]) opt_X = self.min_max_encode(opt_X) opt_new = ~(opt_X[:,None] == temp_X[None].cpu()).all(-1).any(1) for i, x in enumerate(opt_X): if opt_new[i]: if self.verbose: print(f"New point at pos {i}: {back_transformed_x_options[i], x_options[i]}") self.rand_prev = False return [back_transformed_x_options[i]] print('backup from initial rand search') back_transformed_x_options = [self.transform_back(x) for x in x_rs] opt_X = np.array([self.transform(deepcopy(transformed_x_options)) for transformed_x_options in back_transformed_x_options]) opt_X = self.min_max_encode(opt_X) opt_new = ~(opt_X[:,None] == temp_X[None].cpu()).all(-1).any(1) for i, x in enumerate(opt_X): if opt_new[i]: if self.verbose: print(f"New point at pos {i}: {back_transformed_x_options[i], x_rs[i]} with ei {x_rs_eis[i]}") self.rand_prev = False return [back_transformed_x_options[i]] print("No new points found, random suggestion") return self.random_suggest() except Exception as e: raise e def transform(self, X_dict): X_tf = [] for i, feature in enumerate(X_dict.keys()): X_dict[feature] = self.transform_feature(X_dict[feature], i) X_tf.append(X_dict[feature]) return X_tf def transform_inverse(self, X_list): X_tf = {} for i, hp_name in enumerate(self.hp_names): X_tf[hp_name] = self.transform_feature_inverse(X_list[i], i) return X_tf def observe(self, X, y): """Feed an observation back. Parameters ---------- X : list of dict-like Places where the objective function has already been evaluated. Each suggestion is a dictionary where each key corresponds to a parameter being optimized. y : array-like, shape (n,) Corresponding values where objective has been evaluated """ # Update the model with new objective function observations # ... # No return statement needed if np.isinf(y) and y > 0: y[:] = 1e10 if not np.isnan(y) and not np.isinf(y): assert len(y) == 1 and len(X) == 1, "Only one suggestion at a time is supported" X = {key: value for key, value in sorted(X[0].items())} assert list(X.keys()) == list(self.api_config.keys()) == list(self.hp_names) == list( self.space_x.param_list) if self.verbose: print(f"{X=}, {y=}") X = self.transform(X) if self.verbose: print(f"transformed {X=}") self.X.append(X) if self.minimize: self.y.append(-y[0]) else: self.y.append(y[0]) else: assert False def test(): from bayesmark.experiment import _build_test_problem, run_study, OBJECTIVE_NAMES #function_instance = _build_test_problem(model_name='ada', dataset='breast', scorer='nll', path=None) function_instance = _build_test_problem(model_name='kNN', dataset='boston', scorer='mse', path=None) # Setup optimizer api_config = function_instance.get_api_config() import os # check is file config = { "pfn_file": 'final_models/model_hebo_morebudget_9_unused_features_3.pt', "minimize": 1, "device": "cpu:0", "fit_encoder_from_step": None, 'pre_sample_size': 1_000, 'num_grad_steps': 15_000, 'num_candidates': 10, 'rand_bool': True, } opt = PFNOptimizer(api_config, verbose=True, **config) function_evals, timing, suggest_log = run_study( opt, function_instance, 50, 1, callback=None, n_obj=len(OBJECTIVE_NAMES), ) if __name__ == "__main__": import uuid from bayesmark.serialize import XRSerializer from bayesmark.cmd_parse import CmdArgs import bayesmark.cmd_parse as cmd import bayesmark.constants as cc description = "Run a study with one benchmark function and an optimizer" args = cmd.parse_args(cmd.experiment_parser(description)) run_uuid = uuid.UUID(args[CmdArgs.uuid]) # set global logging level logging.basicConfig(level=logging.DEBUG) # This is the entry point for experiments, so pass the class to experiment_main to use this optimizer. # This statement must be included in the wrapper class file: experiment_main(PFNOptimizer, args=args)