Optuna Artifacts チュートリアル

Optuna の artifact モジュールは、ファイルなどの形式でトライアルごとに比較的大きな属性を保存するためのモジュールです。Optuna v3.3 から導入されたこのモジュールは、大規模モデルのスナップショットを利用したハイパーパラメータチューニング、大量の化学構造の最適化、画像や音声を用いた人間参加型の最適化など、幅広い用途に活用できます。Optuna の artifact モジュールを使用することで、データベースに保存するには大きすぎるデータも扱うことが可能です。さらに、optuna-dashboard と連携することで、保存したアーティファクトをウェブ UI 上で自動的に可視化でき、実験管理の手間を大幅に削減できます。

概要

• artifact モジュールは、トライアルに関連する大規模なデータを簡単に保存・利用するための機能を提供します。

• optuna-dashboard のウェブページにアクセスするだけで、保存したアーティファクトを可視化でき、ダウンロードも容易に行えます。

• artifact モジュールの抽象化により、バックエンド（ファイルシステム、AWS S3 など）を簡単に切り替えることが可能です。

• artifact モジュールは Optuna と密接に連携しているため、実験管理を Optuna エコシステムだけで完結させることができ、コードベースを簡素化できます。

概念

Fig 1. 「artifact」の概念図

「artifact」は Optuna のトライアルに関連付けられます。Optuna では、目的関数を順次評価して最大（または最小）値を探索します。このように順次実行される目的関数の評価をトライアルと呼びます。通常、トライアルとその関連属性はストレージオブジェクトを介してファイルや RDB などに保存されます。実験管理のために、各トライアルごとに optuna.trial.Trial.user_attrs を保存・利用することも可能です。ただし、これらの属性は整数や短い文字列などの小規模なデータを想定して設計されており、大規模なデータの保存には適していません。Optuna の artifact モジュールを使用すれば、モデルのスナップショット、化学構造、画像・音声データなど、各トライアルに関連する大規模なデータを保存できます。

また、このチュートリアルでは触れませんが、トライアルだけでなくスタディに関連するアーティファクトも管理可能です。興味のある方は、公式ドキュメント を参照してください。

アーティファクトが有用な場面

アーティファクトは、各トライアルごとに RDB に保存するには大きすぎるデータを保存したい場合に有用です。例えば、以下のような場面で artifact モジュールが役立ちます：

• 機械学習モデルのスナップショット保存：LLM のような大規模機械学習モデルのハイパーパラメータチューニングを行う場合、モデルは非常に大きく、学習の各ラウンド（Optuna ではトライアルに相当）に時間がかかります。データセンターの停電やスケジューラによる計算ジョブの強制終了など、予期せぬ事態に備えて、各トライアルの途中でモデルのスナップショットを保存したい場合があります。これらのスナップショットは通常大きく、RDB に保存するよりファイルとして保存する方が適しています。このような場合に artifact モジュールが有用です。

• 化学構造の最適化：ブラックボックス最適化問題として安定な化学構造の探索を行う場合、1 つの化学構造を評価することが Optuna の 1 トライアルに相当します。これらの化学構造データは複雑で大規模であるため、RDB に保存するのは適切ではありません。特定のファイル形式で保存することが考えられますが、その場合に artifact モジュールが有用です。

• 画像生成モデルの Human-in-the-loop 最適化：画像を生成する generative model のプロンプトを Optuna で最適化する場合、Optuna でプロンプトをサンプリングし、generative model で画像を生成し、Human-in-the-loop 最適化プロセスで人間が画像を評価するとします。生成される画像は大容量データであるため、RDB に保存するのは適切ではなく、そのような場合に artifact モジュールが適しています。

トライアルとアーティファクトの記録方法

これまで説明したように、artifact モジュールは各トライアルごとに大容量データを保存したい場合に有用です。本節では、以下の 2 つのシナリオにおける artifact の動作を説明します：1) SQLite + local file system ベースの artifact バックエンドを使用する場合（最適化サイクル全体をローカルで完了する場合に適しています）、2) MySQL + AWS S3 ベースの artifact バックエンドを使用する場合（データをリモートに保存したい場合に適しています）。

シナリオ 1: SQLite + file system ベースの artifact ストア

Fig 2. SQLite + file system ベースの artifact ストア

