Optuna Artifacts 教程

Optuna 的 artifact 模块是一个旨在按 trial 保存相对较大的属性的模块，形式为文件等。该模块从 Optuna v3.3 引入，应用广泛，例如利用大型模型快照进行超参数调优、优化海量化学结构，甚至利用图像或声音进行人机协同优化。使用 Optuna 的 artifact 模块可以处理数据库无法存储的大型数据。此外，通过与 optuna-dashboard 集成，保存的 artifacts 可以通过 Web UI 自动可视化，大大减轻了实验管理的工作量。

摘要

• Artifact 模块提供了一种简单的方式来保存和使用与 trial 相关的大型数据。

• 通过访问 optuna-dashboard 的网页即可可视化保存的 artifacts，下载也非常方便。

• 得益于 artifact 模块的抽象，后端（文件系统、AWS S3）可以轻松切换。

• 由于 artifact 模块与 Optuna 紧密集成，实验管理可以通过 Optuna 生态系统单独完成，从而简化了代码库。

概念

图 1. “Artifact” 的概念。

“Artifact” 与 Optuna trial 相关联。在 Optuna 中，目标函数被顺序评估以搜索最大值（或最小值）。顺序重复的目标函数的每次评估称为一个 trial。通常，trial 及其关联的属性通过存储对象保存到文件或 RDB 等。对于实验管理，您还可以为每个 trial 保存和使用 optuna.trial.Trial.user_attrs。然而，这些属性假定为整数、短字符串或其他小数据，不适合存储大型数据。通过 Optuna 的 artifact 模块，用户可以为每个 trial 保存大型数据（如模型快照、化学结构、图像和音频数据等）。

另外，虽然本教程未涉及，但也可以管理不仅与 trials 而且与 studies 关联的 artifacts。如果您感兴趣，请参考 官方文档。

Artifacts 有用的场景

当您想为每个 trial 保存 RDB 无法存储的大型数据时，Artifacts 非常有用。例如，在以下情况下，artifact 模块将非常方便：

• 保存机器学习模型快照：假设您正在对 LLM 等大型机器学习模型进行超参数调优。模型非常大，每次学习（对应 Optuna 中的一个 trial）都需要花费时间。为了应对训练过程中可能发生的意外情况（如数据中心停电或调度程序抢占计算作业），您可能希望为每个 trial 在训练过程中保存模型快照。这些快照通常很大，更适合保存为某种文件而不是存储在 RDB 中。在这种情况下，artifact 模块很有用。

• 优化化学结构：假设您正在将寻找稳定化学结构的问题制定为黑盒优化问题进行探索。评估一个化学结构对应 Optuna 中的一个 trial，而该化学结构本身是一个复杂且庞大的结构。将此类化学结构数据存储在 RDB 中是不合适的。可以设想将化学结构数据以特定的文件格式保存，在这种情况下，artifact 模块很有用。

• 图像的人机协同优化：假设您正在优化生成式模型输出图像的提示。您使用 Optuna 采样提示，使用生成式模型输出图像，并让人类对图像进行评分，以进行人机协同优化。由于输出图像是大型数据，因此不适合使用 RDB 存储，在这种情况下，使用 artifact 模块非常合适。

Trial 和 Artifacts 的记录方式

如前所述，当您想为每个 trial 保存大型数据时，artifact 模块很有用。在本节中，我们将解释 artifacts 在以下两种场景中的工作方式：第一种是使用基于 SQLite + 本地文件系统的 artifact 后端（适用于整个优化周期在本地完成的情况），第二种是使用基于 MySQL + AWS S3 的 artifact 后端（适用于希望将数据保留在远程位置的情况）。

场景 1：基于 SQLite + 文件系统的 artifact 存储

图 2. 基于 SQLite + 文件系统的 artifact 存储。

首先，我们将解释一个在本地完成优化的简单场景。

