こんにちは、CCCMKホールディングスTECH LAB三浦です。
すっかりと暖かくなり、春らしくなりました。近所の学校や保育園で卒業式や卒園式が行われているのを見ると、新しい季節がやって来るんだなぁとしみじみ感じます。
さて、今回は前回に引き続き画像生成の話を紹介させて頂きます。
はじめに
前回Hugging Faceの"Diffusers"というライブラリを使って画像生成で利用される"Diffusion Model"を作ってみた話をまとめました。
今回はモデルの構造を色々変えてみた時に、画像生成にどのような影響が出るのかが気になったので、実際に試してみました。
特にモデルにAttention Blockをどれだけ含めるのか、クラスラベルの条件を与えるとどうなるのか、の2つに着目して調べています。生成される画像の質は、本当はFIDなどの指標で定量的に評価をするべきなのですが、今回は生成された画像を目視確認するだけに留まっています。定量的な評価の仕組みの構築は今後の課題だと認識しています。
使用したデータ
前回と同じく、Hugging Faceで公開されている人の手の画像のデータセットを使わせて頂きました。
このデータセットには約2,000件の画像が含まれていますが、その中からランダムで1,000件抽出して学習に使用しました。
また、各画像にはその画像に含まれる手の形に応じた8種類のラベルが付与されています。 (たとえばラベル0は"closed fist", つまり"握りこぶし") この情報も今回はクラスラベルによる条件付きの画像生成の際に利用します。
また、画像サイズは64x64にリサイズして使用しました。モデルの出力も64x64にしています。
クラスラベル条件付き画像生成
画像を生成する際にランダムノイズに加え、どのクラス(種類)の画像を生成するのかを追加情報として与えることが出来ます。これを実現するために、まずUNetを定義する際にパラメータnum_class_embedsにクラスラベル数を指定する必要があります。今回使用したデータセットには8つのラベルが含まれているので、8を指定します。
model = UNet2DModel(
...
num_class_embeds=8, #class conditions
...
)
モデルへの入力は次のようにclass_labelsパラメータの指定が必要になります。
clean_images = batch["image"] labels = batch["label"] ... noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps) with accelerator.accumulate(model): noise_pred = model( noisy_images, timesteps, class_labels=labels, # 生成したい画像のクラスの指定が必要 return_dict=False )[0] ...
Attention層の数の変更
UNetのUpSampling, DownSampling BlockをそれぞれDownBlock2DからAttnDownBlock2D, UpBlock2DからAttnUpBlock2Dに変更するとどうなるのかを確認します。
ベースのモデル構造は以下です。
model = UNet2DModel(
sample_size=64,
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"
)
)
Attention Blockを2つにした場合は以下です。
from diffusers import UNet2DModel model = UNet2DModel( ... down_block_types=( "DownBlock2D", "DownBlock2D", "AttnDownBlock2D", "DownBlock2D", "AttnDownBlock2D", "DownBlock2D" ), up_block_types=( "UpBlock2D", "AttnUpBlock2D", "UpBlock2D", "AttnUpBlock2D", "UpBlock2D", "UpBlock2D" ) )
さらに全てAttention Blockに変えたパターンも試してみました。
from diffusers import UNet2DModel model = UNet2DModel( ... down_block_types=( "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D" ), up_block_types=( "AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D" ) )
評価用画像の生成
100epochの学習期間の中で10epochずつ16枚の画像を生成させます。
各パターンによる画像生成結果
クラスラベル条件ありなし、Attention Blockの個数を変えて計6パターンで学習を行い、生成した評価用画像をまとめました。
気付いたこと
今回試した限りだとなかなか評価が難しい・・・と感じました。
Attention Blockを増やすことによる効果
Attention Blockを増やすと、全体的に暗い画像になる印象を受けました。Attention Blockの追加によって生成画像の品質が向上する、という傾向は今回は見られませんでした。全体としてのBlock数の最適化、各Blockで出力するチャネル数の最適化などが必要なのかもしれません。
クラスラベル条件による効果
こちらはクラスラベル条件を与えることで、生成される画像に変化が見られたと思います。具体的には条件を与えた方が画像に含まれる手の輪郭が明確になる傾向が見られました。学習時に画像だけでなく、その画像がどういったカテゴリに属するものなのかという情報を与えることが、画像生成において影響を与えていることが伺えました。
まとめ
今回はひたすら"Diffusers"でシンプルなDiffusion Modelを色々な設定で学習させた時の様子を調べてまとめてみました。Attention Blockを増やすと品質が向上する、と思っていたのですが、今回試した範囲ではその傾向が見られなかったのが意外でした。単純に増やすのではなく、増やし方にもなんらかのルールが存在するのかもしれません。一方でクラスラベルによる条件は、一定の効果がありそうです。そして画像系の検証は結果が目に見えて理解できるので、楽しいな、と思いました。