まず、最適化をローカルで完了するシンプルなケースについて説明します。

通常、Optuna の最適化履歴はストレージオブジェクトを介して何らかのデータベースに永続化されます。ここでは、軽量な RDB 管理システムである SQLite をバックエンドとして使用する場合を考えます。SQLite では、データは単一のファイル（例：./example.db）に保存されます。最適化履歴には、各トライアルでサンプリングされたパラメータ、そのパラメータの評価値、各トライアルの開始/終了時刻などが含まれます。このファイルは SQLite 形式であり、大容量データの保存には適していません。大容量データのエントリを書き込むとパフォーマンスが低下する可能性があります。なお、SQLite は分散並列最適化には適していません。そのような場合は、後述する MySQL を使用するか、JournalStorage を使用してください。

そこで、artifact モジュールを使用して大容量データを別の形式で保存します。各トライアルで生成されたデータを保存したい場合（例：画像データの場合は png 形式）、その保存先はローカルファイルシステム上の任意のディレクトリ（例：./artifacts ディレクトリ）を指定できます。目的関数を定義する際には、artifact モジュールを使用してデータの保存と参照を行うだけで済みます。

上記のケースの簡単な擬似コードは以下のようになります。

import os

import optuna
from optuna.artifacts import FileSystemArtifactStore
from optuna.artifacts import upload_artifact
from optuna.artifacts import download_artifact


base_path = "./artifacts"
os.makedirs(base_path, exist_ok=True)
artifact_store = FileSystemArtifactStore(base_path=base_path)


def objective(trial: optuna.Trial) -> float:
    ... = trial.suggest_float("x", -10, 10)

    # アーティファクトの作成と書き込み
    file_path = generate_example(...)  # この関数は何らかのファイルを生成します
    artifact_id = upload_artifact(
        artifact_store=artifact_store,
        file_path=file_path,
        study_or_trial=trial,
    )  # 戻り値はアーティファクト ID です
    trial.set_user_attr(
        "artifact_id", artifact_id
    )  # 後で参照できるように RDB に ID を保存します

    return ...


study = optuna.create_study(study_name="test_study", storage="sqlite:///example.db")
study.optimize(objective, n_trials=100)

# 最適トライアルに関連するアーティファクトをダウンロード
best_artifact_id = study.best_trial.user_attrs.get("artifact_id")
download_file_path = ...  # ダウンロードしたアーティファクトを保存するパスを設定
download_artifact(
    artifact_store=artifact_store, file_path=download_file_path, artifact_id=best_artifact_id
)
with open(download_file_path, "rb") as f:
    content = f.read().decode("utf-8")
print(content)

シナリオ 2: リモート MySQL RDB サーバー + AWS S3 artifact ストア

Fig 3. リモート MySQL RDB サーバー + AWS S3 artifact ストア

次に、データをリモートで読み書きするケースについて説明します。

最適化の規模が大きくなると、すべてをローカルで計算することが難しくなります。Optuna のストレージ オブジェクトは URL を指定することでデータをリモートに永続化でき、分散最適化が可能になります。ここでは MySQL をリモートのリレーショナル データベース サーバーとして使用します。MySQL はオープンソースのリレーショナル データベース管理システムで、さまざまな用途で 使用されている有名なソフトウェアです。MySQL を Optuna で使用する場合、チュートリアル が参考になります。ただし、MySQL のようなリレーショナル データベースで大量のデータを読み書きするのは適切ではありません。

Optuna では、このようなデータを各トライアルごとに読み書きしたい場合、artifact モジュールを使用することが一般的です。シナリオ 1 とは異なり、計算ノード間で最適化を分散するため、ローカル ファイル システムベースのバックエンドは使用できません。代わりに、 オンライン クラウド ストレージ サービスである AWS S3 と、Python から操作するためのフレームワークである Boto3 を使用します。 v3.3 以降、Optuna にはこの Boto3 バックエンドを備えた組み込みの artifact store があります。

データの流れは図 3 のようになります。各トライアルで計算された情報（アーティファクト情報を除く最適化履歴）は MySQL サーバーに 書き込まれます。一方、アーティファクト情報は AWS S3 に書き込まれます。分散最適化を行うすべてのワーカーは、それぞれに対して 並列に読み書きでき、レースコンディションなどの問題は Optuna のストレージ モジュールと artifact モジュールによって自動的に 解決されます。その結果、実際のデータの場所がアーティファクト情報と非アーティファクト情報で異なりますが（前者は AWS S3、後者は MySQL RDB）、ユーザーは透過的にデータを読み書きできます。上記のプロセスを簡単な擬似コードに翻訳すると次のようになります。

