Skip to main content
Version: Next

Optuna Sweeper plugin

PyPI PyPI - License PyPI - Python Version PyPI - DownloadsΒ Example (Click Here)Β Example (Click Here)

This plugin enables Hydra applications to utilize Optuna for the optimization of the parameters of experiments.

Installation​

This plugin requires hydra-core>=1.1.0. Please install it with the following command:

pip install hydra-core --upgrade

You can install the plugin via pip:

pip install hydra-optuna-sweeper --upgrade

There are several standard approaches for configuring plugins. Check this page for more information.

Usage​

Please set hydra/sweeper to optuna in your config file.

defaults:
- override hydra/sweeper: optuna

Alternatively, add hydra/sweeper=optuna option to your command line.

The default configuration is here.

Example 1: Single-Objective Optimization​

We include an example in this directory. example/sphere.py implements a simple benchmark function to be minimized.

You can discover the Optuna sweeper parameters with:

python example/sphere.py hydra/sweeper=optuna --cfg hydra -p hydra.sweeper
# @package hydra.sweeper
sampler:
_target_: optuna.samplers.TPESampler
seed: 123
consider_prior: true
prior_weight: 1.0
consider_magic_clip: true
consider_endpoints: false
n_startup_trials: 10
n_ei_candidates: 24
multivariate: false
warn_independent_sampling: true
_target_: hydra_plugins.hydra_optuna_sweeper.optuna_sweeper.OptunaSweeper
direction: minimize
storage: null
study_name: sphere
n_trials: 20
n_jobs: 1
params:
x: range(-5.5,5.5,step=0.5)
y: choice(-5,0,5)

The function decorated with @hydra.main() returns a float which we want to minimize, the minimum is 0 and reached for:

x: 0
y: 0

To run optimization, clone the code and run the following command in the plugins/hydra_optuna_sweeper directory:

python example/sphere.py --multirun

You can also override the search space parametrization:

python example/sphere.py --multirun 'x=interval(-5.0, 5.0)' 'y=interval(0, 10)'

You might find the optimization_results.yaml file (i.e. best params and best value) under multirun logs folder:

name: optuna
best_params:
x: 0.0
'y': 0
best_value: 0.0

Sampler configuration​

This plugin supports Optuna's samplers. You can change the sampler by overriding hydra/sweeper/sampler or change sampler settings within hydra.sweeper.sampler.

Search space configuration​

This plugin supports Optuna's distributions to configure search spaces. They can be defined either through commandline override or config file.

Configuring through commandline override​

Hydra provides a override parser that support rich syntax. Please refer to OverrideGrammer/Basic and OverrideGrammer/Extended for details.

Interval override​

By default, interval is converted to UniformDistribution. You can use IntUniformDistribution, LogUniformDistribution or IntLogUniformDistribution by casting the interval to int and tagging it with log.

Example for interval override
python example/sphere.py --multirun 'x=int(interval(-5.0, 5.0))' 'y=tag(log, interval(1, 10))'

The output is as follows:

[HYDRA] Study name: sphere
[HYDRA] Storage: None
[HYDRA] Sampler: TPESampler
[HYDRA] Directions: ['minimize']
[HYDRA] Launching 1 jobs locally
[HYDRA] #0 : x=-3 y=1.6859762540733367
[HYDRA] Launching 1 jobs locally
[HYDRA] #1 : x=1 y=5.237816870668193
...
[HYDRA] Best parameters: {'x': 0, 'y': 1.0929184723430116}
[HYDRA] Best value: 1.1944707871885822

Range override​

range is converted to IntUniformDistribution. If you apply shuffle to range, CategoricalDistribution is used instead. If any of range's start, stop or step is of type float, it will be converted to DiscreteUniformDistribution

Example for range override
python example/sphere.py --multirun 'x=range(-5.0, 5.0)' 'y=shuffle(range(-5, 5))'

The output is as follows:

[HYDRA] Study name: sphere
[HYDRA] Storage: None
[HYDRA] Sampler: TPESampler
[HYDRA] Directions: ['minimize']
[HYDRA] Launching 1 jobs locally
[HYDRA] #0 : x=-3 y=-4
[HYDRA] Launching 1 jobs locally
[HYDRA] #1 : x=1 y=-1
...
[HYDRA] Best parameters: {'x': 0, 'y': -1}
[HYDRA] Best value: 1.0

Choice override​

