NovelAI Aspect Ratio Bucketing の翻訳

これは NovelAI Aspect Ratio Bucketing の翻訳だ。

要約すると、Stable Diffusion は任意の解像度とアスペクト比で学習可能だが、バッチごとに解像度は固定する必要がある（固定しないと学習が遅い）。なのでバッチを実行する際にアスペクト比を選び、そのバッチでは選択したアスペクト比の画像のみを学習させる。原文では、アスペクト比の一致しない学習画像の加工処理を実行時にしているように見えるが、実際は前処理でスケール＆クロップしているものと思われる。

NovelAI Aspect Ratio Bucketing

aspect ratio bucketing を使ったトレーニングは出力品質を大きく改善させる（個人的にはセンタークロップで学習させたベースモデルはこれ以上増えて欲しくない）。なので、bucketing code を permissive MIT ライセンスで公開することにした。

このリポジトリは NovelAI Improvements on Stable Diffusion で説明した、画像生成モデルを学習させるための aspect ratio bucketing の実装を提供する。

開発の経緯

現在の画像生成モデルの共通する課題は、生成される画像が不自然にクロップされることだ。これは正方形画像を生成するようにモデルが訓練されるために起こる。しかし、ほとんどの写真やイラストは正方形ではない。

このミスマッチは GPU の最適化に原因がある。GPU を効率的に利用するためには同時に複数の画像を学習させなければならない。しかしそのためにはモデルを訓練するさいの画像は同じサイズにする必要がある。妥協案として、正方形のアスペクト比が選択され、中央を切り出した画像をモデルが学習することになる。

その結果モデルは、頭や足がフレームアウトした人間や、剣先がフレームアウトした剣の画像を学ぶことになる。我々はストーリーテリングに付随する画像生成モデルを作っているので、ノーカットでキャラクターを生成できることが重要で、生成された騎士が金属っぽい直線が無限に伸びているようなものを持っていてはならない。

クロップされた画像で訓練することのほかの問題はテキストと画像とのミスマッチだ。たとえば crown タグのついた画像はクロップ後の画像に王冠が含まれてないことがよくある。

テキストと画像とのミスマッチの問題はランダムクロップによってわずかに改善されることが分かった。

Stable Diffusion はさまざまな画像サイズを出力できる。しかし 512x512 から離れると要素が繰り返し現れ、小さい解像度では使い物にならない。

しかしこれは、可変サイズの画像でモデルを学習することが可能であることを示している。さまざまなアスペクト比の画像を１枚ずつ学習させることは可能だが、学習に時間がかかりすぎる。加えてミニバッチによる正則化ができないので学習が不安定になりがちだ。

Custom Batch Generation

この問題に対する解決策が存在するように見えないので、カスタムバッチ生成コードを書くことにした。これは画像のアスペクト比で画像を分類し、バッチを作成する。我々はこれを aspect ratio bucketing と呼んでいる。ほかのアイデアとしては、画像サイズを固定し、そのサイズに合うように画像をスケールし、トレーニング時に余白をマスクする方法が考えられる。この方法は学習時に余計な計算が発生するので、使わなかった。

バケツサイズの定義

初めにバケツサイズを定義しなければならない。我々は最大総ピクセル数 393,216（512x768）で最大長辺 1024 ピクセルを選択した。最大画像サイズが 512x512 より大きいので、VRAM 消費量が大きく、GPU あたりのバッチサイズも少なくなるが、gradient accumulation によって相殺できる。

以下のアルゴリズムでバケツを生成する：

幅を 256 に設定する
幅が 1024 以下の間以下をループ：

高さが 1024 以下かつ総ピクセル数 393,216 以下という条件で、最大高さを探す
見つけた高さと幅でバケツを作成
幅を 64 増加

同じことを幅と高さとを入れ替えて実行する。重複を削除し、最後に 512x512 のサイズのバケツを追加する。