import os

import boto3
from botocore.config import Config
import optuna
from optuna.artifact import upload_artifact
from optuna.artifact import download_artifact
from optuna.artifact.boto3 import Boto3ArtifactStore


artifact_store = Boto3ArtifactStore(
    client=boto3.client(
        "s3",
        aws_access_key_id=os.environ[
            "AWS_ACCESS_KEY_ID"
        ],  # これらの環境変数が正しく設定されていると仮定します。以下も同様です。
        aws_secret_access_key=os.environ["AWS_SECRET_ACCESS_KEY"],
        endpoint_url=os.environ["S3_ENDPOINT"],
        config=Config(connect_timeout=30, read_timeout=30),
    ),
    bucket_name="example_bucket",
)


def objective(trial: optuna.Trial) -> float:
    ... = trial.suggest_float("x", -10, 10)

    # アーティファクトの作成と書き込み。
    file_path = generate_example(...)  # この関数は何らかのファイルを返す。
    artifact_id = upload_artifact(
        artifact_store=artifact_store,
        file_path=file_path,
        study_or_trial=trial,
    )  # 戻り値はアーティファクト ID です。
    trial.set_user_attr(
        "artifact_id", artifact_id
    )  # 後で参照できるように ID を RDB に保存します。

    return ...


study = optuna.create_study(
    study_name="test_study",
    storage="mysql://USER:PASS@localhost:3306/test",  # 適切な URL を設定してください。
)
study.optimize(objective, n_trials=100)

# 最良のトライアルに関連付けられたアーティファクトのダウンロード。
best_artifact_id = study.best_trial.user_attrs.get("artifact_id")
download_file_path = ...  # ダウンロードしたアーティファクトを保存するパスを設定します。
download_artifact(
    artifact_store=artifact_store, file_path=download_file_path, artifact_id=best_artifact_id
)
with open(download_file_path, "rb") as f:
    content = f.read().decode("utf-8")
print(content)

例: 化学構造の最適化

このセクションでは、artifact モジュールを活用して化学構造を最適化する Optuna の例を紹介します。比較的小規模な構造を対象としますが、 複雑な構造の場合も同様のアプローチが適用できます。

特定の分子が他の物質に吸着する過程を考えます。この過程では、吸着分子の吸着位置によって吸着反応の容易さが変化します。吸着反応の容易さは、 吸着エネルギー（吸着前後の系のエネルギー差）で評価できます。吸着分子の位置関係を入力とし、吸着エネルギーを出力する目的関数の最小化問題として 定式化することで、この問題をブラックボックス最適化問題として解くことができます。

まず、必要なモジュールをインポートし、いくつかの補助関数を定義します。化学構造を扱う ASE ライブラリを Optuna に加えてインストールする必要が ありますので、pip install ase でインストールしてください。

from __future__ import annotations

import io
import logging
import os
import sys
import tempfile

from ase import Atoms
from ase.build import bulk, fcc111, molecule, add_adsorbate
from ase.calculators.emt import EMT
from ase.io import write, read
from ase.optimize import LBFGS
import numpy as np
from optuna.artifacts import FileSystemArtifactStore
from optuna.artifacts import upload_artifact
from optuna.artifacts import download_artifact
from optuna.logging import get_logger
from optuna import create_study
from optuna import Trial


# メッセージを表示するため stdout のストリームハンドラを追加
get_logger("optuna").addHandler(logging.StreamHandler(sys.stdout))


def get_opt_energy(atoms: Atoms, fmax: float = 0.001) -> float:
    calculator = EMT()
    atoms.set_calculator(calculator)
    opt = LBFGS(atoms, logfile=None)
    opt.run(fmax=fmax)
    return atoms.get_total_energy()


def create_slab() -> tuple[Atoms, float]:
    calculator = EMT()
    bulk_atoms = bulk("Pt", cubic=True)
    bulk_atoms.calc = calculator

    a = np.mean(np.diag(bulk_atoms.cell))
    slab = fcc111("Pt", a=a, size=(4, 4, 4), vacuum=40.0, periodic=True)
    slab.calc = calculator
    E_slab = get_opt_energy(slab, fmax=1e-4)
    return slab, E_slab


