optuna-memorydump で InMemoryStorage の結果を定期的にSQLite3に書き出す。

CPUたくさん積んだマシン1台で、Optunaを回すときにstorage backendの選択は悩ましい点があります。 結果をDashboardやJupyterで詳しく確認したり、あとから最適化を再開するためにRDB storage backendを使って永続化しておきたいのですが、マシン1台しか使わないときにわざわざMySQLとかPostgreSQL立てるのは結構面倒なので、自然とSQLite3を選びたくなります。

ただSQLite3はマルチスレッドやマルチプロセスによる最適化時に、あまりおすすめできません。 そもそもLock errorを完全に防ぐことが難しいという話もありますが、 enqueue_trial など一部の機能は 行ロックのできないSQLite3 では並列最適化時に正しく動作しません。またRDB storage backendは性能面の課題もあり、RandomSamplerやCmaEsSamplerで実行時間の短い目的関数を数千回評価したいときにはほぼ確実にRDB storage backedがボトルネックになるかと思います。

そこで今回は高速に動作するIn-memory storage backendで最適化を回しつつ、評価回数が500回や1000回など指定したintervalに達したときにSQLite3に書き出すツールを用意してみました。冪等性を保証していて何回書き出しても差分だけを適用します。思いつきでサクッと作ってみたツールですが、思ったよりも便利そうなのでよければ使ってみてください。

optuna-memorydump の使い方

GitHubリポジトリはこちらです。まだPyPIには上げていないのでGitHubからインストールしてください。

github.com

$ pip install git+https://github.com/c-bata/optuna-memorydump.git

使い方はこんな感じでCallback関数としてOptunaに登録します。

import optuna
from optuna_memorydump import MemoryDumpCallback


def objective(trial):
    x1 = trial.suggest_uniform("x1", -100, 100)
    x2 = trial.suggest_uniform("x2", -100, 100)
    return 100 * (x2 - x1 ** 2) ** 2 + (x1 - 1) ** 2


if __name__ == "__main__":
    study = optuna.create_study(
        study_name="dumped", sampler=optuna.samplers.CmaEsSampler(),
    )
    study.optimize(
        objective, timeout=30, n_jobs=4, gc_after_trial=False,
        callbacks=[MemoryDumpCallback("sqlite:///db.sqlite3", interval=100)],
    )
    print("Best value: {} (params: {})\n".format(study.best_value, study.best_params))

intervalを100にすると、100回評価が終わるたびにcallbackが反応して指定したstorageに書き出します。intervalを指定できるようにしたのは、仮に1日中回し続けたいといったケースで12時間後にエラーがでて最初からやり直しとかを避けるためです。

$ python examples/rosenbrock.py
[I 2020-04-12 18:33:39,200] Finished trial#0 with value: ...
[I 2020-04-12 18:33:39,204] Finished trial#2 with value: ...
[I 2020-04-12 18:33:39,205] Finished trial#1 with value: ...
:
[I 2020-04-12 18:33:39,924] Finished trial#100 with value: ...
[I 2020-04-12 18:33:39,931] memorydump is triggered at trial 100 (thread=123145615126528).
[I 2020-04-12 23:40:20,013] memorydump of trial 100 is finished in 1.938s (thread=123145458638848).
:
[I 2020-04-12 18:33:39,924] Finished trial#200 with value: ...
[I 2020-04-12 23:40:21,505] memorydump of trial 200 is finished in 0.105s (thread=123145458638848).
[I 2020-04-12 23:40:24,061] memorydump is triggered at trial 300 (thread=123145425059840).
: