たった数枚の画像で Stable Diffusion をファインチューニングできる効率的な Amazon SageMaker JumpStart の使い方

2022 年 11 月、AWS のお客様が Amazon SageMaker JumpStart で Stable Diffusion モデルを使ってテキストから画像を生成できるようになったことを発表しました。Stable Diffusion は、リアルで高品質な画像や印象的なアートをわずか数秒で生成することができるディープラーニングモデルです。印象的な画像の作成は、アートから NFT など幅広い業界で活用が見込めますが、今日ではユースケースにパーソナライズされた AI が期待されています。本日、Amazon SageMaker JumpStart でカスタムデータセットでファインチューニングすることにより、ユースケースに合わせて画像生成モデルをパーソナライズすることが可能になりました。これは、アート、ロゴ、カスタムデザイン、NFT などを作成するときや、ペットのカスタム AI 画像や自分のアバターを生成するなどのポテンシャルのある興味深いユースケースに役立ちます。

この記事では、SageMaker Python SDK で利用できる JumpStart API を使ったプログラムによる方法と、Amazon SageMaker Studio の JumpStart のユーザーインターフェース（UI）という2つの方法で、Stable Diffusion モデルをファインチューニングする方法の概要を説明します。また、データセットの品質、トレーニングデータセットのサイズ、ハイパーパラメータ値の選択、複数のデータセットへの適用性などの設計方法についても説明します。最後に、最近 JumpStart に追加された、異なる入力言語やスタイルを持つ 80 以上の一般利用可能なファインチューニング済みモデルについて説明します。

Stable Diffusion と転移学習

Stable Diffusion は、テキストから画像への変換モデルで、テキストプロンプトから写真のようにリアルな画像を作成することが可能です。Diffusion モデルは、実際の画像に加えられたノイズを除去するプロセスを学習するように訓練されています。このノイズ除去処理により、リアルな画像が生成されます。また、これらのモデルは、生成処理をテキストで条件付けすることで、テキストのみから画像を生成することも可能です。（詳細については割愛しますが、仕組みついて興味のある方はこちらの記事をご参照ください）

機械学習（ML）において，ある領域で学習した知識を別の領域に移すことを転移学習といいます。転移学習を用いると、少ないデータセット、少ない学習時間で正確なモデルを作成することができます。転移学習を用いると、わずか5枚の画像で、Stable Diffusion モデルを自身のデータセットでファインチューニングすることができます。下図の例では、”Dog named Doppler”（ドップラーという名前の犬）と数枚のドップラーの画像を用意し学習することで「ビーチにいるドップラー」や「鉛筆でスケッチしたドップラー」の画像を生成しています。

次に示す左の例では ”my chair in white color“ （白い私の椅子）というテキストと複数の椅子の画像で「私の椅子」についてファインチューニングしたモデルで赤い椅子を生成しています。右の例では、”my ottoman“（私のオットマン）というテキストと複数のオットマンの画像で「私のオットマン」についてファインチューニングしたモデルで「私のオットマンに座る猫」の画像を生成しています。

Stable Diffusion のような大規模なモデルのファインチューニングには、通常、トレーニングスクリプト（学習用のコード）を提供する必要があります。また、メモリ不足の問題、ペイロードサイズの問題など、様々な問題があります。さらに、スクリプト、モデル、利用するインスタンスが効率的に動作することを確認するために、エンドツーエンドテストを実行する必要があります。SageMaker JumpStart は、堅牢にテストされたすぐに使えるスクリプトを提供することで、このプロセスを簡素化します。Stable Diffusion モデル用の JumpStart ファインチューニングスクリプトは、DreamBooth の手法に沿った学習スクリプトをあらかじめ提供することで、上記のような数枚でのファインチューニングを可能にしています。これらのスクリプトは、Studio UI からワンクリックで、または JumpStart API からごくわずかな行のコードでアクセスすることができます。

Stable Diffusion モデルを使用する場合、CreativeML Open RAIL++-M ライセンスに同意する必要があることに注意してください。

Studio UI から JumpStart にアクセスする

このセクションでは、Studio UI を使用して JumpStart モデルをトレーニングおよびデプロイする方法を説明します。次の動画は、JumpStart で事前トレーニングされた Stable Diffusion モデルを見つけ、それをトレーニングし、デプロイする方法を示しています。モデルページには、モデルやその使用方法に関する貴重な情報が記載されています。SageMaker のトレーニングインスタンスを設定した後、Train を選択します。モデルのトレーニングが完了したら、推論用インスタンスを選び Deploy を選択してトレーニング済みのモデルをデプロイできます。エンドポイントが「In service」のステージになると、推論リクエストに応答する準備が整います。

推論にかかる時間を短縮するために、JumpStart には、新しく作成したエンドポイントで推論を実行する方法を示すサンプルノートブックが用意されています。Studio でノートブックにアクセスするには、モデルエンドポイントページの「Use Endpoint from Studio」セクションで「Open Notebook」を選択します。

JumpStart には、Stable Diffusion モデルをファインチューニングし、その結果得られたファインチューニング済みモデルをデプロイするために使用できる簡単なノートブックも用意されています。これを使用して、犬の面白い画像を生成することができます。このノートブックにアクセスするには、JumpStart の検索バーで「 Generate Fun images of your dog（犬の面白い画像を生成）」を検索してください。ノートブックを実行するには、わずか5枚のトレーニング画像を使用し、ローカル Studio フォルダにアップロードします。5枚以上の画像がある場合は、それらもアップロードすることができます。ノートブックは、トレーニング画像を S3 にアップロードし、データセットでモデルをトレーニングし、結果のモデルをデプロイします。トレーニングは 20 分程度で終了します。トレーニングのステップ数を変更することで、トレーニングのスピードを上げることができます。ノートブックには、デプロイされたモデルで試すためのサンプルプロンプトがいくつか用意されていますが、好きなプロンプトを試すことができます。また、このノートブックを使って、あなた自身やペットのアバターを作ることもできます。例えば、犬の代わりに猫の画像を最初のステップでアップロードし、プロンプトを犬から猫に変更すれば、モデルは猫の画像を生成します。

SageMaker SDK を使用したプログラムによる JumpStart の使用

このセクションでは、SageMaker Python SDK を使用してモデルをトレーニングおよびデプロイする方法について説明します。JumpStart で適切な事前学習済みモデルを選択し、SageMaker の学習ジョブでモデルを学習し、学習済みモデルを SageMaker のエンドポイントにデプロイします。さらに、デプロイされたエンドポイントで推論を実行しますが、これらはすべて SageMaker Python SDK を使用しています。以下の例にはコードスニペットが含まれています。

このデモのすべてのステップを含む完全なコードについては、「 Introduction to JumpStart – Text to Image example notebook 」を参照してください。このデモのコードを実行する際は、SageMaker Studio にデモノートブックをダウンロードして実行してください。SageMaker Studio Lab を利用している場合はこちらからも起動できます。デモの実行にかかるコストは約 1 ドルです。（ml.g4dn.2xlarge の場合、学習 ~30 分：$0.4、推論：$0.94/hour）

Stable Diffusion モデルのトレーニングとファインチューニング

各モデルは、一意の model_id で識別されます。次のコードは、カスタムトレーニングデータセット上で、model_id model-txt2img-stabilityai-stable-diffusion-v2-1-base で識別される Stable Diffusion 2.1 ベースモデルをファインチューニングする方法を示しています。model_id の全リストとファインチューニング可能なモデルについては、Built-in Algorithms with pre-trained Model Table を参照してください。各 model_id について、SageMaker Python SDK の Estimator クラスを通じて SageMaker のトレーニングジョブを起動するには、SageMaker で提供されるユーティリティ関数を通じて、Docker イメージ URI、トレーニングスクリプト URI、および事前学習済みモデル URI を取得する必要があります。トレーニングスクリプト URI には、データ処理、事前学習済みモデルの読み込み、モデルのトレーニング、学習済みモデルを推論のために保存するのに必要なすべてのコードが含まれています。事前学習済みモデル URI には、事前学習済みモデルのアーキテクチャ定義とモデルパラメータが含まれています。事前学習済みモデルの URI は特定のモデルに固有のものです。事前学習済みモデルの tarball は Hugging Face から事前にダウンロードされ、適切なモデル署名とともに Amazon Simple Storage Service (Amazon S3) のバケットに保存されており、トレーニングジョブが隔離されたネットワークで実行されるようになっています。以下のコードをご覧ください。