def create_mol() -> tuple[Atoms, float]:
    calculator = EMT()
    mol = molecule("CO")
    mol.calc = calculator
    E_mol = get_opt_energy(mol, fmax=1e-4)
    return mol, E_mol


def atoms_to_json(atoms: Atoms) -> str:
    f = io.StringIO()
    write(f, atoms, format="json")
    return f.getvalue()


def json_to_atoms(atoms_str: str) -> Atoms:
    return read(io.StringIO(atoms_str), format="json")


def file_to_atoms(file_path: str) -> Atoms:
    return read(file_path, format="json")

各関数の説明

• get_opt_energy: 化学構造を受け取り、局所的に安定な構造へ遷移させた後のエネルギーを計算して返す

• create_slab: 吸着対象物質のスラブ構造を作成する

• create_mol: 吸着対象分子を作成する

• atoms_to_json: 化学構造をJSON形式の文字列に変換する

• json_to_atoms: JSON形式の文字列を化学構造に変換する

• file_to_atoms: ファイルからJSON形式の文字列を読み込み、化学構造に変換する

これらの関数を用いて、Optunaで吸着構造を探索するコードは以下の通りである。目的関数はアーティファクトストアを保持するため クラス Objective として定義されている。 __call__ メソッドでは、吸着対象物質（slab）と吸着分子（mol）を取得し、 Optunaで位置関係をサンプリングした後（複数の trial.suggest_xxx メソッドを使用）、``add_adsorbate``関数で吸着反応を 実行し、局所的に安定な構造へ遷移させる。その後、構造をアーティファクトストアに保存し、吸着エネルギーを返却する。

main 関数では Study の作成と最適化の実行コードを記述している。 Study 作成時にはSQLiteをストレージとして指定し、 アーティファクトストアのバックエンドにはローカルファイルシステムを使用している。つまり、前節で説明したシナリオ1に 該当する。100回の最適化試行後、最良の試行の情報を表示し、最終的に化学構造を best_atoms.png として保存する。 得られた best_atoms.png は図4に示す。

class Objective:
    def __init__(self, artifact_store: FileSystemArtifactStore) -> None:
        self._artifact_store = artifact_store

    def __call__(self, trial: Trial) -> float:
        slab = json_to_atoms(trial.study.user_attrs["slab"])
        E_slab = trial.study.user_attrs["E_slab"]

        mol = json_to_atoms(trial.study.user_attrs["mol"])
        E_mol = trial.study.user_attrs["E_mol"]

        phi = 180.0 * trial.suggest_float("phi", -1, 1)
        theta = np.arccos(trial.suggest_float("theta", -1, 1)) * 180.0 / np.pi
        psi = 180 * trial.suggest_float("psi", -1, 1)
        x_pos = trial.suggest_float("x_pos", 0, 0.5)
        y_pos = trial.suggest_float("y_pos", 0, 0.5)
        z_hig = trial.suggest_float("z_hig", 1, 5)
        xy_position = np.matmul([x_pos, y_pos, 0], slab.cell)[:2]
        mol.euler_rotate(phi=phi, theta=theta, psi=psi)

        add_adsorbate(slab, mol, z_hig, xy_position)
        E_slab_mol = get_opt_energy(slab, fmax=1e-2)

        write(f"./tmp/{trial.number}.json", slab, format="json")
        artifact_id = upload_artifact(
            artifact_store=self._artifact_store,
            file_path=f"./tmp/{trial.number}.json",
            study_or_trial=trial,
        )
        trial.set_user_attr("structure", artifact_id)

        return E_slab_mol - E_slab - E_mol


