無条件画像生成

無条件画像生成モデルは、学習時にテキストや画像による条件付けを行いません。学習データの分布に類似した画像のみを生成します。

このガイドでは、train_unconditional.py トレーニングスクリプトを詳しく見ていき、その仕組みを理解した上で、ご自身のユースケースに合わせてどのようにカスタマイズできるかを解説します。

スクリプトを実行する前に、ソースからライブラリをインストールしていることを確認してください。

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

次に、トレーニングスクリプトを含むサンプルフォルダに移動し、必要な依存関係をインストールします。

cd examples/unconditional_image_generation
pip install -r requirements.txt

🤗 Accelerate は、複数 GPU/TPU での学習や混合精度学習を支援するためのライブラリです。ハードウェアや環境に基づいて、学習環境を自動的に設定します。詳細については、🤗 Accelerate のクイックツアーをご覧ください。

🤗 Accelerate の環境を初期化します:

accelerate config

設定を選択せずにデフォルトの🤗 Accelerate環境を設定するには:

accelerate config default

または、ノートブックのような対話型シェルをサポートしていない環境では、以下を使用できます:

from accelerate.utils import write_basic_config

write_basic_config()

最後に、独自のデータセットでモデルをトレーニングしたい場合は、トレーニング用のデータセットを作成するガイドを参照して、トレーニングスクリプトと互換性のあるデータセットの作成方法を確認してください。

スクリプトのパラメータ

以下のセクションでは、スクリプトの修正方法を理解するために重要な部分に焦点を当てていますが、スクリプトのすべての側面を詳細に説明するものではありません。さらに詳しく知りたい場合は、スクリプトをぜひお読みいただき、ご質問やご懸念事項があればお知らせください。

トレーニングスクリプトには、トレーニング実行をカスタマイズするための多くのパラメータが用意されています。すべてのパラメータとその説明は、parse_args()関数に記載されています。各パラメータには、トレーニングバッチサイズや学習率などのデフォルト値が設定されていますが、必要に応じてトレーニングコマンドで独自の値を設定することもできます。

例えば、bf16 形式の混合精度でトレーニングを高速化するには、トレーニングコマンドに --mixed_precision パラメータを追加します。

accelerate launch train_unconditional.py \
  --mixed_precision="bf16"

指定すべき基本的かつ重要なパラメータには、以下が含まれます。

--dataset_name: Hub上のデータセット名、または学習に使用するデータセットのローカルパス
--output_dir: 学習済みモデルを保存する場所
--push_to_hub: 学習したモデルを Hub にプッシュするかどうか
--checkpointing_steps: モデルの学習中にチェックポイントを保存する頻度です。学習が中断された場合に有用で、学習コマンドに--resume_from_checkpointを追加することで、そのチェックポイントから学習を再開できます。

データセットを用意するだけで、それ以外のすべては学習スクリプトが処理します!

学習スクリプト

データセットの前処理と学習ループのコードは、main()関数にあります。学習スクリプトを調整する必要がある場合は、ここで変更を加える必要があります。

モデル設定を指定しない場合、train_unconditional スクリプトは UNet2DModel を初期化します。必要に応じて、ここで UNet を設定できます。

model = UNet2DModel(
    sample_size=args.resolution,
    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",
    ),
)

次に、スクリプトはschedulerとoptimizerを初期化します。

# Initialize the scheduler
accepts_prediction_type = "prediction_type" in set(inspect.signature(DDPMScheduler.__init__).parameters.keys())
if accepts_prediction_type:
    noise_scheduler = DDPMScheduler(
        num_train_timesteps=args.ddpm_num_steps,
        beta_schedule=args.ddpm_beta_schedule,
        prediction_type=args.prediction_type,
    )
else:
    noise_scheduler = DDPMScheduler(num_train_timesteps=args.ddpm_num_steps, beta_schedule=args.ddpm_beta_schedule)

# Initialize the optimizer
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

その後、データセットを読み込み、前処理する方法を指定できます:

dataset = load_dataset("imagefolder", data_dir=args.train_data_dir, cache_dir=args.cache_dir, split="train")

augmentations = transforms.Compose(
    [
        transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
        transforms.CenterCrop(args.resolution) if args.center_crop else transforms.RandomCrop(args.resolution),
        transforms.RandomHorizontalFlip() if args.random_flip else transforms.Lambda(lambda x: x),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5]),
    ]
)

最後に、training loopは、画像にノイズを追加すること、ノイズ残差を予測すること、損失を計算すること、指定されたステップでチェックポイントを保存すること、モデルを保存して Hub にプッシュすることなど、その他のすべての処理を担当します。training loop の動作について詳しく知りたい場合は、デノイジングプロセスの基本的なパターンを解説している Understanding pipelines, models and schedulers チュートリアルをご覧ください。

スクリプトを実行します

すべての変更が完了したか、デフォルトの設定で問題ない場合は、トレーニングスクリプトの起動準備が整いました! 🚀

完全なトレーニングの実行には、4xV100 GPUで2時間かかります。

単一GPU

マルチGPU

トレーニングスクリプトは、リポジトリ内にチェックポイントファイルを作成して保存します。これで、トレーニング済みのモデルを読み込んで推論に使用できます。

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("anton-l/ddpm-butterflies-128").to("cuda")  # or "mps", "xpu", "cpu"
image = pipeline().images[0]

Update on GitHub

Diffusers

無条件画像生成

スクリプトのパラメータ

学習スクリプト

スクリプトを実行します