サンプルコード

min_height = 256
max_height = 1024
min_width = 256
max_width = 1024
step = 64
max_pixel_count = 512*768
buckets = [[512, 512]]
for width in range(min_width, max_width, step):
    h = max(filter(lambda x: x*width <= max_pixel_count, range(min_height, max_height+1, step)))
    if not [width, h] in buckets:
        buckets.append([width,h])
    if not [h, width] in buckets:
        buckets.append([h,width])
buckets.sort(key=lambda x: (x[0], x[0]/x[1]))
print(buckets)

画像の振り分け

次に対応するバケツに画像を入れる。

バケツ解像度を NumPy array に入れ、アスペクト比を計算する
データセット内の画像のアスペクト比を計算する
バケツのアスペクト比から画像のアスペクト比を引いて、その絶対値が最も小さいバケツに画像を入れる

image_bucket = argmin(abs(bucket_aspects - image_aspect))

画像のバケツ番号とアイテム ID とが関連付けられてデータセット内に保存される。画像のアスペクト比がバケツサイズとかけ離れている場合は、その画像はデータセットから除外される。

サンプルコード

aspects = list(map(lambda x: x[0]/x[1], buckets))
image_width = 3421
image_height = 7644
image_aspect = image_width/image_height
print(min(map(lambda x: [abs(x[1]-image_aspect), x[0]], enumerate(aspects)), key=lambda x:x[0]))

アイテムをプロセスへ振り分け

複数の GPU で訓練する場合、それぞれのエポックの前に、各 GPU が同じサイズの個別のサブセットで作業できるように、データセットを分割する必要がある。これをするにはまず、データセット内のアイテム ID のリストをコピーしてシャッフルする。コピーされたリストが GPU 数 * バッチサイズで割り切れない場合は、割り切れるようにリストの末尾を除外する。

そしてそのリストを「全 GPU 数 * バッチサイズ（world_size*bsz）」個に分け、現在のプロセスの GPU ID（global rank）に基づいて分配する。残りのカスタムバッチ生成はそれぞれの以下のプロセスを個別に実行する。

各プロセスでの処理

まずアイテム ID を見て画像をバケツへ振り分ける。それぞれのバケツ内の画像数がバッチサイズで割り切れない場合、割り切れるようにそのバケツ内から画像を除外する。除外した画像は catch-all バケツに入れておく。振り分けられたアイテム ID の数がバッチサイズで割り切れるので、catch-all バケツを含むすべてのバケツの画像数はバッチサイズで割り切れる。

バッチが要求されると、重み分布（weighted distribution）に基づいてバケツを選択する。バケツのウェイトは、選択されたバケツの画像数 / 残りのバケツの画像数の合計、で計算される。これによって、バケツのサイズが異なっていても、訓練中に強いバイアスが発生しなくなる。もしウェイトなしでバケツを選択した場合、小さいバケツがすぐに空になってしまい、訓練の終了間際にはサイズの大きいバケツしか残らなくなる。

アイテムは選択されたバケツから取り出され、取り出されたアイテムはバケツから削除される。バケツが空なら、そのエポックではそのバケツは削除される。選択されたアイテム ID と選択されたバケツの解像度は image-loading 関数に渡される。

Image Loading

image loading コードはこのリポジトリには含まれていないが、実装は簡単だ。

それぞれのアイテム ID の画像はロード後にバケツ解像度に合うように処理される。これには２つのアプローチ方法がある。

ひとつは単純にスケールする方法だ。これはわずかに画像がゆがむ。なので我々は別の方法を選択した。

アスペクト比を維持したまま余白が出ないよう画像をスケールし

バケツのアスペクト比と画像のアスペクト比が一致しているならそのまま
アスペクト比が一致しないなら、ランダムクロップ

我々のケースではイメージあたりの平均アスペクト比エラーは 0.033 であり、実際の画像では 32 ピクセルより小さかった。