Optuna 工件教程

Optuna 的工件模块 (artifact module) 是一个模块，旨在按试验 (trial) 逐个保存相对较大的属性，例如以文件的形式。此模块从 Optuna v3.3 开始引入，应用范围广泛，例如利用大型模型快照进行超参数调优、优化大型化学结构，甚至利用图像或声音进行人机协作优化。使用 Optuna 的工件模块可以处理那些因过大而无法存储在数据库中的数据。此外，通过与 optuna-dashboard 集成，保存的工件可以自动通过 Web UI 可视化，这显著降低了实验管理的难度。

要点速览

• 工件模块提供了一种保存和使用与试验 (trial) 相关的大型数据的简单方法。

• 使用 optuna-dashboard 访问网页即可可视化保存的工件，下载也很方便。

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

• 由于工件模块与 Optuna 紧密关联，实验管理只需依靠 Optuna 生态系统即可完成，从而简化了代码库。

概念

图 1. “工件”的概念。

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

此外，虽然本教程并未涉及，但也可以管理不仅与试验相关联的工件，还与研究 (study) 相关联的工件。如果您对此感兴趣，请参考 官方文档。

工件适用的场景

当您希望为每个试验 (trial) 保存大到无法存储在 RDB 中的数据时，工件非常有用。例如，工件模块在以下情况中会很方便：

• 保存机器学习模型快照：假设您正在为大型机器学习模型（如 LLM）调优超参数。模型非常庞大，每一轮学习（对应于 Optuna 中的一次试验）都需要时间。为了应对训练过程中的意外事件（如数据中心停电或调度器抢占计算任务），您可能希望为每个试验保存训练过程中的模型快照。这些快照通常较大，与其存储在 RDB 中，不如保存为某种文件格式更合适。在这种情况下，工件模块就很有用了。

• 优化化学结构：假设您正在构建并探索将寻找稳定化学结构的问题作为黑盒优化问题。评估一个化学结构对应于 Optuna 中的一次试验 (trial)，该化学结构复杂且庞大。不适合将此类化学结构数据存储在 RDB 中。可以考虑将化学结构数据保存为特定的文件格式，在这种情况下，工件模块就很有用了。

• 利用人机协作优化图像：假设您正在优化用于生成图像的生成模型提示。您使用 Optuna 采样提示，使用生成模型输出图像，然后让人类对图像进行评分以进行人机协作优化过程。由于输出图像是大型数据，不适合使用 RDB 存储它们，在这种情况下，使用工件模块非常合适。

试验 (Trial) 和工件的记录方式

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

场景 1：SQLite + 文件系统作为工件存储

图 2. SQLite + 文件系统作为工件存储。

首先，我们解释优化在本地完成的简单情况。

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

因此，让我们使用工件模块以不同的格式保存大型数据。假设数据是为每个试验 (trial) 生成的，并且您希望以某种格式（例如，如果是图像则为 png 格式）保存它。保存工件的具体目的地可以是本地文件系统上的任何目录（例如，./artifacts 目录）。在定义目标函数时，您只需使用工件模块保存和引用数据即可。

上述情况的简单伪代码如下所示：

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 作为工件存储

图 3. 远程 MySQL RDB 服务器 + AWS S3 作为工件存储。

接下来，我们解释数据在远程读取和写入的情况。

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

在 Optuna 中，当您希望为每个试验 (trial) 读写此类数据时，通常会使用工件模块。与场景 1 不同，我们将优化分布在多个计算节点上，因此基于本地文件系统的后端将无法工作。相反，我们将使用在线云存储服务 AWS S3，以及从 Python 与其交互的框架 Boto3。从 v3.3 开始，Optuna 内置了使用此 Boto3 后端的工件存储。

数据流如图 3 所示。每个试验 (trial) 中计算的信息，即优化历史（不包括工件信息），会写入 MySQL 服务器。另一方面，工件信息会写入 AWS S3。所有进行分布式优化的 worker 都可以并行地向各自的存储读写数据，竞争条件等问题由 Optuna 的存储模块和工件模块自动解决。因此，尽管工件信息和非工件信息的实际数据位置不同（前者在 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)

示例：化学结构的优化

在本节中，我们介绍一个利用工件模块使用 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 搜索吸附结构的代码如下所示。目标函数被定义为一个 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
/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.00021064099479506382 and parameters: {'phi': 0.985940775286323, 'theta': -0.49787283680664096, 'psi': -0.22818892236873456, 'x_pos': 0.19526958382976595, 'y_pos': 0.4328956559427395, 'z_hig': 4.97356841142053}. Best is trial 0 with value: -0.00021064099479506382.
Trial 1 finished with value: -0.9850453891170909 and parameters: {'phi': 0.13798061637862102, 'theta': 0.4413034282816244, 'psi': 0.18410830756761332, 'x_pos': 0.2965422097428139, 'y_pos': 0.12274714554705346, 'z_hig': 4.242065658661653}. Best is trial 1 with value: -0.9850453891170909.
Trial 2 finished with value: -0.9841725584949605 and parameters: {'phi': 0.05174409930375168, 'theta': -0.22528620942198785, 'psi': -0.19963688187623996, 'x_pos': 0.01021822118390664, 'y_pos': 0.05508777488438721, 'z_hig': 2.0948722649337594}. Best is trial 1 with value: -0.9850453891170909.
Best trial is #1
    Its adsorption energy is -0.9850453891170909
    Its adsorption position is
        phi  : 0.13798061637862102
        theta: 0.4413034282816244
        psi. : 0.18410830756761332
        x_pos: 0.2965422097428139
        y_pos: 0.12274714554705346
        z_hig: 4.242065658661653
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 进行化学结构优化时，使用工件模块很方便。对于小型结构或较少试验次数的情况，将其转换为字符串并直接保存在 RDB 中也可以。但是，在处理复杂结构或进行大规模搜索时，最好将其保存在 RDB 外部以避免过载，例如外部文件系统或 AWS S3 中。

结论

当您希望为每个试验 (trial) 保存相对较大的数据时，工件模块是一个有用的特性。它可以用于各种目的，例如保存机器学习模型快照、优化化学结构以及利用图像和声音进行人机协作优化。它是使用 Optuna 进行黑盒优化的强大助手。此外，如果 Optuna 提交者们尚未注意到的使用方法，请在 GitHub discussions 上告知我们。祝您使用 Optuna 获得愉快的优化体验！