Optunaではハイパーパラメーターや評価値といった試行履歴を保存するストレージクラスがあり、OptunaのデフォルトのストレージクラスはInMemoryStorageです。しかしInMemoryStorageは試行履歴を永続化せず、分散最適化にも利用できません。そのためOptunaでは複数のストレージクラスが用意されています。

Optuna 4.0ではストレージクラスの一つであるJournalStorageおよびJournalFileBackendが正式にサポートされました。本ブログでは、これらの技術的なポイントに加えて活用事例や使用方法について紹介します。

JournalStorageとは

JournalStorageはOptunaのストレージクラスのひとつです。名前の由来はOptunaの操作ログを積み重ねていく記録方式であることによります。導入時の主要なモチベーションは、様々なバックエンド（データベースなど）を容易にOptunaストレージとして利用可能にしたいという点にありました。これを達成するために、JournalStorageクラスがOptunaのストレージとして振る舞う一方で、バックエンドへの読み書きを担うクラスが別個に設けられており、責任を分割する設計がなされています。そしてJournalStorageクラスは各バックエンド向けに用意されたクラスのオブジェクトを初期化時に受け取る仕様になっています。（図1）

図1: JournalStorageと関連クラスの関係性を矢印で示した図。上段にグレーで示したJournalStorageクラスは、中段に青で示したバックエンドのクラスオブジェクトを引数に取る。バックエンドのなかには、更にOptional引数でロック機構に関するクラス（緑色）を引数に取れるものもある。

JournalStorageはv3.1から試験的に導入されていました。v4.0ではこの正式サポートを開始しました。正式サポートに伴い、ログファイルの後方互換性が今後保証されます。その他にはクラス名やモジュールパスの再編成・安定性向上のための改善が入っています。

JournalStorageを使ったシンプルなコード例として、以下のように書くことができます。

import optuna
from optuna.storages import JournalStorage
from optuna.storages.journal import JournalFileBackend


def objective(trial):
    xs = [trial.suggest_float(f"x{n}", -1.0, 1.0) for n in range(3)]
    return sum(x ** 2 for x in xs)


storage = JournalStorage(JournalFileBackend("./optuna_journal_storage.log"))
study = optuna.create_study(storage=storage)
study.optimize(objective, n_trials=20)

前述の通りJournalStorageクラスはバックエンドにアクセスする機能を実装したクラスオブジェクト（このコードではJournalFileBackendクラス）を初期化時の引数に取ります。この引数には上記の例で使われているJournalFileBackend以外にも有名なNoSQLであるRedisをOptunaストレージとして利用するためのJournalRedisBackendクラスも指定可能です。ただし今回はexperimentalを外す対象をJournalFileBackendに絞っており、JournalRedisBackendはv4.0でも引き続きexperimentalとなります。

JournalStorageのさらなる詳細については過去のPFN Tech blogもご参照ください。

JournalFileBackendとは

JournalFileBackendは分散最適化に対応したストレージ機能を提供するクラスです。上記のサンプルコードのようにJournalStorageクラスの初期化時にオブジェクトを渡すことで利用できます。JournalFileBackendの最大の長所は、Network File System（NFS）経由の分散最適化が可能になる点にあります。これはNFSの仕様でatomicと規定されているシステムコールを使って排他制御を実装しているからです。ロックを取得する2種類の手法がそれぞれクラスとして実装されており、JournalFileBackendクラスの初期化時にオプショナル引数で渡すことで切り替え可能です。利用方法についてはドキュメントを、メカニズムのさらなる詳細についてはJournalStorageの詳細とともに過去のPFN Tech blogにて解説されていますので、ぜひご参照ください。

OptunaのStudyを保存する方法として、JournalStorageのほかにRDBStorageがあります。RDBStorageはMySQLなどのほかにSQLite3にも対応しています。SQLite3もJournalFileBackendと同様に単一ファイルをストレージとして用いることができ、手軽さが長所です。しかしSQLite3のファイルがNFS上に存在する場合、複数のノードないし複数のプロセスからの同時アクセスがうまく機能しないことが知られています。この点についてはSQLiteの公式FAQでも言及されています。以上の点から、NFS上のファイルをストレージとして用いてOptunaの分散最適化を行うには、JournalFileBackendが唯一の方法と言えます。

活用事例

社内活用事例を1つ紹介します。Optunaを使用したとある案件で、RDBStorage + MySQLを用いて大規模な分散最適化を実行した後でその最適化結果を解析する際に、DBサーバーへの負荷がボトルネックとなっていました。解析作業を効率化しDBサーバーへの負荷を軽減するために、MySQLに保存されていたStudyをoptuna.copy_study関数で別のStorageに移行することを検討しました。RedisやSQLite3などのStorageへの変換も検討しましたが、当該ワークロードにおいてはJournalStorage + JournalFileBackendへの変換が最も高速で適していることが分かりました。これにより、DBサーバーへの負荷を避けつつ、解析作業を迅速に行うことができました。

Optuna 4.0での安定化に伴う変更

Optuna 4.0ではAPIをより分かりやすいものにするため、クラス名やモジュールパスを再編成しました。上記の使用例のコードはv4.0の仕様に準拠したコードとなっています。（抜粋して再掲：）

from optuna.storages import JournalStorage
from optuna.storages.journal import JournalFileBackend
storage = JournalStorage(JournalFileBackend("./optuna_journal_storage.log"))

v3系の記法で書かれたコードにつきましては、当面そのままでも後方互換性を保って動作しますが非推奨となります。新しいモジュールパスの詳細についてはドキュメントおよびマイグレーションガイドをご覧ください。

v3.1の提供開始時に作成されたログファイルはv4.0以降もそのまま利用可能です！　上記コード例と同様にJournalFileBackendクラスからファイルを読み込み使用してください。