from sagemaker import image_uris, model_uris, script_uris

# Currently, not all the stable diffusion models in jumpstart support finetuning. Thus, we manually select a model
# which supports finetuning.
train_model_id, train_model_version, train_scope = (
"model-txt2img-stabilityai-stable-diffusion-v2-1-base",
"*",
"training",
)

# Tested with ml.g4dn.2xlarge (16GB GPU memory) and ml.g5.2xlarge (24GB GPU memory) instances. Other instances may work as well.
# If ml.g5.2xlarge instance type is available, please change the following instance type to speed up training.
training_instance_type = "ml.g4dn.2xlarge"

# Retrieve the docker image
train_image_uri = image_uris.retrieve(
region=None,
framework=None,  # automatically inferred from model_id
model_id=train_model_id,
model_version=train_model_version,
image_scope=train_scope,
instance_type=training_instance_type,
)

# Retrieve the training script. This contains all the necessary files including data processing, model training etc.
train_source_uri = script_uris.retrieve(
model_id=train_model_id, model_version=train_model_version, script_scope=train_scope
)

# Retrieve the pre-trained model tarball to further fine-tune
train_model_uri = model_uris.retrieve(
model_id=train_model_id, model_version=train_model_version, model_scope=train_scope
)

これらのモデル固有のトレーニング成果物を用いて、Estimator クラスのオブジェクトを構築することができます。

# Create SageMaker Estimator instance
sd_estimator = Estimator(
    role=aws_role,
    image_uri=train_image_uri,
    source_dir=train_source_uri,
    model_uri=train_model_uri,
    entry_point="transfer_learning.py",  # Entry-point file in source_dir and present in train_source_uri.
    instance_count=1,
    instance_type=training_instance_type,
    max_run=360000,
    hyperparameters=hyperparameters,
    output_path=s3_output_location,
    base_job_name=training_job_name,
)

# Launch a SageMaker Training job by passing s3 path of the training data
sd_estimator.fit({"training": training_dataset_s3_path}, logs=True)

トレーニングデータセット

トレーニング用データのフォーマットについて、以下に説明します。

入力 – インスタンス画像と dataset_info.json を含むディレクトリ。以下の条件に注意
- 画像は .png, .jpg, .jpeg のいずれかの形式である
- dataset_info.json は、{'instance_prompt':<<instance_prompt>>} という形式である
出力 – 推論用にデプロイ可能な学習済みモデル

S3 のパスは、s3://bucket_name/input_directory/ のようになります。末尾の / は必須であることに注意してください。以下はトレーニングデータのフォーマット例です。

input_directory
    |---instance_image_1.png
    |---instance_image_2.png
    |---instance_image_3.png
    |---instance_image_4.png
    |---instance_image_5.png
    |---dataset_info.json

Prior preservation を使用する際のデータフォーマットについては、この投稿の Prior preservation のセクションを参照してください。

デフォルトのデータセットとして、犬の画像を用意しました。これは、クラス画像（訳注：後述の Prior preservation セクションで解説する）を持たない1匹の犬の複数の画像（インスタンスプロンプトに対応するインスタンス画像）で構成されています。デフォルトのデータセットを使う場合は、デモ用ノートブックで推論を行いながら、プロンプト「 a photo of a Doppler dog 」を試してみてください。

ライセンス: MIT

ハイパーパラメータ

次に、カスタムデータセットで転移学習を行う場合、トレーニング用ハイパーパラメータのデフォルト値を必要に応じて変更することも可能です。hyperparameters.retrieve_default を呼び出すことで、これらのハイパーパラメータの Python 辞書とそのデフォルト値を取得し、必要に応じて更新した後、Estimator クラスへ渡すことができます。以下のコードをご覧ください。

from sagemaker import hyperparameters
# Retrieve the default hyper-parameters for fine-tuning the model
hyperparameters = hyperparameters.retrieve_default(
model_id=train_model_id, model_version=train_model_version
)

# [Optional] Override default hyperparameters with custom values
hyperparameters["max_steps"] = "400"

ファインチューニングアルゴリズムでは，以下のハイパーパラメータがサポートされています。

with_prior_preservation – Prior preservation の損失を追加するフラグ。Prior preservation は、オーバーフィッティングを回避するための正則化。(選択肢： ["True", "False"], デフォルト："False" )
num_class_images – Prior preservation のための最小クラス画像． with_prior_preservation = True で， class_data_dir に十分な画像が存在しない場合、 class_prompt で追加画像がサンプリングされます． (値：正の整数，デフォルト：100 )
epochs – ファインチューニングアルゴリズムが、トレーニングデータセットを何周するか。(値: 正の整数，デフォルト : 20 )
max_steps – 実行されるトレーニングステップの総数。None でない場合、epochs より優先されます。(値 : "None" または整数の文字列、デフォルト : "None")
batch_size – モデルの重みが更新される前に実行されるトレーニングサンプルの数。with_prior_preservation = True の場合のクラス画像生成時のバッチサイズと同じ。 (値：正の整数、デフォルト：1 )
learning_rate – トレーニングサンプルの各バッチを処理した後、モデルの重みが更新される割合。(値: 正の浮動小数点数，デフォルト : 2e-06)
prior_loss_weight – Prior preservation 損失の重み。(値: 正の浮動小数点数，デフォルト : 1.0 )
center_crop – 目的の解像度にリサイズする前に，画像をクロップするかどうか。 (選択肢：["True"/"False"]、デフォルト："False" )
lr_scheduler – Learning Rate Scheduler の種類。 (選択肢 : ["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"], デフォルト : "constant" ) 詳細については、Learning Rate Scheduler を参照してください。
adam_weight_decay – AdamW optimizer のバイアスおよび LayerNorm の重みを除くすべてのレイヤに適用する重み減衰（0でない場合）。(値 : 浮動小数点数, デフォルト : 1e-2 )
adam_beta1 – AdamW optimizer の beta1 ハイパーパラメータ（第一モーメント推定値の指数関数的減衰率）。(値：浮動小数点数，デフォルト：0.9 )
adam_beta2 – AdamW optimizer の beta2 ハイパーパラメータ（第一モーメント推定値の指数関数的減衰率）。(値 : 浮動小数点数, デフォルト : 0.999 )
adam_epsilon – AdamW optimizer の epsilon ハイパーパラメータ。通常 0 による除算を避けるために小さな値に設定されます。（値：浮動小数点数、デフォルト：1e-8 )
gradient_accumulation_steps – backward/update パスを実行する前に蓄積する更新ステップの数。 (値: 整数、デフォルト : 1 )
max_grad_norm – 最大勾配ノルム（勾配クリッピングのため）。(値 : 浮動小数点数, デフォルト : 1.0 )
seed – 学習で再現性のある結果を得るためにランダムな状態を固定する際のシード。 (値 : 整数, デフォルト : 0 )

ファインチューニングされたモデルのデプロイ

モデルのトレーニングが終了したら、モデルを直接永続的なリアルタイムエンドポイントにデプロイすることができます。必要な Docker Image URI とスクリプト URI を取得し、モデルをデプロイします。以下のコードをご覧ください。

inference_instance_type = "ml.g4dn.2xlarge"

# Retrieve the inference docker container uri
deploy_image_uri = image_uris.retrieve(
    region=None,
    framework=None,  # automatically inferred from model_id
    image_scope="inference",
    model_id=train_model_id,
    model_version=train_model_version,
    instance_type=inference_instance_type,
)