choice is converted to CategoricalDistribution.

Example for choice override
python example/sphere.py --multirun 'x=choice(-5.0, 0.0, 5.0)' 'y=choice(0, 1, 2, 3, 4, 5)'

The output is as follows:

[HYDRA] Study name: sphere
[HYDRA] Storage: None
[HYDRA] Sampler: TPESampler
[HYDRA] Directions: ['minimize']
[HYDRA] Launching 1 jobs locally
[HYDRA] #0 : x=5.0 y=5
[HYDRA] Launching 1 jobs locally
[HYDRA] #1 : x=5.0 y=2
...
[HYDRA] Best parameters: {'x': 0.0, 'y': 0}
[HYDRA] Best value: 0.0

Configuring through config file​

The syntax in config file is consistent with the above commandline override. For example, a commandline override x=range(1,4) can be expressed in config file as x: range(1,4) under the hydra.sweeper.params node.

Example 2: Multi-Objective Optimization​

In the same directory, example/multi-objective.py implements a simple benchmark function, which has two objective values. We want to minimize two objectives simultaneously.

You can discover the Optuna sweeper parameters with:

python example/multi-objective.py hydra/sweeper=optuna --cfg hydra -p hydra.sweeper
Configuration of the multi-objective optimization example
# @package hydra.sweeper
sampler:
_target_: optuna.samplers.NSGAIISampler
seed: 123
population_size: 50
mutation_prob: null
crossover_prob: 0.9
swapping_prob: 0.5
constraints_func: null
_target_: hydra_plugins.hydra_optuna_sweeper.optuna_sweeper.OptunaSweeper
direction:
- minimize
- minimize
storage: null
study_name: multi-objective
n_trials: 20
n_jobs: 1
params:
x: range(0, 5, step=0.5)
y: range(0, 3, step=0.5)

To run the optimization, use the following command in the plugins/hydra_optuna_sweeper directory:

python example/multi-objective.py --multirun

For problems with trade-offs between two different objectives, there may be no single solution that simultaneously minimizes both objectives. Instead, we obtained a set of solutions, namely Pareto optimal solutions, that show the best trade-offs possible between the objectives. In the following figure, the blue dots show the Pareto optimal solutions in the optimization results.

Pareto-optimal solutions

EXPERIMENTAL: Custom-Search-Space Optimization​

Hydra's Optuna Sweeper allows users to provide a hook for custom search space configuration. This means you can work directly with the optuna.trial.Trial object to suggest parameters. To use this feature, define a python function with signature Callable[[DictConfig, optuna.trial.Trial], None] and set the hydra.sweeper.custom_search_space key in your config to target that function.

You can find a full example in the same directory as before, where example/custom-search-space-objective.py implements a benchmark function to be minimized. The example shows the use of Optuna's pythonic search spaces in combination with Hydra. Part of the search space configuration is defined in config files, and part of it is written in Python.

Example: Custom search space configuration
defaults:
- override hydra/sweeper: optuna

hydra:
sweeper:
sampler:
seed: 123
direction: minimize
study_name: custom-search-space
storage: null
n_trials: 20
n_jobs: 1

params:
x: range(-5.5, 5.5, 0.5)
y: choice(-5, 0, 5)
# `custom_search_space` should be a dotpath pointing to a
# callable that provides search-space configuration logic:
custom_search_space: .custom-search-space-objective.configure

x: 1
y: 1
z: 100
max_z_difference_from_x: 0.5
# example/custom-search-space-objective.py

...

def configure(cfg: DictConfig, trial: Trial) -> None:
x_value = trial.params["x"]
trial.suggest_float(
"z",
x_value - cfg.max_z_difference_from_x,
x_value + cfg.max_z_difference_from_x,
)
trial.suggest_float("+w", 0.0, 1.0) # note +w here, not w as w is a new parameter

...

The method that custom_search_space points to must accepts both a DictConfig with already set options and a trial object which needs further configuration. In this example we limit z the difference between x and z to be no more than 0.5. Note that this custom_search_space API should be considered experimental and is subject to change.

Order of trial configuration​

Configuring a trial object is done in the following sequence:

  • search space parameters are set from the hydra.sweeper.params key in the config
  • Command line overrides are set
  • custom_search_space parameters are set

It is not allowed to set search space parameters in the custom_search_space method for parameters which have a fixed value from command line overrides. Trial.user_attrs can be inspected to find any of such fixed parameters.