def main():
    study = create_study(
        study_name="test_study",
        storage="sqlite:///example.db",
        load_if_exists=True,
    )

    slab, E_slab = create_slab()
    study.set_user_attr("slab", atoms_to_json(slab))
    study.set_user_attr("E_slab", E_slab)

    mol, E_mol = create_mol()
    study.set_user_attr("mol", atoms_to_json(mol))
    study.set_user_attr("E_mol", E_mol)

    os.makedirs("./tmp", exist_ok=True)

    base_path = "./artifacts"
    os.makedirs(base_path, exist_ok=True)
    artifact_store = FileSystemArtifactStore(base_path=base_path)
    study.optimize(Objective(artifact_store), n_trials=3)
    print(
        f"Best trial is #{study.best_trial.number}\n"
        f"    Its adsorption energy is {study.best_value}\n"
        f"    Its adsorption position is\n"
        f"        phi  : {study.best_params['phi']}\n"
        f"        theta: {study.best_params['theta']}\n"
        f"        psi. : {study.best_params['psi']}\n"
        f"        x_pos: {study.best_params['x_pos']}\n"
        f"        y_pos: {study.best_params['y_pos']}\n"
        f"        z_hig: {study.best_params['z_hig']}"
    )

    best_artifact_id = study.best_trial.user_attrs["structure"]

    with tempfile.TemporaryDirectory() as tmpdir_name:
        download_file_path = os.path.join(tmpdir_name, f"{best_artifact_id}.json")
        download_artifact(
            artifact_store=artifact_store,
            file_path=download_file_path,
            artifact_id=best_artifact_id,
        )

        best_atoms = file_to_atoms(download_file_path)
        print(best_atoms)
        write("best_atoms.png", best_atoms, rotation=("315x,0y,0z"))


if __name__ == "__main__":
    main()

A new study created in RDB with name: test_study
/mnt/nfs-mnj-hot-99-home/mshibata/sandbox/optuna-documentation-plamo-ja/optuna-doc-plamo-translation/tmp-optuna/tutorial/20_recipes/012_artifact_tutorial.py:243: FutureWarning:

Please use atoms.calc = calc

Trial 0 finished with value: -0.9835303337442225 and parameters: {'phi': -0.5875349078108321, 'theta': 0.8376155483350654, 'psi': -0.638161729674908, 'x_pos': 0.3013720554603374, 'y_pos': 0.43972722821182053, 'z_hig': 4.258735326803025}. Best is trial 0 with value: -0.9835303337442225.
Trial 1 finished with value: -0.9841892572010886 and parameters: {'phi': -0.1652613809566772, 'theta': -0.21547558351887375, 'psi': 0.30174478781642766, 'x_pos': 0.4981508298400785, 'y_pos': 0.21405419407519066, 'z_hig': 2.0948233524707973}. Best is trial 1 with value: -0.9841892572010886.
Trial 2 finished with value: -0.9815866581017181 and parameters: {'phi': -0.8129684859646717, 'theta': -0.02940568424469392, 'psi': 0.1930644338597458, 'x_pos': 0.2588995996499024, 'y_pos': 0.3282351581854703, 'z_hig': 2.444184250987972}. Best is trial 1 with value: -0.9841892572010886.
Best trial is #1
    Its adsorption energy is -0.9841892572010886
    Its adsorption position is
        phi  : -0.1652613809566772
        theta: -0.21547558351887375
        psi. : 0.30174478781642766
        x_pos: 0.4981508298400785
        y_pos: 0.21405419407519066
        z_hig: 2.0948233524707973
Atoms(symbols='COPt64', pbc=True, cell=[[11.087434329005065, 0.0, 0.0], [5.5437171645025325, 9.601999791710057, 0.0], [0.0, 0.0, 86.78963916567001]], tags=..., calculator=SinglePointCalculator(...))

Fig 4. 上記コードで得られた化学構造

上記のように、化学構造の最適化を行う際にartifactモジュールを使用すると便利です。 構造が小規模な場合や試行回数が少ない場合は、文字列に変換してRDBに直接保存しても問題ありません。 しかし、複雑な構造を扱う場合や大規模な探索を行う場合は、RDBに負荷をかけないように 外部ファイルシステムやAWS S3などに保存することをお勧めします。

まとめ

artifactモジュールは、各試行ごとに比較的大きなデータを保存したい場合に便利な機能です。 機械学習モデルのスナップショット保存、化学構造の最適化、画像や音声の人間参加型最適化など、 様々な用途に活用できます。Optunaを使ったブラックボックス最適化の強力な味方となるでしょう。 また、Optunaのメンテナが気づいていない活用方法があれば、GitHubのディスカッションで ぜひ教えてください。Optunaを使った最適化ライフを楽しんでください！