# Retrieve the inference script uri. This includes scripts for model loading, inference handling etc.
deploy_source_uri = script_uris.retrieve(
    model_id=train_model_id, model_version=train_model_version, script_scope="inference"
)

# Use the estimator from the previous step to deploy to a SageMaker endpoint
finetuned_predictor = sd_estimator.deploy(
    initial_instance_count=1,
    instance_type=inference_instance_type,
    entry_point="inference.py",  # entry point file in source_dir and present in deploy_source_uri
    image_uri=deploy_image_uri,
    source_dir=deploy_source_uri,
    endpoint_name=endpoint_name,
)

左はモデルのファインチューニングに使用した riobugger という猫のトレーニング画像です（max_steps = 400 以外のパラメータはデフォルト）。中央と右は、riobugger の海岸での画像と鉛筆のスケッチを予測するよう求められたときに、ファインチューニングされたモデルによって生成された画像です。

サポートされているパラメータや応答形式など、推論の詳細については、Generate images from text with the stable diffusion model on Amazon SageMaker JumpStart. を参照してください。

ファインチューニングの考慮点

Stable Diffusion モデルのトレーニングは、すぐにオーバーフィットしてしまう傾向があります。良質な画像を得るためには、トレーニングステップ数や学習率（Learning Rateなど、利用可能なトレーニングハイパーパラメータの間で良いバランスを見つける必要があります。このセクションでは、いくつかの実験結果を示し、これらのパラメータをどのように設定するかのガイダンスを提供します。

推奨事項

以下の推奨事項を検討します。

まず、質の良いトレーニング画像（4-20 枚）を用意します。人間の顔を学習させる場合は、より多くの画像が必要かもしれません。
犬や猫など、人間以外の被写体でトレーニングを行う場合は、200 ～ 400 ステップのトレーニングを行います。人間の顔を学習させる場合は、さらに多くのステップが必要になる場合があります。オーバーフィッティングが発生した場合は、ステップ数を減らしてください。アンダーフィット（ファインチューニングしたモデルが被写体の画像を生成できない）する場合は、ステップ数を増やします。
人間以外の顔を学習する場合、with_prior_preservation = False を設定しても性能に大きな影響はありません。人間の顔の場合は、with_prior_preservation=True を設定する必要があるかもしれません。
with_prior_preservation=True を設定する場合、ml.g5.2xlarge インスタンスタイプを使用します。
複数の被写体を連続して学習させる場合、被写体が非常に似ている場合（例えば、全て犬）、モデルは最後の被写体を保持し、それ以前の被写体を忘れてしまいます。被写体が異なる場合（例えば、最初に猫、次に犬）、モデルは両方の被写体を保持します。
学習率は低めに設定し、満足のいく結果が得られるまでステップ数を徐々に増やしていくことを推奨します。

トレーニングデータセット

ファインチューニングされたモデルの品質は、トレーニング画像の品質に直接影響されます。そのため、良い結果を得るためには高品質の画像を収集する必要があります。ぼやけた画像や低解像度の画像は、ファインチューニングモデルの品質に影響を与えます。以下の追加パラメータに留意してください。

トレーニング画像の数 – 最小4枚のトレーニング画像でモデルのファインチューニングを行うことができます。我々は最小4枚、最大 16 枚のトレーニングデータセットで実験を行いました。どちらの場合も、ファインチューニングによりモデルを被写体に適応させることができました。
データセットのフォーマット – 我々は .png, .jpg, .jpeg のフォーマットの画像でファインチューニングアルゴリズムをテストしました。他のフォーマットでも動作する可能性があります。
画像の解像度 – トレーニング画像はどのような解像度でも構いません。ファインチューニングアルゴリズムは、ファインチューニングを開始する前に、すべてのトレーニング画像をリサイズします。しかし、トレーニングイメージのトリミングやリサイズをよりコントロールしたい場合は、モデルの基本解像度（この例では、512 × 512 ピクセル）に自分でリサイズすることをお勧めします。