通常，Optuna 的优化历史会通过存储对象持久化到某种数据库中。在这里，让我们考虑使用 SQLite（一个轻量级的 RDB 管理系统）作为后端的方法。使用 SQLite，数据存储在单个文件中（例如 ./example.db）。优化历史包括每个 trial 中采样的参数、这些参数的评估值、每个 trial 的开始和结束时间等。此文件是 SQLite 格式的，不适合存储大型数据。写入大型数据条目可能会导致性能下降。请注意，SQLite 不适合分布式并行优化。如果您想执行此操作，请使用 MySQL，如我们稍后将解释的，或者使用 JournalStorage。

因此，让我们使用 artifact 模块以不同的格式保存大型数据。假设数据是为每个 trial 生成的，并且您想以某种格式（例如，如果是图像，则为 png 格式）保存它。保存 artifacts 的特定目标可以是本地文件系统上的任何目录（例如 ./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)

    # Creating and writing an artifact.
    file_path = generate_example(...)  # This function returns some kind of file.
    artifact_id = upload_artifact(
        artifact_store=artifact_store,
        file_path=file_path,
        study_or_trial=trial,
    )  # The return value is the artifact ID.
    trial.set_user_attr(
        "artifact_id", artifact_id
    )  # Save the ID in RDB so that it can be referenced later.

    return ...


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

# Downloading artifacts associated with the best trial.
best_artifact_id = study.best_trial.user_attrs.get("artifact_id")
download_file_path = ...  # Set the path to save the downloaded artifact.
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 存储

图 3. 远程 MySQL RDB 服务器 + AWS S3 artifact 存储。

接下来，我们将解释远程读写数据的情况。

随着优化规模的增加，在本地完成所有计算变得困难。Optuna 的存储对象可以通过指定 URL 将数据持久化到远程，从而实现分布式优化。在这里，我们将使用 MySQL 作为远程关系数据库服务器。MySQL 是一个开源的关系数据库管理系统，是一个用于各种目的的知名软件。对于 Optuna 的 MySQL 使用，教程 是一个很好的参考。然而，在 MySQL 等关系数据库中读写大型数据也是不合适的。

在 Optuna 中，当您想为每个 trial 读写此类数据时，通常使用 artifact 模块。与场景 1 不同，我们将优化分布到计算节点上，因此基于本地文件系统的后端将不起作用。相反，我们将使用 AWS S3（一种在线云存储服务）和 Boto3（一个用于从 Python 与之交互的框架）。从 v3.3 开始，Optuna 已经内置了带有此 Boto3 后端的 artifact 存储。

数据流如图 3 所示。每个 trial 中计算的信息（对应于优化历史，不包括 artifact 信息）写入 MySQL 服务器。另一方面，artifact 信息写入 AWS S3。所有执行分布式优化的工作节点都可以并行地读写彼此，Optuna 的存储模块和 artifact 模块会自动解决诸如竞争条件之类的问题。因此，尽管实际数据位置在 artifact 信息和非 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"
        ],  # Assume that these environment variables are set up properly. The same applies below.
        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)

    # Creating and writing an artifact.
    file_path = generate_example(...)  # This function returns some kind of file.
    artifact_id = upload_artifact(
        artifact_store=artifact_store,
        file_path=file_path,
        study_or_trial=trial,
    )  # The return value is the artifact ID.
    trial.set_user_attr(
        "artifact_id", artifact_id
    )  # Save the ID in RDB so that it can be referenced later.

    return ...


study = optuna.create_study(
    study_name="test_study",
    storage="mysql://USER:PASS@localhost:3306/test",  # Set the appropriate URL.
)
study.optimize(objective, n_trials=100)

# Downloading artifacts associated with the best trial.
best_artifact_id = study.best_trial.user_attrs.get("artifact_id")
download_file_path = ...  # Set the path to save the downloaded artifact.
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 优化化学结构的示例。我们将针对相对较小的结构，但对于更复杂的结构，方法是相同的。

