はじめに
こんにちは、CCCMKホールディングスTECH LABの三浦です。
3月ももう中旬です。まだ寒い日もありますが、少しずつ暖かい日も増えてきました。この時期は「これまでの締め」と「新しいことへの準備」が交互にやってくるので頭がグルグルしてきます。
さてここ数回の記事では画像生成で使われるDiffusion Modelについて何回か触れてきましたのですが、自分自身が手を動かしてDiffusion Modelを構築してみる、という経験がありませんでした。
Hugging Faceの"Diffusers"というライブラリを使うと比較的シンプルにDiffusion Modelを学習させることが出来ることが分かり、実際に試してみようと思いました。
今回の記事では、Diffusion Modelを学習するためのコードを一部掲載していますが、基本的にはDiffusersのTutorialに含まれている"Train a diffusion model"の内容に従い、一部分変更を加えたものです。
Diffusion Modelの学習方法
最初にDiffusion Modelの学習の手順について整理してみます。"画像生成"に使われるDiffusion Modelですが、モデル自体はノイズが付与された画像を受け取り、付与されたノイズを推計し出力する動作をします。ノイズは入力画像と同サイズのテンソルで、各成分は平均0, 分散1の標準正規分布に従います。
ノイズを画像に付与する強度も特定のルールに従って変更します。ハイパーパラメータとして指定する"timestep"を用いて0ステップから"timestep"ステップの間でノイズの強度を0に近いほど低く、"timestep"に近いほど強くするように変更します。
これらを踏まえてDiffusion Model学習時の流れをまとめると次のようになります。
- 学習データから画像を取得
- 画像と同サイズのノイズを標準正規分布に従ってランダム生成
- 0～"timestep"の範囲の整数"t"をランダム選択
- ノイズを"t"に基づいた強度で画像に付与
- モデルにノイズ付の画像を入力し出力がノイズと一致するようにパラメータを更新
学習したDiffusion Modelを推論に使う場合は逆の手順を取ります。ノイズを入力し、モデルが出力した結果をノイズから取り除きます。さらにそれをモデルに入力し、出力した結果を取り除き、といったことを繰り返していくことでノイズのない、学習データの分布に従った画像が生成出来ます。
学習に使用したデータセット
DiffusersのTutorialでは"Smithsonian Butterflies dataset"という蝶々の画像が使用されています。
せっかくなので違うデータセットで試してみることにしました。Hugging Faceのデータセットで公開されている"trashsock/hands-images"という人の手の画像1,996枚で構成されたデータセットです。
"Smithsonian Butterflies dataset"の枚数と揃えるため、ランダムサンプリングで1,000枚の画像のデータセットを抜き出し、学習用データセットとして使用しました。
利用環境
Azure Databricksで行いました。"Standard_NC24ads_A100_v4"というタイプのVMを使いました。A100のGPUが1基搭載されています。学習時のメトリクスはMLflowに記録しました。
コードについて
ここからは学習時に使用したコードについて、ポイントを絞ってご紹介します。
学習の設定
学習時に使用する情報をdataclassを使って設定します。基本的な設定はTutorialに従っていますが、学習時のメトリクスをMLflowに送る為の設定を追加し、モデルをHugging Face Hugに送らない設定にしています。 また、学習epochも試しに50から100まで増加してみました。
from dataclasses import dataclass @dataclass class TrainingConfig: image_size = 128 train_batch_size = 16 eval_batch_size = 16 num_epochs = 100 gradient_accumulation_steps = 1 learning_rate = 1e-4 lr_warmup_steps = 500 save_image_epochs = 10 save_model_epochs = 30 mixed_precision = "fp16" output_dir = "hands-images" push_to_hub = False hub_private_repo = None overwrite_output_dir = True seed = 0 mlflow_experiment_name="/Users/user_name/hands-images" #MLflow記録設定 config = TrainingConfig()
1,000件のデータセットの作成
"trashsock/hands-images"は2,000件のデータセットとなっており、ここから学習に使用する1,000件のデータをランダムサンプリングしました。
from datasets import load_dataset config.dataset_name = "trashsock/hands-images" dataset = load_dataset(config.dataset_name, split="train")\ .shuffle(seed=42)\ .select(list(range(1000)))
データセットに含まれる画像を確認します。
import matplotlib.pyplot as plt fig, axs = plt.subplots(1, 4, figsize=(16, 4)) for i, image in enumerate(dataset[:4]["image"]): axs[i].imshow(image.resize((config.image_size, config.image_size))) axs[i].axis("off") fig.show()
画像の変換処理の適用
画像に対しサイズの変更やランダムな左右反転、ピクセルの値の正規化処理を行います。
from torchvision import transforms preprocess = transforms.Compose( [ transforms.Resize((config.image_size, config.image_size)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[0.5], std=[0.5]) ] ) def transform(examples): images = [preprocess(image.convert("RGB")) for image in examples["image"]] return {"image": images} dataset.set_transform(transform)
結果を確認してみます。
fig, axs = plt.subplots(1, 4, figsize=(16, 4)) for i, image in enumerate(dataset[:4]["image"]): axs[i].imshow(image.permute(1, 2, 0)*0.5 + 0.5) axs[i].axis("off") fig.show()
Schedulers
Diffusersには画像に対するノイズ付与を担う"Schedulers"という機能が実装されています。色々な方法でノイズ付与が可能で、"DDPMScheduler"などのクラスで実装されています。
今回はDDPMSchedulerを使用しました。どのようにノイズが付与されるのかを確認してみます。"timesteps"として1,000を設定し、0～999までのステップにおいてどのようにノイズが付与されるのかを可視化します。
import torch from PIL import Image from diffusers import DDPMScheduler sample_image = sample_image.unsqueeze(0) noise_scheduler = DDPMScheduler(num_train_timesteps=1000) noise = torch.randn(sample_image.shape) # ノイズ生成 timesteps = torch.LongTensor([0, 100, 500, 999]) # 4つのtimestepでノイズ付与状況を確認 noisy_images = noise_scheduler.add_noise(sample_image,noise,timesteps) fig, axs = plt.subplots(1, 4, figsize=(16, 4)) for i, image in enumerate(noisy_images): axs[i].imshow(image.permute(1, 2, 0)*0.5 + 0.5) axs[i].axis("off") fig.show()
ステップが1,000に近づくほど、ノイズが強くなっていることが確認出来ます。
モデル構造の定義
Tutorialと同じ、12ブロックのUNetでモデルを定義しました。
from diffusers import UNet2DModel model = UNet2DModel( sample_size=config.image_size, in_channels=3, out_channels=3, layers_per_block=2, block_out_channels=(128, 128, 256, 256, 512, 512), down_block_types=( "DownBlock2D", "DownBlock2D", "DownBlock2D", "DownBlock2D", "AttnDownBlock2D", "DownBlock2D" ), up_block_types=( "UpBlock2D", "AttnUpBlock2D", "UpBlock2D", "UpBlock2D", "UpBlock2D", "UpBlock2D" ) )
学習処理の変更
学習処理はTutorialの内容をほぼ踏襲したのですが、一部メトリクスの記録部分をMLflow向けに変更しました。
まず学習処理が定義されたtrain_loop関数の内部の一部を、以下の様に変更しました。
... progress_bar.update(1) # MLflow向けに変更 logs = { "loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0], } progress_bar.set_postfix(**logs) if accelerator.is_main_process: mlflow.log_metrics(logs,step=global_step) global_step += 1 ...
そして学習実行時の処理を、MLflow向けに書き換えました。
from accelerate import notebook_launcher mlflow.set_experiment(config.mlflow_experiment_name) with mlflow.start_run(): args = (config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler) notebook_launcher(train_loop, args, num_processes=1)
学習結果
メトリクスの推移
100epoch経過後のloss(ノイズとのMSE Loss)と学習率の推移です。
lossは収束しているように見えます。
生成画像の確認
100epoch経過後にtimesteps 1,000で生成させた16枚の画像がこちらです。なんとなく手っぽいです。
10epochの学習完了ごとに生成された画像と比べてみます。
一部は手に近づいている傾向があり、一部はなかなか特徴が捉えられていないような印象を受けました。
今回使用したデータセットの中には様々な形状や向きの手の画像が含まれています。特に手の向きというのが今回の学習では上手く捉え切れていないのかな、と感じました。
学習epochをもっと増やすべきなのかもしれませんし、UNetの構造をもう少し深くする必要があるのかもしれません。この辺りを今後試していきたいと思いました。
まとめ
ということで、今回はHugging FaceのDiffusersというライブラリを使ってDiffusion Modelの構築に取り組んでみた話をまとめてみました。実際に自分で試してみると、Diffusion Modelを使って高品質な画像を生成するためには色々と調整しなければいけない箇所が多いんだな、と感じました。今回の結果よりももっと良い品質の画像を作るためにはどんなアプローチがあるのか、という観点で、Diffusion Modelについてもっと深掘りしていきたいと感じました。