実験の設定

この記事の実験では、ファインチューニングを行う際、特に指定がない限りハイパーパラメータのデフォルト値を使用します。さらに、4つのデータセットのうち1つを使用します。

Dog1-8 – 8 枚の画像からなる Dog 1
Dog1-16 – 16 枚の画像からなる Dog 1
Dog2-4 – 4 枚の画像からなる Dog 2
Cat-8 – 8 枚の画像からなる猫

ごちゃごちゃしないように、各セクションではデータセットの代表的な画像1枚のみをデータセット名と一緒に表示しています。完全なトレーニングデータセットは、この投稿の実験データセットのセクションで見ることができます。

オーバーフィット

Stable Diffusion モデルは、数枚の画像でファインチューニングを行う際にオーバーフィットする傾向があります。そのため、epochs、max_epochs、learning_rate などのパラメータを慎重に選択する必要があります。ここでは、Dog1-16 のデータセットを使用しました。

モデルの性能を評価するために，ファインチューニングを行ったモデルを4つのタスクで評価します。

ファインチューニングしたモデルは、学習させたのと同じ設定で被写体（ドップラーという犬）の画像を生成することができるか？
- コメント – はい、できます。注目すべきは、モデルの性能はトレーニングステップの数によって向上することです。
ファインチューニングされたモデルは、トレーニングされた環境とは異なる環境での被写体の画像を生成できるか？例えば、海岸でのドップラーの画像を生成することができるか？
- コメント – はい、できます。モデルの性能は、ある時点まではトレーニングのステップ数に応じて向上することに留意する必要があります。しかし、モデルの学習が長すぎると、モデルがオーバーフィットする傾向があり、モデル性能が低下します。
ファインチューニングされたモデルは、学習対象が属するクラスの画像を生成できるのか？例えば、一般的な犬の画像を生成することができるか？
- コメント – トレーニングステップ数を増やすと、モデルはオーバーフィットし始めます。その結果、犬という一般的なクラスを忘れてしまい、被写体に関連した画像しか生成されなくなります。
ファインチューニングを行ったモデルは、トレーニングデータセットにないクラスや被写体の画像を生成することができるだろうか？例えば、猫の画像を生成することができるだろうか？
- コメント – トレーニングステップの数を増やすと、モデルはオーバーフィットし始めます。その結果、指定されたクラスに関係なく、被写体に関連する画像しか生成されなくなります。

（max_steps ハイパーパラメータを設定して）異なるステップ数でモデルをファインチューニングし、ファインチューニングした各モデルについて、以下の4つのプロンプト（左から右の順に示す）のそれぞれについて画像を生成しています。

“ドップラーという犬の写真”
“ビーチにいるドップラーという犬の写真”
“犬の写真”
“猫の写真”

以下の画像は、50 ステップで学習させたモデルによるものです。

以下の画像は、100 ステップで学習させたモデルによるものです。

以下の画像は、200 ステップで学習させたモデルによるものです。

以下の画像は、400 ステップで学習させたモデルによるものです。

最後に、以下の画像は、800 ステップで学習させたモデルによるものです。

複数のデータセットでトレーニングする

ファインチューニングを行う際、複数の被写体でファインチューニングを行い、ファインチューニング後のモデルで全ての被写体の画像を生成できるようにしたい場合があります。残念ながら、JumpStart は現在、1つの被写体でのトレーニングに限定されています。複数の被写体で同時にモデルのファインチューニングを行うことはできません。さらに、異なる被写体に対して順次モデルのファインチューニングを行うと、被写体が類似している場合、モデルが最初の被写体を忘れてしまうという問題があります。

本節では、以下のような実験を考えます。

被写体 A のモデルをファインチューニングする。
被写体 B に対して、ステップ 1 のモデルをファインチューニングする。
Step 2 の出力モデルを用いて、被写体 A と被写体 B の画像を生成する。

以下の実験では、以下のことが観察されます。