考虑一种特定分子吸附到另一种物质上的过程。在此过程中，吸附反应的难易程度取决于吸附分子相对于被吸附物质的位置。吸附反应的难易程度可以通过吸附能（吸附后系统能量与吸附前系统能量之差）来评估。通过将问题制定为目标函数的最小化问题，该目标函数以吸附分子的位置关系作为输入并输出吸附能，然后将该问题作为黑盒优化问题解决。

首先，让我们导入必要的模块并定义一些辅助函数。除了 Optuna 之外，您还需要安装 ASE 库来处理化学结构，所以请用 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


# Add stream handler of stdout to show the messages
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_to_atoms：将字符串转换为化学结构。

• file_to_atoms：从文件中读取字符串并将其转换为化学结构。

利用这些函数，使用 Optuna 搜索吸附结构的 Python 代码如下。目标函数被定义为类 Objective，以便传递 artifact 存储。在其 __call__ 方法中，它检索被吸附的物质（slab）和被吸附的分子（mol），然后使用 Optuna（多个 trial.suggest_xxx 方法）对它们的位置关系进行采样，接着使用 add_adsorbate 函数触发吸附反应，将其转换为局部稳定结构，然后将结构保存在 artifact 存储中并返回吸附能。

主函数 main 包含创建 Study 并执行优化的代码。创建 Study 时，使用 SQLite 指定存储，并为 artifact 存储使用本地文件系统作为后端。换句话说，它对应于上一节中解释的场景 1。在执行 100 次 trial 优化后，它显示最佳 trial 的信息，最后将化学结构保存为 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
/home/docs/checkouts/readthedocs.org/user_builds/optuna/checkouts/stable/tutorial/20_recipes/012_artifact_tutorial.py:280: FutureWarning:

Please use atoms.calc = calc

Trial 0 finished with value: -0.9842987953263753 and parameters: {'phi': 0.10240292156029374, 'theta': 0.6885339262754429, 'psi': -0.4639363045451581, 'x_pos': 0.15565936812518077, 'y_pos': 0.4559720207594061, 'z_hig': 4.070953283495404}. Best is trial 0 with value: -0.9842987953263753.
Trial 1 finished with value: -0.9848750394102952 and parameters: {'phi': 0.07268549042464523, 'theta': 0.7671912896795494, 'psi': -0.5736606120110788, 'x_pos': 0.1165935052770703, 'y_pos': 0.33895197714284886, 'z_hig': 4.151708702088685}. Best is trial 1 with value: -0.9848750394102952.
Trial 2 finished with value: -0.9842526642677241 and parameters: {'phi': 0.4984698564845864, 'theta': 0.3626526763696529, 'psi': -0.055747476925396056, 'x_pos': 0.21756966812868578, 'y_pos': 0.3665677703385695, 'z_hig': 2.3706620783631966}. Best is trial 1 with value: -0.9848750394102952.
Best trial is #1
    Its adsorption energy is -0.9848750394102952
    Its adsorption position is
        phi  : 0.07268549042464523
        theta: 0.7671912896795494
        psi. : -0.5736606120110788
        x_pos: 0.1165935052770703
        y_pos: 0.33895197714284886
        z_hig: 4.151708702088685
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(...))

图 4. 上述代码获得的化学结构。

如上所示，在使用 Optuna 进行化学结构优化时，使用 artifact 模块很方便。对于小型结构或 trial 数量较少的情况，将其转换为字符串并直接保存在 RDB 中也是可以的。但是，在处理复杂结构或进行大规模搜索时，最好将其保存在 RDB 之外，以避免使其过载，例如保存在外部文件系统或 AWS S3 中。

结论

当您想为每个 trial 保存相对较大的数据时，artifact 模块是一个非常有用的功能。它可以用于多种目的，例如保存机器学习模型快照、优化化学结构以及图像和声音的人机协同优化。它是 Optuna 黑盒优化功能的强大助手。另外，如果我们 Optuna 提交者尚未注意到其使用方式，请在 GitHub 讨论中告知我们。祝您在 Optuna 上优化生活愉快！