Optuna: Hyperparameter Optimization Framework

If a process running a trial is killed unexpectedly (e.g., by a cluster job scheduler), the trial remains in the RUNNING state. Optuna's experimental heartbeat mechanism can automatically transition these trials to FAIL.

Using study.optimize

When using study.optimize, configure heartbeat_interval and grace_period in RDBStorage. A trial will be marked as failed if no heartbeat is recorded within the grace_period after the heartbeat_interval.

import optuna

def objective(trial):
    # (Very time-consuming computation)

# Recording heartbeats every 60 seconds.
# Trials with no heartbeat for > 120 seconds are automatically failed.
storage = optuna.storages.RDBStorage(url="sqlite:///:memory:", heartbeat_interval=60, grace_period=120)
study = optuna.create_study(storage=storage)
study.optimize(objective, n_trials=100)

Using ask and tell

If you use the manual ask/tell API, the heartbeat mechanism does not automatically update states. You must manually check for stale trials and update them:

from datetime import datetime
import optuna

study = optuna.create_study(storage=...)
# Example: fail trials running for more than 1 day
grace_period = 3600*24
for t in study.get_trials(states=[optuna.trial.TrialState.RUNNING]):
    if (datetime.now() - t.datetime_start).total_seconds() > grace_period:
        study.tell(t, state=optuna.trial.TrialState.FAIL)

Retrying stale trials

You can use RetryHeartbeatStaleTrialCallback to automatically retry trials that have failed due to a lost heartbeat. This callback is invoked at the start of each new trial.

import optuna
from optuna.storages import RetryHeartbeatStaleTrialCallback

storage = optuna.storages.RDBStorage(
    url="sqlite:///:memory:",
    heartbeat_interval=60,
    grace_period=120,
    heartbeat_stale_trial_callback=RetryHeartbeatStaleTrialCallback(max_retry=3),
)

study = optuna.create_study(storage=storage)

To suggest $n$ variables that represent proportions (where $0 \le p[k] \le 1$ and $\sum p[k] = 1$), you can transform variables sampled from a uniform distribution using the following method:

1. Sample $n$ values $u_i$ from $Uni(0, 1)$.
2. Transform them using $x_i = -\log(u_i)$ to follow an exponential distribution $Exp(1)$.
3. Normalize them: $p_i = x_i / \sum x_i$.

This results in a vector $p$ that follows a flat Dirichlet distribution $Dir(1, ..., 1)$.

import numpy as np
import optuna

def objective(trial):
    n = 5
    x = []
    for i in range(n):
        # Sample from uniform and transform to exponential
        x.append(- np.log(trial.suggest_float(f"x_{i}", 0, 1)))

    # Normalize to get Dirichlet proportions
    p = [val / sum(x) for val in x]

    # Store proportions in user_attrs for retrieval
    for i in range(n):
        trial.set_user_attr(f"p_{i}", p[i])

    return 0

study = optuna.create_study(sampler=optuna.samplers.RandomSampler())
study.optimize(objective, n_trials=1000)

Optuna's optimization process is built around two core abstractions:

1. Study: An optimization session based on an objective function. The goal of a study is to find the optimal set of hyperparameter values by running multiple trials.
2. Trial: A single execution of the objective function. During a trial, Optuna suggests specific hyperparameter values which are then used to evaluate the objective function.

By using a define-by-run API, you can dynamically construct search spaces using standard Python control flow (like if statements and for loops) within your objective function.

import optuna
import sklearn.svm
import sklearn.ensemble
import sklearn.datasets
import sklearn.model_selection
import sklearn.metrics

# Define an objective function to be minimized.
def objective(trial):
    # Invoke suggest methods of a Trial object to generate hyperparameters.
    regressor_name = trial.suggest_categorical('classifier', ['SVR', 'RandomForest'])
    if regressor_name == 'SVR':
        svr_c = trial.suggest_float('svr_c', 1e-10, 1e10, log=True)
        regressor_obj = sklearn.svm.SVR(C=svr_c)
    else:
        rf_max_depth = trial.suggest_int('rf_max_depth', 2, 32)
        regressor_obj = sklearn.ensemble.RandomForestRegressor(max_depth=rf_max_depth)

    X, y = sklearn.datasets.fetch_california_housing(return_X_y=True)
    X_train, X_val, y_train, y_val = sklearn.model_selection.train_test_split(X, y, random_state=0)

    regressor_obj.fit(X_train, y_train)
    y_pred = regressor_obj.predict(X_val)

    error = sklearn.metrics.mean_squared_error(y_val, y_pred)

    return error  # An objective value linked with the Trial object.

study = optuna.create_study()  # Create a new study.
study.optimize(objective, n_trials=100)  # Invoke optimization of the objective function.

Optuna provides several sampling strategies via the optuna.samplers module. Choosing the right sampler depends on your parameter types (Float, Integer, Categorical), whether you are doing multi-objective or multivariate optimization, and your computational budget.

Key considerations:

• If you are unsure which sampler to use: Consider using AutoSampler from OptunaHub, which automatically selects a sampler during optimization.
• For standard optimization: TPESampler is a common choice for single-objective problems. RandomSampler is a baseline.
• For multi-objective optimization: Use NSGAIISampler or NSGAIIISampler.
• For exhaustive search: Use GridSampler or BruteForceSampler.
• For high-dimensional continuous spaces: CmaEsSampler or GPSampler may be appropriate.

Note on Reproducibility: Samplers use a seed argument for initialization. However, if n_jobs != 1 in study.optimize, samplers reseed to avoid duplicate parameters. This makes results difficult to reproduce in parallel settings. For distributed optimization, use seed=None or different seeds for each process.