はじめに
DreaBoothはgoogleらがCVPR2023で発表したdiffusion model系の再学習のしくみ。
[0] N. Ruiz, et. al. "DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation" CVPR2023.
再学習させるデータに過適合するのを防ぎ、また構文の意味を損なわないように配慮して再学習させる。
- dreamboothのロジックは以下にざっくりまとめまています
https://qiita.com/masataka46/items/134e08af7fcb4d7a726c
既に多くの解説が出回っているので、ここではハマりどころを重点的にまとめる。
1. 環境構築
以下の２つを試した。
- 自分のRTX3090パソコンにDockerで環境構築
- colabで環境構築
colab系は既に多くの資料があるので、ここでは前者を記述。
1.1 Ubuntuに基本的なものを入れる
- Nvidia driver
- GPU対応のDocker
- その他、pycharmやvs-code、iterm など
ここでは省略。
1.2 Dockerを立ち上げる
例えば以下のようなDockerfileを用いる。
FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04
ENV TZ=Asia/Tokyo
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime
RUN echo $TZ > /etc/timezone
RUN apt-get update
RUN apt-get install -y apt-file
RUN apt-file update
RUN apt-get install -y git
RUN apt-get install -y wget
RUN apt-get install -y curl
RUN apt-get install -y vim
# python
RUN apt-get install -y python3-pip
# alias
RUN echo alias python="python3" >> /root/.bashrc
RUN echo alias pip="pip3" >> /root/.bashrc
# Install python packages
RUN pip install --upgrade pip
RUN pip install opencv-python
RUN pip install scipy
RUN pip install scikit-learn
RUN pip install poetry
# Install VIM
RUN apt-get update
RUN apt-get install -y apt-file
RUN apt-file update
RUN apt-get install -y vim
# insatll langchain, and else
RUN pip install langchain==0.1.14
RUN pip install faiss-gpu tiktoken sentence_transformers
RUN pip install wandb
RUN CMAKE_ARGS="-DLLAMA_CUBLAS=on" FORCE_CMAKE=1 pip install llama-cpp-python
RUN pip install flash_attn
RUN pip install openai==0.28.0
RUN pip install pandas==2.2.2
RUN pip install -U langchain-community
RUN pip install openpyxl==3.1.2
RUN pip install python-docx==1.1.2
RUN pip install pypdf==5.1.0
RUN pip install git+https://github.com/huggingface/diffusers
RUN pip install accelerate>=0.16.0
RUN pip install torchvision==0.20.1
RUN pip install ftfy
RUN pip install tensorboard
RUN pip install peft==0.7.0
WORKDIR /root/share
EXPOSE 8080/TCP
必要ないものも多分に含まれています。
diffusersもinstallしています。
# in the case to use cuda
IMAGE_REPOSITORY=dreambooth
IMAGE_TAG=latest
IMAGE_FULLNAME=${IMAGE_REPOSITORY}:${IMAGE_TAG}
# docker build if changed.
docker build -t ${IMAGE_FULLNAME} .
# allow display connection
xhost local:
# run container
docker run \
--interactive \
--tty \
--rm \
--mount=type=bind,src="$(pwd)",dst=/root/share \
--mount=type=bind,src=/etc/group,dst=/etc/group,readonly \
--mount=type=bind,src=/etc/passwd,dst=/etc/passwd,readonly \
--mount=type=bind,src=/tmp/.X11-unix,dst=/tmp/.X11-unix \
$( if [ -e $HOME/.Xauthority ]; then echo "--mount=type=bind,src=$HOME/.Xauthority,dst=/root/.Xauthority"; fi ) \
--gpus=all \
--env=QT_X11_NO_MITSHM=1 \
--env=DISPLAY=${DISPLAY} \
--env=NVIDIA_DRIVER_CAPABILITIES=all \
--net=host \
--publish=8080:8080 \
--publish=8080:8080/udp \
--name=${IMAGE_REPOSITORY}_${IMAGE_TAG}_$(date "+%Y_%m%d_%H%M%S") \
${IMAGE_FULLNAME}
このDockerRun.shとDockerfileを同じディレクトリ内に配置し、Dockerをbuild, runする。
# DockerRun.sh, Dockerfileがある階層にて
bash DockerRun.sh
ポイントは、GPUで演算させる場合、torch_xla をインストールしないこと。このパッケージはTPU等利用する場合に必要のようだが、GPUで演算させる場合はエラーとなる。
2. accelerateの設定をする
accelerate config
でGPU周りなどの設定をする。ここが一番のハマりどころ。いろいろ聞かれるが、設定を間違えると動作しない。
私の場合、
- 自分のlocalパソコン
- GPUは１個
なので、
In which compute environment are you running?
Please select a choice using the arrow or number keys, and selecting with enter
➔ This machine
AWS (Amazon SageMaker)
This machineでリターン。
Which type of machine are you using?
Please select a choice using the arrow or number keys, and selecting with enter
➔ No distributed training
multi-CPU
multi-XPU
multi-GPU
multi-NPU
multi-MLU
multi-MUSA
TPU
私の場合、GPU１個なので、いちばん上の「No distributed training」でリターン。
Do you want to run your training on CPU only (even if a GPU / Apple Silicon / Ascend NPU device is available)? [yes/NO]:
GPU使うので、NO。
Do you wish to optimize your script with torch dynamo?[yes/NO]:
よくわからないのでNO。
Do you want to use DeepSpeed? [yes/NO]:
よくわからないのでNO。
What GPU(s) (by id) should be used for training on this machine as a comma-seperated list? [all]:
ここではGPUのIDを入力する。私の場合、
nvidia-smi で
nvidia-smi
Wed Jan 1 09:20:16 2025
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.154.05 Driver Version: 535.154.05 CUDA Version: 12.2 |
|-----------------------------------------+----------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+======================+======================|
| 0 NVIDIA GeForce RTX 3090 On | 00000000:05:00.0 On | N/A |
| 0% 43C P8 28W / 350W | 310MiB / 24576MiB | 0% Default |
| | | N/A |
+-----------------------------------------+----------------------+----------------------+
+---------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=======================================================================================|
+---------------------------------------------------------------------------------------+
と表示されたので、GPU欄に記載の０を入力。
Would you like to enable numa efficiency? (Currently only supported on NVIDIA hardware). [yes/NO]:
よくわからないのでNO。
Do you wish to use mixed precision?
Please select a choice using the arrow or number keys, and selecting with enter
➔ no
fp16
bf16
fp8
学習時のbit数のことか？float32で学習させる気なので最上部でリターン。
うまく設定できれば
accelerate configuration saved at /root/.cache/huggingface/accelerate/default_config.yaml
と表示される。
3. 学習させる画像を用意する
適当な画像を適当なディレクトリに配置する。私は特定の種類の犬画像を５種類用意した。
4. train_dreambooth.pyを持ってくる
下記リンクの diffusers/examples/dreambooth/train_dreambooth.py をコピペして適当な位置に配置。
5. outputのディレクトリを作成
（自動で作成してくれるのかもしれないが）outputのディレクトリを作成する。
例えば
mkdir ../output
6. 環境変数の設定
今回、diffusionのv1-4を使う。
export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export INSTANCE_DIR="PATH/TO/IMAGE_DIR"
export OUTPUT_DIR="PATH/TO/OUTPUT_DIR"
7. 学習開始
以下のようなコマンドで学習を開始。
accelerate launch train_dreambooth.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--instance_data_dir=$INSTANCE_DIR \
--output_dir=$OUTPUT_DIR \
--instance_prompt="a photo of sks dog" \
--resolution=512 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--learning_rate=5e-6 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--max_train_steps=400
instance_promptは適当なものに設定する。その他、目的に合わせて諸々ハイパラを変更する。
以下のようにVRAM 19GB近く使ってました。減らすには上記の量子化すればよい？
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.154.05 Driver Version: 535.154.05 CUDA Version: 12.2 |
|-----------------------------------------+----------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+======================+======================|
| 0 NVIDIA GeForce RTX 3090 On | 00000000:05:00.0 On | N/A |
| 43% 61C P2 303W / 350W | 18953MiB / 24576MiB | 92% Default |
| | | N/A |
+-----------------------------------------+----------------------+----------------------+
+---------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=======================================================================================|
| 0 N/A N/A 5384 G /usr/lib/xorg/Xorg 168MiB |
| 0 N/A N/A 5517 G /usr/bin/gnome-shell 75MiB |
| 0 N/A N/A 6653 G ...2509800,12662543674627072428,262144 55MiB |
| 0 N/A N/A 762978 C /usr/bin/python3 18638MiB |
+---------------------------------------------------------------------------------------+
10分程度で学習は終了する。
outputしたディレクトリに
feature_extractor logs model_index.json safety_checker scheduler text_encoder tokenizer unet vae
のように諸々保存される。
8. 学習したモデルを用いて推論
今回、以下のリンクにある diffusers/src/diffusers/pipelines/のREADME.mdに従って推論させる。
以下の
# make sure you're logged in with `huggingface-cli login`
from diffusers import StableDiffusionPipeline, LMSDiscreteScheduler
pipe = StableDiffusionPipeline.from_pretrained("../output/")
pipe = pipe.to("cuda")
prompt = "a photo of sks dog"
image = pipe(prompt).images[0]
image.save("tmp_dog.png")
を記述した infer.pyを作成する。
- 冒頭pipeのpretrain読み込む部分は学習時に
OUTPUT_DIRで指定したディレクトリを指定 - promptは学習時に指定したpromptを指定
- 保存する画像名はテキトーに指定
python infer.py
で推論させる。それっぽい画像が生成される。