Pruning unpromising trials

Pruning often cuts tuning time several-fold for iterative models like neural nets and boosting.

Many trials are obviously bad after a few epochs. Pruning stops them early so the budget goes to promising ones.

Inside an iterative objective:

for epoch in range(n_epochs):
    score = train_one_epoch_and_validate()
    trial.report(score, step=epoch)
    if trial.should_prune():
        raise optuna.TrialPruned()

report records an intermediate value; should_prune asks the study's pruner whether this trial is worse than others at the same step. Raising optuna.TrialPruned marks the trial PRUNED.

Pruners, set with create_study(pruner=...):

MedianPruner (default): prune if the value is worse than the median of previous trials at the same step. Tune n_startup_trials (trials completed before pruning starts) and n_warmup_steps (steps before a trial can be pruned).
PercentilePruner: like median but with a chosen percentile.
SuccessiveHalvingPruner and HyperbandPruner: allocate resources in rungs, keeping only the top fraction at each; Hyperband is a strong general choice.
PatientPruner: wraps another pruner and waits a few steps of no improvement.
ThresholdPruner: prune when values cross fixed bounds, such as a diverging loss.
WilcoxonPruner: for objectives evaluated over many instances or folds rather than epochs.

Pruning is a no-op unless the objective calls report. Steps must be comparable across trials (same epoch numbering).

Example

import optuna
from sklearn.datasets import load_digits
from sklearn.linear_model import SGDClassifier
from sklearn.model_selection import train_test_split
optuna.logging.set_verbosity(optuna.logging.WARNING)
X, y = load_digits(return_X_y=True)
X_tr, X_va, y_tr, y_va = train_test_split(X / 16, y, random_state=0)

def objective(trial):
    alpha = trial.suggest_float("alpha", 1e-6, 1e0, log=True)
    clf = SGDClassifier(alpha=alpha, random_state=0)
    for step in range(20):                       # "epochs"
        clf.partial_fit(X_tr, y_tr, classes=list(range(10)))
        acc = clf.score(X_va, y_va)
        trial.report(acc, step)
        if trial.should_prune():
            raise optuna.TrialPruned()
    return acc

pruner = optuna.pruners.HyperbandPruner(min_resource=2, max_resource=20)
study = optuna.create_study(direction="maximize", pruner=pruner)
study.optimize(objective, n_trials=30)
n_pruned = sum(t.state.name == "PRUNED" for t in study.trials)
print(n_pruned, "of 30 trials pruned")   # typically more than half
print(round(study.best_value, 3))                          # about 0.95

Remember

Call trial.report(value, step) and check trial.should_prune() each step
Raise optuna.TrialPruned() to stop a trial
MedianPruner is the default; Hyperband is a strong general choice

Watch out for

Setting a pruner but never calling trial.report, so nothing is ever pruned.

Check yourself

You set pruner=MedianPruner() but no trial is ever pruned. What is the most likely cause?

Official docs: optuna.readthedocs.io/en/stable/reference/pruners.html

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks