diff --git a/README.md b/README.md index ebac551..7f98224 100644 --- a/README.md +++ b/README.md @@ -30,16 +30,26 @@ Install the Python dependencies: pip install -r requirements.txt ``` +The requirement file installs [BoTorch](https://botorch.org/) and +[GPyTorch](https://gpytorch.ai/), which must match the local PyTorch +installation. If you are using a GPU-specific wheel, install the desired +`torch`/`torchvision` versions first so that `pip install -r +requirements.txt` selects compatible builds. + ### Usage Run the script with desired arguments: ``` -python train.py [-h] [--random_hparams | --llm] [--rounds ROUNDS] [--search_space SEARCH_SPACE] - [--dataset_dir DATASET_DIR] [--arch ARCH] [--save_dir SAVE_DIR] +python train.py [-h] [--random_hparams | --llm] [--llm_rounds LLM_ROUNDS] + [--bo_rounds BO_ROUNDS] [--bo_batch_size BO_BATCH_SIZE] [--hybrid_bo] + [--search_space SEARCH_SPACE] [--dataset_dir DATASET_DIR] + [--arch ARCH] [--save_dir SAVE_DIR] [--train_batch_size TRAIN_BATCH_SIZE] [--eval_batch_size EVAL_BATCH_SIZE] - [--learning_rate LEARNING_RATE] [--weight_decay WEIGHT_DECAY] [--label_smoothing LABEL_SMOOTHING] - [--optimizer OPTIMIZER] [--num_train_epochs NUM_TRAIN_EPOCHS] [--seed SEED] - [--max_train_batches MAX_TRAIN_BATCHES] [--max_eval_batches MAX_EVAL_BATCHES] + [--learning_rate LEARNING_RATE] [--weight_decay WEIGHT_DECAY] + [--label_smoothing LABEL_SMOOTHING] [--optimizer OPTIMIZER] + [--num_train_epochs NUM_TRAIN_EPOCHS] [--seed SEED] + [--max_train_batches MAX_TRAIN_BATCHES] + [--max_eval_batches MAX_EVAL_BATCHES] ``` To run the LLM-based hyperparameter search with a single seed: @@ -48,6 +58,17 @@ python train.py --llm --seed SEED ``` By default, this tunes the optimizer, learning rate, batch size, weight decay, and label smoothing. GPT-4 generated some ranges which we found reasonable. To reproduce our results, run the command with five different random seeds. +### Hybrid LLM + Bayesian optimization workflow + +Enable `--hybrid_bo` to hand off the LLM suggestions to a [BoTorch](https://botorch.org/)-powered Bayesian optimization stage. The `--llm_rounds` flag controls the number of configurations proposed by the LLM, while `--bo_rounds` determines how many additional candidates are generated by Bayesian optimization. The BoTorch helper evaluates candidates in batches of `--bo_batch_size` (default 1). + +```bash +python train.py --llm --hybrid_bo --llm_rounds 4 --bo_rounds 6 --bo_batch_size 2 \ + --seed 0 --search_space constrained +``` + +The JSON artifact saved under `--save_dir` records which method (LLM or BoTorch) produced each configuration along with the configured round counts. + To run random hyperparameter search: ``` python train.py --random_hparams --seed SEED diff --git a/cifar/train.py b/cifar/train.py index aeb66e1..17c3cae 100644 --- a/cifar/train.py +++ b/cifar/train.py @@ -7,7 +7,7 @@ import sys import time from pathlib import Path -from typing import Tuple +from typing import Dict, List, Tuple import numpy as np import torch @@ -27,6 +27,10 @@ sys.path.insert(0, str(repo_root)) from ollama_client import OllamaChatClient +from utils.botorch_search import ( + get_default_botorch_search_space, + suggest_next_candidates, +) from cifar.pipeline import construct_resnet9, get_cifar10_dataset from cifar.vit import ViT @@ -149,7 +153,7 @@ def update_configs(self, config, validation_error, validation_loss): self.configs.append((config, validation_error, validation_loss)) -def llm_hyperparameter_search(tuner, rounds=10): +def llm_hyperparameter_search(tuner, args, rounds=10): """Perform hyperparameter search using the LLM tuner.""" results = [] training_exception = None @@ -157,31 +161,94 @@ def llm_hyperparameter_search(tuner, rounds=10): config = None try: config = tuner.suggest_hyperparameters(training_exception=training_exception) - validation_loss, validation_error = train_and_evaluate(config) - tuner.update_configs(config, validation_error, validation_loss) - # save and print results + validation_loss, validation_error, metrics = train_and_evaluate(config, args) + resolved_config = metrics.get("config", config) + tuner.update_configs(resolved_config, validation_error, validation_loss) results.append({ - "config": config, - "accuracy": 1-validation_error, - "loss": validation_loss - + "config": resolved_config, + "accuracy": metrics["eval_acc"], + "loss": validation_loss, + "method": "llm", }) - print(f"Round {i+1}: Validation Error = {validation_error}, Validation Loss = {validation_loss}") - print(f"Suggested Config: {config}") + print( + f"Round {i+1}: Validation Error = {validation_error}, Validation Loss = {validation_loss}" + ) + print(f"Suggested Config: {resolved_config}") training_exception = None except Exception as e: print(f"Error during training: {e}") print("Continuing to the next round...") - # When an error occurs, append None for accuracy and loss results.append({ "config": config, "accuracy": None, - "loss": None + "loss": None, + "method": "llm", + "error": str(e), }) training_exception = str(e) return results +def run_botorch_hyperparameter_search( + tuner: LLMHyperparameterTuner, + args, + total_rounds: int, + results: List[Dict[str, float]], +) -> None: + if total_rounds <= 0: + return + + search_space = get_default_botorch_search_space(args.search_space) + remaining = total_rounds + logger = logging.getLogger(__name__) + + while remaining > 0: + batch_size = min(args.bo_batch_size, remaining) + try: + candidates = suggest_next_candidates( + tuner.configs, + search_space, + num_candidates=batch_size, + ) + except Exception as exc: + logger.warning("Failed to generate BoTorch candidates: %s", exc) + break + + if not candidates: + logger.warning("No BoTorch candidates were generated; stopping early.") + break + + for config in candidates: + try: + validation_loss, validation_error, metrics = train_and_evaluate(config, args) + resolved_config = metrics.get("config", config) + tuner.update_configs(resolved_config, validation_error, validation_loss) + results.append( + { + "config": resolved_config, + "accuracy": metrics["eval_acc"], + "loss": validation_loss, + "method": "botorch", + } + ) + logger.info( + "BoTorch candidate evaluated with error %.4f and loss %.4f", + validation_error, + validation_loss, + ) + except Exception as exc: + logger.exception("BoTorch candidate evaluation failed: %s", exc) + results.append( + { + "config": config, + "accuracy": None, + "loss": None, + "method": "botorch", + "error": str(exc), + } + ) + remaining -= len(candidates) + def parse_args(): parser = argparse.ArgumentParser(description="Train ResNet-9 model on CIFAR-10 dataset.") tuner_group = parser.add_mutually_exclusive_group() @@ -189,7 +256,11 @@ def parse_args(): tuner_group.add_argument('--random_hparams', action='store_true', help='Randomly sample hyperparameters', default=False) tuner_group.add_argument('--llm', action='store_true', help='Use LLM for hyperparameter search', default=False) # llm tuning specific hyperparams - parser.add_argument('--rounds', type=int, default=10, help='Number of times we interact with the LLM to get hyperparameters') + parser.add_argument('--rounds', type=int, default=None, help='Deprecated: use --llm_rounds. Number of LLM proposals to evaluate.') + parser.add_argument('--llm_rounds', type=int, default=None, help='Number of LLM-driven proposals to evaluate during tuning.') + parser.add_argument('--bo_rounds', type=int, default=0, help='Number of BoTorch proposals to evaluate after the LLM stage.') + parser.add_argument('--bo_batch_size', type=int, default=1, help='Number of BoTorch candidates to evaluate per optimization step.') + parser.add_argument('--hybrid_bo', action='store_true', default=False, help='Enable the hybrid LLM + BoTorch workflow.') parser.add_argument('--search_space', type=str, default='constrained', help='Search space for LLM tuning') # general hparams parser.add_argument("--dataset_dir", type=str, default="./data", help="A folder to download or load CIFAR-10 dataset.") @@ -223,8 +294,88 @@ def parse_args(): ) parser.add_argument("--seed", type=int, default=1004, help="A seed for reproducible training pipeline.") args = parser.parse_args() + if args.llm_rounds is None: + args.llm_rounds = args.rounds if args.rounds is not None else 10 + if args.llm_rounds < 0: + raise ValueError('--llm_rounds must be non-negative') + if args.bo_rounds < 0: + raise ValueError('--bo_rounds must be non-negative') + if args.bo_batch_size <= 0: + raise ValueError('--bo_batch_size must be positive') return args + +def extract_hyperparameters_from_args(args) -> Dict[str, float]: + return { + "learning_rate": args.learning_rate, + "weight_decay": args.weight_decay, + "train_batch_size": args.train_batch_size, + "label_smoothing": args.label_smoothing, + "optimizer": args.optimizer, + } + + +def resolve_hyperparameters(args, overrides: Dict[str, float]) -> Dict[str, object]: + resolved = { + "learning_rate": float(overrides.get("learning_rate", args.learning_rate)), + "weight_decay": float(overrides.get("weight_decay", args.weight_decay)), + "train_batch_size": int(overrides.get("train_batch_size", args.train_batch_size)), + "label_smoothing": float(overrides.get("label_smoothing", args.label_smoothing)), + "optimizer": str(overrides.get("optimizer", args.optimizer)).lower(), + } + for key, value in overrides.items(): + if key not in resolved: + resolved[key] = value + return resolved + + +def run_training_job(args, hyperparameters: Dict[str, object]) -> Dict[str, object]: + logger = logging.getLogger(__name__) + + resolved_hyperparameters = resolve_hyperparameters(args, hyperparameters) + + train_dataset = get_cifar10_dataset( + split="train", dataset_dir=args.dataset_dir, + ) + model = train( + dataset=train_dataset, + batch_size=resolved_hyperparameters['train_batch_size'], + num_train_epochs=args.num_train_epochs, + learning_rate=resolved_hyperparameters['learning_rate'], + weight_decay=resolved_hyperparameters['weight_decay'], + label_smoothing=resolved_hyperparameters['label_smoothing'], + optimizer=resolved_hyperparameters['optimizer'], + arch=args.arch, + hyps=resolved_hyperparameters, + max_train_batches=args.max_train_batches, + ) + + eval_train_dataset = get_cifar10_dataset(split="eval_train", dataset_dir=args.dataset_dir) + train_loss, train_acc = evaluate( + model=model, + dataset=eval_train_dataset, + batch_size=args.eval_batch_size, + max_batches=args.max_eval_batches, + ) + logger.info(f"Train loss: {train_loss}, Train Accuracy: {train_acc}") + + eval_dataset = get_cifar10_dataset(split="valid", dataset_dir=args.dataset_dir) + eval_loss, eval_acc = evaluate( + model=model, + dataset=eval_dataset, + batch_size=args.eval_batch_size, + max_batches=args.max_eval_batches, + ) + logger.info(f"Evaluation loss: {eval_loss}, Evaluation Accuracy: {eval_acc}") + + return { + "train_loss": train_loss, + "train_acc": train_acc, + "eval_loss": eval_loss, + "eval_acc": eval_acc, + "config": resolved_hyperparameters, + } + def train( dataset: data.Dataset, batch_size: int, @@ -402,7 +553,7 @@ def main(args): if args.seed is not None: torch.manual_seed(args.seed) np.random.seed(args.seed) - + save_filename = os.path.join(args.save_dir, f"results_trial_{args.seed}.json") print("Checking if results already saved in ", save_filename) if args.random_hparams: @@ -414,74 +565,37 @@ def main(args): # randomly sample hyperparameters hyps = sample_hyperparameters(args.arch) else: - lr = args.learning_rate - wd = args.weight_decay - bs = args.train_batch_size - ls = args.label_smoothing - optimizer = args.optimizer - hyps = { - "learning_rate": lr, - "weight_decay": wd, - "train_batch_size": bs, - "label_smoothing": ls, - "optimizer": optimizer - } + hyps = extract_hyperparameters_from_args(args) print("Hyperparameters:") for k, v in hyps.items(): print(f"{k}: {v}") logging.basicConfig(level=logging.INFO) - logger = logging.getLogger() - - train_dataset = get_cifar10_dataset( - split="train", dataset_dir=args.dataset_dir, - ) - model = train( - dataset=train_dataset, - batch_size=hyps['train_batch_size'], - num_train_epochs=args.num_train_epochs, - learning_rate=hyps['learning_rate'], - weight_decay=hyps['weight_decay'], - label_smoothing=hyps['label_smoothing'], - optimizer=hyps['optimizer'], - arch=args.arch, - hyps=hyps, - max_train_batches=args.max_train_batches, - ) - - eval_train_dataset = get_cifar10_dataset(split="eval_train", dataset_dir=args.dataset_dir) - train_loss, train_acc = evaluate( - model=model, - dataset=eval_train_dataset, - batch_size=args.eval_batch_size, - max_batches=args.max_eval_batches, - ) - logger.info(f"Train loss: {train_loss}, Train Accuracy: {train_acc}") + metrics = run_training_job(args, hyps) - eval_dataset = get_cifar10_dataset(split="valid", dataset_dir=args.dataset_dir) - eval_loss, eval_acc = evaluate( - model=model, - dataset=eval_dataset, - batch_size=args.eval_batch_size, - max_batches=args.max_eval_batches, - ) - logger.info(f"Evaluation loss: {eval_loss}, Evaluation Accuracy: {eval_acc}") - if args.random_hparams: - # save hyperparameters and results - save_results(hyps, {'train_loss': train_loss, 'train_acc': train_acc, 'eval_loss': eval_loss, 'eval_acc': eval_acc}, args.seed, save_filename) - - return train_loss, train_acc, eval_loss, eval_acc + save_results( + hyps, + { + 'train_loss': metrics['train_loss'], + 'train_acc': metrics['train_acc'], + 'eval_loss': metrics['eval_loss'], + 'eval_acc': metrics['eval_acc'], + }, + args.seed, + save_filename, + ) + + return ( + metrics['train_loss'], + metrics['train_acc'], + metrics['eval_loss'], + metrics['eval_acc'], + ) -def train_and_evaluate(hyperparameters): - args = parse_args() - args.learning_rate = hyperparameters["learning_rate"] - args.weight_decay = hyperparameters["weight_decay"] - args.train_batch_size = hyperparameters["train_batch_size"] - args.label_smoothing = hyperparameters["label_smoothing"] - args.optimizer = hyperparameters["optimizer"] - train_loss, train_acc, eval_loss, eval_acc = main(args) - eval_error = 1 - eval_acc - return eval_loss, eval_error +def train_and_evaluate(hyperparameters: Dict[str, object], args) -> Tuple[float, float, Dict[str, object]]: + metrics = run_training_job(args, hyperparameters) + eval_error = 1 - metrics["eval_acc"] + return metrics["eval_loss"], eval_error, metrics prompt_end = """You will get the validation error rate and loss before you need to specify the next configuration. The goal is to find the configuration that minimizes the error rate with the given budget, so you should explore different parts of the search space if the loss is not changing. Provide a config in JSON format. Do not put new lines or any extra characters in the response, only provide the config. Example config: { @@ -517,7 +631,7 @@ def train_and_evaluate(hyperparameters): args = parse_args() np.random.seed(args.seed) torch.manual_seed(args.seed) - + # performs hyperparameter search with an Ollama-hosted LLM, check if already run if args.llm: if not os.path.exists(args.save_dir): @@ -530,14 +644,27 @@ def train_and_evaluate(hyperparameters): tuner = LLMHyperparameterTuner(initial_prompt_constrained) elif args.search_space == 'unconstrained': tuner = LLMHyperparameterTuner(initial_prompt) - results = llm_hyperparameter_search(tuner, rounds=args.rounds) + else: + raise ValueError(f"Unknown search space: {args.search_space}") + + logging.basicConfig(level=logging.INFO) + results: List[Dict[str, float]] = [] + if args.llm_rounds > 0: + results.extend( + llm_hyperparameter_search(tuner, args, rounds=args.llm_rounds) + ) + if args.hybrid_bo and args.bo_rounds > 0: + run_botorch_hyperparameter_search(tuner, args, args.bo_rounds, results) results_dict = { "results": results, "search_space": args.search_space, - "rounds": args.rounds + "llm_rounds": args.llm_rounds, + "bo_rounds": args.bo_rounds if args.hybrid_bo else 0, + "hybrid_bo": args.hybrid_bo, + "bo_batch_size": args.bo_batch_size, } with open(save_filename, 'w') as f: - json.dump(results_dict, f, indent=4) + json.dump(results_dict, f, indent=4) # performs a single training run with given hyperparameters else: main(args) diff --git a/requirements.txt b/requirements.txt index 381f0ad..9aaf1fc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,6 @@ numpy matplotlib torch torchvision +gpytorch==1.10 +botorch==0.8.5 einops diff --git a/tests/test_botorch_search.py b/tests/test_botorch_search.py new file mode 100644 index 0000000..69a5277 --- /dev/null +++ b/tests/test_botorch_search.py @@ -0,0 +1,46 @@ +import sys +from pathlib import Path + +sys.path.append(str(Path(__file__).resolve().parents[1])) + +from utils.botorch_search import get_default_botorch_search_space, suggest_next_candidates + + +def test_suggest_next_candidates_returns_valid_configs(): + search_space = get_default_botorch_search_space("constrained") + history = [ + ( + { + "learning_rate": 1e-3, + "weight_decay": 1e-4, + "label_smoothing": 0.05, + "train_batch_size": 128, + "optimizer": "sgd", + }, + 0.35, + 0.8, + ), + ( + { + "learning_rate": 3e-3, + "weight_decay": 5e-5, + "label_smoothing": 0.02, + "train_batch_size": 256, + "optimizer": "adam", + }, + 0.28, + 0.7, + ), + ] + + candidates = suggest_next_candidates(history, search_space, num_candidates=2) + + assert len(candidates) == 2 + for candidate in candidates: + assert set(candidate) == set(search_space.keys()) + assert search_space["optimizer"]["choices"].count(candidate["optimizer"]) == 1 + assert search_space["train_batch_size"]["choices"].count(candidate["train_batch_size"]) == 1 + for name, spec in search_space.items(): + if spec["type"] == "continuous": + low, high = spec["bounds"] + assert low <= candidate[name] <= high diff --git a/tests/test_llm_hyperparameter_search.py b/tests/test_llm_hyperparameter_search.py index 561c9bf..305dc08 100644 --- a/tests/test_llm_hyperparameter_search.py +++ b/tests/test_llm_hyperparameter_search.py @@ -1,5 +1,6 @@ import sys from pathlib import Path +from types import SimpleNamespace from unittest import mock sys.path.append(str(Path(__file__).resolve().parents[1])) @@ -23,12 +24,13 @@ def update_configs(self, *args, **kwargs): def test_llm_hyperparameter_search_handles_failed_suggestions(): tuner = ErrorOnceTuner() + args = SimpleNamespace() with mock.patch.object(cifar_train, "train_and_evaluate", autospec=True) as mocked_eval: - results = cifar_train.llm_hyperparameter_search(tuner, rounds=1) + results = cifar_train.llm_hyperparameter_search(tuner, args, rounds=1) mocked_eval.assert_not_called() assert results == [ - {"config": None, "accuracy": None, "loss": None} + {"config": None, "accuracy": None, "loss": None, "method": "llm", "error": "LLM failure"} ] diff --git a/utils/botorch_search.py b/utils/botorch_search.py new file mode 100644 index 0000000..83d0c0c --- /dev/null +++ b/utils/botorch_search.py @@ -0,0 +1,211 @@ +"""BoTorch helpers for hybrid LLM + Bayesian optimization tuning.""" + +from __future__ import annotations + +import math +from collections import OrderedDict +from typing import Dict, Iterable, List, Sequence + +import torch +from botorch.acquisition.monte_carlo import qExpectedImprovement +from botorch.fit import fit_gpytorch_mll +from botorch.models import SingleTaskGP +from botorch.optim.optimize import optimize_acqf +from botorch.utils.transforms import normalize, unnormalize +from gpytorch.mlls import ExactMarginalLogLikelihood +from torch.quasirandom import SobolEngine + +def get_default_botorch_search_space(search_space: str = "constrained") -> OrderedDict[str, Dict[str, object]]: + """Return a default search space compatible with CIFAR training.""" + + space: OrderedDict[str, Dict[str, object]] = OrderedDict( + [ + ( + "learning_rate", + {"type": "continuous", "bounds": (1e-4, 1e-1), "transform": "log"}, + ), + ( + "weight_decay", + {"type": "continuous", "bounds": (1e-5, 1e-1), "transform": "log"}, + ), + ("label_smoothing", {"type": "continuous", "bounds": (0.0, 0.2)}), + ( + "train_batch_size", + {"type": "categorical", "choices": [32, 64, 128, 256, 512]}, + ), + ("optimizer", {"type": "categorical", "choices": ["sgd", "adam"]}), + ] + ) + + if search_space == "unconstrained": + space["learning_rate"]["bounds"] = (1e-5, 5e-1) + space["weight_decay"]["bounds"] = (1e-6, 5e-1) + space["label_smoothing"]["bounds"] = (0.0, 0.3) + space["train_batch_size"]["choices"] = [32, 64, 128, 256, 512, 1024] + elif search_space != "constrained": + raise ValueError(f"Unsupported search space: {search_space}") + + return space + + +def _encode_config(config: Dict[str, object], space: OrderedDict[str, Dict[str, object]]) -> List[float] | None: + vector: List[float] = [] + for name, spec in space.items(): + if spec["type"] == "continuous": + if name not in config: + return None + value = float(config[name]) + if value <= 0 and spec.get("transform") == "log": + return None + if spec.get("transform") == "log": + value = math.log(value) + vector.append(value) + elif spec["type"] == "categorical": + if name not in config: + return None + value = config[name] + choices: Sequence = spec["choices"] + index = _categorical_index(value, choices) + if index is None: + return None + vector.append(float(index)) + else: + raise ValueError(f"Unsupported parameter type: {spec['type']}") + return vector + + +def _categorical_index(value, choices: Sequence) -> int | None: + if isinstance(value, str): + normalized_value = value.lower() + mapping = {str(choice).lower(): idx for idx, choice in enumerate(choices)} + return mapping.get(normalized_value) + try: + return choices.index(value) + except ValueError: + if isinstance(value, (int, float)): + rounded = int(value) + if rounded in choices: + return choices.index(rounded) + return None + + +def _decode_config(vector: torch.Tensor, space: OrderedDict[str, Dict[str, object]]) -> Dict[str, object]: + config: Dict[str, object] = {} + for idx, (name, spec) in enumerate(space.items()): + raw_value = float(vector[idx]) + if spec["type"] == "continuous": + low, high = spec["bounds"] + value = raw_value + if spec.get("transform") == "log": + value = math.exp(value) + value = min(max(value, low), high) + config[name] = value + elif spec["type"] == "categorical": + choices: Sequence = spec["choices"] + index = int(round(raw_value)) + index = max(0, min(index, len(choices) - 1)) + config[name] = choices[index] + else: + raise ValueError(f"Unsupported parameter type: {spec['type']}") + return config + + +def _build_bounds(space: OrderedDict[str, Dict[str, object]], device: torch.device) -> torch.Tensor: + lowers: List[float] = [] + uppers: List[float] = [] + for spec in space.values(): + if spec["type"] == "continuous": + low, high = spec["bounds"] + if spec.get("transform") == "log": + low = math.log(low) + high = math.log(high) + lowers.append(low) + uppers.append(high) + else: + lowers.append(0.0) + uppers.append(float(len(spec["choices"]) - 1)) + return torch.tensor([lowers, uppers], dtype=torch.double, device=device) + + +def _prepare_training_data( + history: Iterable[tuple[Dict[str, object], float | None, float | None]], + space: OrderedDict[str, Dict[str, object]], +) -> tuple[torch.Tensor | None, torch.Tensor | None]: + encoded_configs: List[List[float]] = [] + targets: List[List[float]] = [] + for config, error_rate, _ in history: + if error_rate is None or not math.isfinite(error_rate): + continue + encoded = _encode_config(config, space) + if encoded is None: + continue + encoded_configs.append(encoded) + targets.append([-float(error_rate)]) + + if not encoded_configs: + return None, None + + train_X = torch.tensor(encoded_configs, dtype=torch.double) + train_Y = torch.tensor(targets, dtype=torch.double) + return train_X, train_Y + + +def _sample_random_candidates( + space: OrderedDict[str, Dict[str, object]], + num_candidates: int, +) -> List[Dict[str, object]]: + dim = len(space) + sobol = SobolEngine(dimension=dim, scramble=True) + samples = sobol.draw(num_candidates).to(dtype=torch.double) + bounds = _build_bounds(space, device=samples.device) + unnormalized = unnormalize(samples, bounds) + return [_decode_config(vector, space) for vector in unnormalized] + + +def suggest_next_candidates( + history: Sequence[tuple[Dict[str, object], float | None, float | None]], + search_space: OrderedDict[str, Dict[str, object]] | None = None, + num_candidates: int = 1, +) -> List[Dict[str, object]]: + """Suggest the next batch of hyperparameters using BoTorch.""" + + if num_candidates <= 0: + raise ValueError("num_candidates must be positive") + + space = search_space or get_default_botorch_search_space() + + train_X, train_Y = _prepare_training_data(history, space) + if train_X is None or train_Y is None or train_X.size(0) < 1: + return _sample_random_candidates(space, num_candidates) + + device = train_X.device + bounds = _build_bounds(space, device=device) + normalized_X = normalize(train_X, bounds) + + model = SingleTaskGP(normalized_X, train_Y) + mll = ExactMarginalLogLikelihood(model.likelihood, model) + fit_gpytorch_mll(mll) + + acquisition = qExpectedImprovement(model, best_f=train_Y.max()) + unit_bounds = torch.stack( + [ + torch.zeros(normalized_X.size(-1), dtype=torch.double, device=device), + torch.ones(normalized_X.size(-1), dtype=torch.double, device=device), + ] + ) + + candidates_normalized, _ = optimize_acqf( + acquisition, + bounds=unit_bounds, + q=num_candidates, + num_restarts=max(5, 2 * normalized_X.size(-1)), + raw_samples=256, + ) + candidates = unnormalize(candidates_normalized, bounds) + return [_decode_config(vector, space) for vector in candidates] + + +__all__ = [ + "get_default_botorch_search_space", + "suggest_next_candidates", +]