Optuna hyperparameter search not reproducible with interrupted / resumed studies

Question:

For big ML models with many parameters, it is helpful if one can interrupt and resume the hyperparameter optimization search.
Optuna allows doing that with the RDB backend, which stores the study in a SQlite database (https://optuna.readthedocs.io/en/stable/tutorial/20_recipes/001_rdb.html#sphx-glr-tutorial-20-recipes-001-rdb-py).

However, when interrupting and resuming a study, the results are not the same as that of an uninterrupted study.

Expect: For a fixed seed, the results from an optimization run with n_trials = x are identical to a study with n_trials = x/5, that is resumed 5 times and a study, that is interrupted with KeyboardInterrupt 5 times and resumed 5 times until n_trials = x.

Actual: The results are equal up to the point of the first interruption. From then on, they differ.

The figures show the optimization history of all trials in a study. The left-most figure (A) shows the uninterrupted run, the center one shows a run interrupted by keyboard (B), the right-most figure shows the run interrupted by n_iter (C). In B and C, the red dotted line shows the point where the first study was first interrupted. Left of the line, the results are equal to the uninterrupted study, to the right they differ.

Is it possible to interrupt and resume a study, so that another study with the same seed that has not been interrupted generates exactly the same result?
(Obviously assuming that the objective function behaves in a non-deterministic way.)

Minimal working example to reproduce:

import optuna
import logging
import sys
import numpy as np

def objective(trial):
    x = trial.suggest_float("x", -10, 10)
    return (x - 4) ** 2

def set_study(db_name, 
                study_name, 
                seed, 
                direction="minimize"):
    '''
    Creates a new study in a sqlite database located in results/ .
    The study can be resumed after keyboard interrupt by simple creating it
    using the same command used for the initial creation.
    '''

    # Add stream handler of stdout to show the messages
    optuna.logging.get_logger("optuna").addHandler(logging.StreamHandler(sys.stdout))
    sampler = optuna.samplers.TPESampler(seed = seed, n_startup_trials = 0)
    storage_name = f"sqlite:///{db_name}.db"
    storage = optuna.storages.RDBStorage(storage_name, heartbeat_interval=1)

    study = optuna.create_study(storage=storage, 
                                study_name=study_name, 
                                sampler=sampler, 
                                direction=direction, 
                                load_if_exists=True)
    return study


study = set_study('optuna_test', 'optuna_test_study', 1)

try:
    # Press CTRL+C to stop the optimization.
    study.optimize(objective, n_trials=100)  
except KeyboardInterrupt:
    pass


df = study.trials_dataframe(attrs=("number", "value", "params", "state"))

print(df)

print("Best params: ", study.best_params)
print("Best value: ", study.best_value)
print("Best Trial: ", study.best_trial)
# print("Trials: ", study.trials)


fig = optuna.visualization.plot_optimization_history(study)
fig.show()
Asked By: Maximilian Wirth

||

Answers:

Found whats causing the problem: The random number generator in the sampler is initialized using the seed, but of course it returns a different number if the study is interrupted and resumed (it is then reinitialised)
This is especially bad using random search with fixed seed, as the search then basically starts from new.

If one really needs reproducible runs, one can simply extract the rng into a binary file after a run or keyboard interrupt, and resume by overwriting the newly generated rng of the sampler with the saved one.

Answered By: Maximilian Wirth