A が犬 1 で B が犬 2 の場合、Step 3 で生成された画像は全て犬 2 に似ている。
A が犬 2 で B が犬 1 の場合、Step 3 で生成された画像は全て犬 1 に似ている。
A が犬 1 で B が猫の場合、犬のプロンプトで生成された画像は犬 1、猫のプロンプトで生成された画像は猫に似ている。

犬 1、犬 2 の順に学習させる

Step1 では、犬 1 の画像 8 枚で 200 ステップのファインチューニングを行い、Step2 では、犬 2 の画像 4 枚でさらに 200 ステップのファインチューニングを行いました。

以下は、Step 2 終了時にファインチューニングしたモデルで、異なるプロンプトに対して生成した画像です。

犬 2、犬 1 の順に学習させる

Step 1 では、犬 2 の画像 4 枚で 200 ステップのファインチューニングを行い、Step 2 では、犬 1 の画像 8 枚でさらに 200 ステップのファインチューニングを行いました。

以下は、Step 2 終了時にファインチューニングしたモデルで、プロンプトを変えて生成した画像です。

犬と猫で学習させる

ステップ 1 では、8 枚の猫の画像で 200 ステップのファインチューニングを行い、次に犬 1 の画像を 8 枚使って、200 ステップでさらにモデルをファインチューニングしました。

以下は、Step 2 終了時にファインチューニングしたモデルで生成した画像です。猫に関するプロンプトが表示された画像は、ファインチューニングの Step 1 における猫のように見え、犬に関するプロンプトが表示された画像は、ファインチューニングの Step 2 における犬のように見えます。

Prior preservation

Prior preservation とは、学習しようとするクラスと同じ画像を追加で使用する手法です。例えば、学習データが特定の犬の画像で構成されている場合、Prior preservation では、一般的な犬のクラス画像を取り込みます。特定の犬に対する学習を行う際に、異なる犬の画像を表示することで、オーバーフィッティングを回避しようとする手法です。インスタンスプロンプトに存在する特定の犬を示すタグは、クラスプロンプトでは欠落しています。例えば、インスタンスプロンプトは “a photo of a riobugger cat” で、クラスプロンプトは “a photo of a cat” である可能性があります。ハイパーパラメータ with_prior_preservation = True を設定することで、Prior preservation を有効にすることができます。with_prior_preservation = True に設定すると、dataset_info.json に class_prompt を含める必要があり、利用可能な任意のクラス画像を含めることができます。with_prior_preservation = True と設定した場合のトレーニングデータセットの形式は以下の通りです。

Input – インスタンス画像、dataset_info.json、（オプション）class_data_dir ディレクトリを含むディレクトリ。以下の点に注意
- 画像は、 .png, .jpg, .jpeg のいずれかの形式である
- dataset_info.json は、{'instance_prompt':<<instance_prompt>>,'class_prompt':<<class_prompt>>} という形式である。
- class_data_dir ディレクトリにクラス画像があること。class_data_dir がない場合、あるいは class_data_dir に十分な画像がない場合、class_prompt で追加画像をサンプリングします。

猫や犬のようなデータセットでは、Prior preservation はファインチューニングされたモデルの性能に大きな影響を与えないため、回避することが可能です。しかし、顔に関する学習を行う場合は必要です。詳しくは、“Training Stable Diffusion with Dreambooth using Diffusers.” を参照してください。

インスタンスタイプ

Stable Diffusion モデルのファインチューニングには、GPU をサポートしたインスタンスによる計算の高速化が必要です。今回は ml.g4dn.2xlarge (16 GB CUDA memory, 1 GPU) と ml.g5.2xlarge (24 GB CUDA memory, 1 GPU) インスタンスでファインチューニングの実験を行いました。クラス画像を生成する際には必要なメモリがより多くなります。そのため、with_prior_preservation=True とした場合、ml.g4dn.2xlarge インスタンスではトレーニングで CUDA のメモリ不足の問題が発生するため、ml.g5.2xlarge インスタンスタイプの使用を推奨します。g5 インスタンスを現在利用できるリージョン（例：us-east–1, us-west-2）については SageMaker Pricing ページをご確認ください。JumpStart ファインチューニングスクリプトは現在シングル GPU を使用しているため、マルチ GPU インスタンスでファインチューニングを行っても性能は上がりません。各インスタンスタイプの詳細については、Amazon EC2 Instance Types を参照してください。

制限とバイアス

Stable Diffusion は画像生成において素晴らしい性能を発揮しますが、いくつかの制限とバイアスを抱えています。これらには以下のものが含まれますが、これらに限定されるものではありません。

トレーニングデータに特徴を持つ十分な画像が含まれていないため、モデルは正確な顔や手足を生成しない可能性があります。
このモデルは LAION-5B データセットでトレーニングされましたが、このデータセットにはアダルトコンテンツが含まれており、さらなる検討なしに製品として使用するには適さない可能性があります。
このモデルは英語のテキストでトレーニングされているため、英語以外の言語ではうまく機能しない可能性があります。
このモデルは、画像内のテキストをうまく生成できません。

制限とバイアスの詳細については、Stable Diffusion v2-1- ベースモデルカードを参照してください。事前学習済みモデルのこれらの制限は、ファインチューニングされたモデルにも引き継がれる可能性があります。

クリーンアップ

ノートブックの実行が終わったら、課金を確実に停止するために、プロセスで作成されたすべてのリソースを削除することを確認します。エンドポイントをクリーンアップするコードは、付属する「 Introduction to JumpStart – Text to Image example notebook 」で提供されています。

JumpStart で公開されているファインチューニング済みモデル

StabilityAI によってリリースされた Stable Diffusion モデルが素晴らしいパフォーマンスを持っていたとしても、それがトレーニングされた言語やドメインの面で制限があります。例えば、Stable Diffusion モデルは英語のテキストで学習されましたが、英語以外のテキストから画像を生成する必要がある場合があります。また、Stable Diffusion モデルは写真のようにリアルな画像を生成するようにトレーニングされていますが、アニメーションやアーティスティックな画像を生成する必要がある場合もあります。

JumpStart では、さまざまな言語とテーマを持つ 80 以上のモデルが公開されています。これらのモデルの多くは、StabilityAI によってリリースされた Stable Diffusion モデルからファインチューニングされたバージョンです。あなたのユースケースが、ファインチューニングされたモデルの1つと一致する場合、独自のデータセットを収集し、それをファインチューニングする必要はありません。Studio UI または使いやすい JumpStart API を使用して、これらのモデルの1つをデプロイするだけでよいのです。JumpStart で事前学習済みの Stable Diffusion モデルをデプロイするには、”Generate images from text with the stable diffusion model on Amazon SageMaker JumpStart” を参照してください。

JumpStart で利用可能なさまざまなモデルによって生成された画像の例を以下に示します。

これらのモデルは、JumpStart スクリプトや DreamBooth スクリプトを使用してファインチューニングされたものではないことに注意してください。一般に公開されているファインチューニング済みモデルの全リストとプロンプト例は、ここからダウンロードできます。

これらのモデルから生成された画像の例については、付録の「オープンソース化されたファインチューニングモデル」のセクションを参照してください。

まとめ

この投稿では、テキストから画像を生成する Stable Diffusion モデルをファインチューニングし、JumpStart を使用してそれを展開する方法を示しました。さらに、モデルをファインチューニングする際に考慮すべき点と、ファインチューニングされたモデルの性能にどのような影響を与えるかについて説明しました。また、JumpStart で利用できる 80 以上のすぐに使えるファインチューニング済みモデルについても説明しました。このデモのすべてのステップを含む完全なコードについては、「Introduction to JumpStart – Text to Image」を参照してください。ぜひこのソリューションを試してみて、コメントをお寄せください。

Stable Diffusion モデルと DreamBooth のファインチューニングについては、次のリソースを参照してください。

JumpStart の詳細については、以下のブログ記事をご覧ください。