Use optuna terminators for early stopping in optuna.OptunaSearch?

1

Optuna provide its own library of schedulers: optuna.terminator, and I’m wondering if there is a way to use these with OptunaSearch?

Background

It does not seem possible to use ASHAScheduler with searching a space defined within OptunaSearch (which is necessary when doing multi-objective EHVI). So the following fails:

from ray.tune.search.optuna import OptunaSearch
import optuna

config['gpus_per_trial'] = 1
config['cpus_per_trial'] = 2
config['num_samples'] = 50
config['ashascheduler_grace'] = 5
config['ashascheduler_reduction'] = 2
config['max_epochs'] = 25

space = {
    "lr": optuna.distributions.LogUniformDistribution(1e-7, 1e-5),
    "alpha": optuna.distributions.FloatDistribution(1, 1e4),
}

scheduler = ASHAScheduler(
        max_t=config['max_epochs'],
        grace_period=config['ashascheduler_grace'],
        reduction_factor=config['ashascheduler_reduction'])

algo = OptunaSearch(
    space,
    metric=["profile_loss", "count_loss", "count_corr", "prof_corr"],
    mode=["min", "min", "max", "max"],
)

tuner = tune.Tuner(
    tune.with_resources(
        tune.with_parameters(train_scBPnet),
        resources={"cpu": config['cpus_per_trial'], "gpu": config['gpus_per_trial']}
    ),
    tune_config=tune.TuneConfig(
        num_samples=config['num_samples'],
        search_alg=algo,
        scheduler=scheduler,
    ),
    param_space=config,
)
results = tuner.fit()

with open(f'{save_dir}/results.pkl', 'wb') as f:
    pickle.dump(results, f)