Last active
October 28, 2023 01:29
-
-
Save c-bata/449f2e90ac50a1285b7fe210ab51eae6 to your computer and use it in GitHub Desktop.
PyCon APAC 2023 - ハイパーパラメータ最適化フレームワーク Optunaの最新機能紹介
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# 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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
動作の様子
実行方法
推奨実行環境
本Exampleを動かす際は Linux / NVIDIA GPUの利用を推奨 しています。もしGPUがなく動かすのが大変という方は、Optuna Dashboard公式ドキュメントの下記チュートリアルをお試しください。こちらはCPUのみで簡単に動かすことができ、Human-in-the-loop最適化のためのコードの理解には十分です。
https://optuna-dashboard.readthedocs.io/en/latest/tutorials/preferential-optimization.html
依存関係のインストール
入力画像の用意
下記画像は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がデフォルトでは利用出来ないことから、次のようなエラーにぶつかることがあります。この場合は
libomp
を別途インストールしたり、別のCコンパイラに切り替えることでインストールが可能です。例えば後者は次のようにできます。main.py
の変更macOSで動かしたい方は、
main.py
に次のpatchを適用してください。pipe = pipe.to(device) -pipe.enable_xformers_memory_efficient_attention()