PyCon APAC 2023 - ハイパーパラメータ最適化フレームワーク Optunaの最新機能紹介

generator.py

# Requirements:
# $ pip install optuna optuna-dashboard diffusers transformers accelerate scipy safetensors xformers botorch
#
# Process A: Launch a process that suggest new params and generate images.
# $ python generator.py
#
# Process B: Launch an Optuna Dashboard process.
# $ optuna-dashboard sqlite:///db.sqlite3 --artifact-dir ./artifact
import os
import time
import tempfile

import optuna
import torch
from optuna.trial import TrialState
from optuna.artifacts import upload_artifact, FileSystemArtifactStore
from optuna_dashboard.preferential import create_study
from optuna_dashboard.preferential.samplers.gp import PreferentialGPSampler
from optuna_dashboard import register_preference_feedback_component
from diffusers import StableDiffusionImg2ImgPipeline
from PIL import Image

device = "cuda:0"
torch_dtype = torch.float16

init_img = Image.open("./input.png")
init_img = init_img.resize((768, 768))
init_img = init_img.convert("RGB")
rng = torch.Generator(device)

# 画像(Artifact)のアップロード先の設定 - 今回はFileSystemArtifactStoreを用いて、 "artifact" ディレクトリ以下に保存。
base_path = os.path.join(os.path.dirname(__file__), "artifact")
os.makedirs(base_path, exist_ok=True)
artifact_store = FileSystemArtifactStore(base_path=base_path)


def suggest_and_generate_image(study: optuna.Study, pipe: StableDiffusionImg2ImgPipeline, tmpdir: str) -> None:
    # OptunaのTrialを生成し、パラメーターおよびプロンプトに含めるキーワードをサンプル
    trial = study.ask()
    guidance_scale = trial.suggest_float("guidance_scale", 1, 50)
    strength = trial.suggest_float("strength", 0.70, 1.0)
    num_inference_steps = trial.suggest_int("num_inference_steps", 5, 100)

    prompts = ["a mascot character with two eyes and a mouth"]
    prompts.append(trial.suggest_categorical("adjectives", ["cute", "funny", "memorable", "charming", "entertaining"]))
    prompts.append(trial.suggest_categorical("style", ["anime", "photo", "painting", ""]))
    prompts.append(trial.suggest_categorical("facial-expression", ["smiling", "frowning", "grinning", ""]))
    negative_prompt = []
    negative_prompt.append(trial.suggest_categorical("negative-quality", ["unnatural", "low-quality", ""]))
    negative_prompt.append(trial.suggest_categorical("negative-adjectives", ["dull", "boring", "unfriendly", ""]))

    # img2imgの実行
    image = pipe(
        ", ".join(prompts),
        negative_prompt=", ".join(negative_prompt),
        generator=rng,
        strength=strength,
        image=init_img,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        num_images_per_prompt=1,
    ).images[0]

    image_path = os.path.join(tmpdir, f"sample-{trial.number}.png")
    image.save(image_path)
    
    # Artifact Storeにアップロード
    artifact_id = upload_artifact(trial, image_path, artifact_store)
    trial.set_user_attr("image", artifact_id)


def main():
    study = create_study(
        study_name="preferential_diffusion",
        storage="sqlite:///db.sqlite3",
        sampler=PreferentialGPSampler(),
        load_if_exists=True,
        n_generate=3,
    )
    register_preference_feedback_component(study, component_type="artifact", artifact_key="image")

    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
        "stabilityai/stable-diffusion-2-1",
        torch_dtype=torch_dtype,
    )
    pipe = pipe.to(device)
    pipe.enable_xformers_memory_efficient_attention()

    # Start Preferential Optimization
    with tempfile.TemporaryDirectory() as tmpdir:
        while True:
            if not study.should_generate():
                time.sleep(0.1)  # Avoid busy-loop
                continue
            suggest_and_generate_image(study, pipe, tmpdir)


if __name__ == '__main__':
    main()

動作の様子

実行方法

推奨実行環境

本Exampleを動かす際は Linux / NVIDIA GPUの利用を推奨 しています。もしGPUがなく動かすのが大変という方は、Optuna Dashboard公式ドキュメントの下記チュートリアルをお試しください。こちらはCPUのみで簡単に動かすことができ、Human-in-the-loop最適化のためのコードの理解には十分です。

https://optuna-dashboard.readthedocs.io/en/latest/tutorials/preferential-optimization.html

• M1 Macbook Proでも動作しましたが、依存関係のインストールがやや大変な点やPyTorch MPS (Metal Performance Shaders) バックエンドではあまり速度が出なかったことから推奨しておりません。

依存関係のインストール

$ pip install optuna optuna-dashboard[preferential] diffusers transformers accelerate scipy safetensors xformers

入力画像の用意

下記画像はOptunaのロゴに手足のようなものを足したものです。 input.png というファイル名で main.py 実行時のlocationと同じディレクトリ内に保存してください。 （※ 本画像の利用はデモの動作確認のみにとどめてください）

実行

次の2つのコマンドをそれぞれ実行してください。

• python generator.py
• optuna-dashboard sqlite:///db.sqlite3 --artifact-dir ./artifact

macOSをお使いの方

xformersのインストール

本プログラムの実行に必要な xformers はmacOS向けのwheelバイナリを配布していないため、macOSではsdist (ソース配布パッケージ) からビルドされます。その際、Appleが提供しているClangにおいてOpenMPがデフォルトでは利用出来ないことから、次のようなエラーにぶつかることがあります。

$ pip install xformers
 ...
      clang -Wno-unused-result -Wsign-compare -Wunreachable-code -DNDEBUG -g -fwrapv -O3 -Wall -I...  -O3 -fopenmp -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE="_clang" -DPYBIND11_STDLIB="_libcpp" -DPYBIND11_BUILD_ABI="_cxxabi1002" -DTORCH_EXTENSION_NAME=_C -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++17
      clang: error: unsupported option '-fopenmp'
      error: command '/usr/bin/clang' failed with exit code 1
      [end of output]

この場合は libomp を別途インストールしたり、別のCコンパイラに切り替えることでインストールが可能です。例えば後者は次のようにできます。

$ brew install gcc
$ ls /usr/local/bin/gcc-*
/usr/local/bin/gcc-12 ...
$ CC=gcc-12 pip install xformers
...
Successfully installed xformers-0.0.16

main.py の変更

macOSで動かしたい方は、 main.py に次のpatchを適用してください。

-device = "cuda:0"
-torch_dtype = torch.float16
+device = "mps"
+torch_dtype = torch.float32

 pipe = pipe.to(device)
-pipe.enable_xformers_memory_efficient_attention()