Entraîner un modèle de diffusion

La génération d’images non conditionnée est une application populaire des modèles de diffusion qui produit des images ressemblant à celles du jeu de données utilisé pour l’entraînement. En général, les meilleurs résultats sont obtenus en affinant un modèle préentraîné sur un jeu de données spécifique. Vous pouvez trouver de nombreux de ces points de contrôle sur le Hub, mais si vous n’en trouvez pas qui vous convient, vous pouvez toujours entraîner le vôtre !

Ce tutoriel vous apprendra à entraîner un UNet2DModel à partir de zéro sur un sous-ensemble du jeu de données Smithsonian Butterflies afin de générer vos propres papillons 🦋.

💡 Ce tutoriel d'entraînement est basé sur le notebook Training with 🧨 Diffusers. Pour plus de détails et de contexte sur les modèles de diffusion, comme leur fonctionnement, consultez le notebook !

Avant de commencer, assurez-vous d'avoir installé 🤗 Datasets pour charger et prétraiter les jeux de données d'images, ainsi que 🤗 Accelerate, afin de simplifier l'entraînement sur n'importe quel nombre de GPU. La commande suivante installera également TensorBoard pour visualiser les métriques d'entraînement (vous pouvez également utiliser Weights & Biases pour suivre votre entraînement).

# uncomment to install the necessary libraries in Colab
#!pip install diffusers[training]

Nous vous encourageons à partager votre modèle avec la communauté. Pour ce faire, vous devrez vous connecter à votre compte Hugging Face (en créer un ici si vous n’en avez pas encore !). Vous pouvez vous connecter depuis un notebook et saisir votre jeton lorsque cela vous est demandé. Assurez-vous que votre jeton dispose du rôle d’écriture.

>>> from huggingface_hub import notebook_login

>>> notebook_login()

Ou connectez-vous depuis le terminal :

hf auth login

Les points de contrôle du modèle étant assez volumineux, installez Git-LFS pour versionner ces fichiers de grande taille :

!sudo apt -qq install git-lfs
!git config --global credential.helper store

Configuration de l'entraînement

Pour plus de commodité, créez une classe TrainingConfig contenant les hyperparamètres d'entraînement (n'hésitez pas à les ajuster) :

>>> from dataclasses import dataclass

>>> @dataclass
... class TrainingConfig:
...     image_size = 128  # the generated image resolution
...     train_batch_size = 16
...     eval_batch_size = 16  # how many images to sample during evaluation
...     num_epochs = 50
...     gradient_accumulation_steps = 1
...     learning_rate = 1e-4
...     lr_warmup_steps = 500
...     save_image_epochs = 10
...     save_model_epochs = 30
...     mixed_precision = "fp16"  # `no` for float32, `fp16` for automatic mixed precision
...     output_dir = "ddpm-butterflies-128"  # the model name locally and on the HF Hub

...     push_to_hub = True  # whether to upload the saved model to the HF Hub
...     hub_model_id = "<your-username>/<my-awesome-model>"  # the name of the repository to create on the HF Hub
...     hub_private_repo = None
...     overwrite_output_dir = True  # overwrite the old model when re-running the notebook
...     seed = 0


>>> config = TrainingConfig()

Charger le jeu de données

Vous pouvez facilement charger le jeu de données Smithsonian Butterflies avec la bibliothèque 🤗 Datasets :

>>> from datasets import load_dataset

>>> config.dataset_name = "huggan/smithsonian_butterflies_subset"
>>> dataset = load_dataset(config.dataset_name, split="train")

💡 Vous pouvez trouver des jeux de données supplémentaires dans le HugGan Community Event ou utiliser votre propre jeu de données en créant un ImageFolder local. Définissez config.dataset_name sur l'identifiant du dépôt du jeu de données s'il provient du HugGan Community Event, ou sur imagefolder si vous utilisez vos propres images.

🤗 Datasets utilise la fonctionnalité Image pour décoder automatiquement les données d'image et les charger en tant que PIL.Image, ce qui nous permet de les visualiser :

>>> 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)
...     axs[i].set_axis_off()
>>> fig.show()

Les images ont toutes des tailles différentes, il faut donc les prétraiter d'abord :

Resize modifie la taille de l'image pour qu'elle corresponde à celle définie dans config.image_size.
RandomHorizontalFlip augmente le jeu de données en inversant aléatoirement les images.
Normalize est important pour remettre à l'échelle les valeurs des pixels dans l'intervalle [-1, 1], ce que le modèle attend.

>>> from torchvision import transforms

>>> preprocess = transforms.Compose(
...     [
...         transforms.Resize((config.image_size, config.image_size)),
...         transforms.RandomHorizontalFlip(),
...         transforms.ToTensor(),
...         transforms.Normalize([0.5], [0.5]),
...     ]
... )

Utilisez la méthode set_transform de 🤗 Datasets pour appliquer la fonction preprocess à la volée pendant l’entraînement :

>>> def transform(examples):
...     images = [preprocess(image.convert("RGB")) for image in examples["image"]]
...     return {"images": images}


>>> dataset.set_transform(transform)

N’hésitez pas à visualiser à nouveau les images pour confirmer qu’elles ont été redimensionnées. Vous êtes maintenant prêt à encapsuler le jeu de données dans un DataLoader pour l’entraînement !

>>> import torch

>>> train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True)

Créer un UNet2DModel

Les modèles préentraînés dans 🧨 Diffusers sont facilement créés à partir de leur classe de modèle avec les paramètres souhaités. Par exemple, pour créer un UNet2DModel :

>>> from diffusers import UNet2DModel

>>> model = UNet2DModel(
...     sample_size=config.image_size,  # the target image resolution
...     in_channels=3,  # the number of input channels, 3 for RGB images
...     out_channels=3,  # the number of output channels
...     layers_per_block=2,  # how many ResNet layers to use per UNet block
...     block_out_channels=(128, 128, 256, 256, 512, 512),  # the number of output channels for each UNet block
...     down_block_types=(
...         "DownBlock2D",  # a regular ResNet downsampling block
...         "DownBlock2D",
...         "DownBlock2D",
...         "DownBlock2D",
...         "AttnDownBlock2D",  # a ResNet downsampling block with spatial self-attention
...         "DownBlock2D",
...     ),
...     up_block_types=(
...         "UpBlock2D",  # a regular ResNet upsampling block
...         "AttnUpBlock2D",  # a ResNet upsampling block with spatial self-attention
...         "UpBlock2D",
...         "UpBlock2D",
...         "UpBlock2D",
...         "UpBlock2D",
...     ),
... )

Il est souvent judicieux de vérifier rapidement que la forme de l'image d'exemple correspond à la forme de la sortie du modèle :

>>> sample_image = dataset[0]["images"].unsqueeze(0)
>>> print("Input shape:", sample_image.shape)
Input shape: torch.Size([1, 3, 128, 128])

>>> print("Output shape:", model(sample_image, timestep=0).sample.shape)
Output shape: torch.Size([1, 3, 128, 128])

Parfait ! Ensuite, vous aurez besoin d'un planificateur pour ajouter du bruit à l'image.

Créer un planificateur

Le comportement du planificateur varie selon que vous utilisez le modèle pour l’entraînement ou pour l’inférence. Lors de l’inférence, le planificateur génère l’image à partir du bruit. Lors de l’entraînement, le planificateur prend une sortie du modèle – ou un échantillon – à un point précis du processus de diffusion et applique du bruit à l’image selon un plan de bruit et une règle de mise à jour.

Regardons le DDPMScheduler et utilisons la méthode add_noise pour ajouter du bruit aléatoire à l’image sample_image de tout à l’heure :

>>> import torch
>>> from PIL import Image
>>> from diffusers import DDPMScheduler

>>> noise_scheduler = DDPMScheduler(num_train_timesteps=1000)
>>> noise = torch.randn(sample_image.shape)
>>> timesteps = torch.LongTensor([50])
>>> noisy_image = noise_scheduler.add_noise(sample_image, noise, timesteps)

>>> Image.fromarray(((noisy_image.permute(0, 2, 3, 1) + 1.0) * 127.5).type(torch.uint8).numpy()[0])

L'objectif d'entraînement du modèle est de prédire le bruit ajouté à l'image. La perte à cette étape peut être calculée par :

>>> import torch.nn.functional as F

>>> noise_pred = model(noisy_image, timesteps).sample
>>> loss = F.mse_loss(noise_pred, noise)

Entraîner le modèle

À ce stade, vous disposez de la plupart des éléments nécessaires pour commencer l’entraînement du modèle, et il ne reste plus qu’à assembler le tout.

Tout d'abord, vous aurez besoin d'un optimiseur et d'un planificateur de taux d'apprentissage :

>>> from diffusers.optimization import get_cosine_schedule_with_warmup

>>> optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
>>> lr_scheduler = get_cosine_schedule_with_warmup(
...     optimizer=optimizer,
...     num_warmup_steps=config.lr_warmup_steps,
...     num_training_steps=(len(train_dataloader) * config.num_epochs),
... )

Ensuite, vous aurez besoin d'une méthode pour évaluer le modèle. Pour l'évaluation, vous pouvez utiliser le DDPMPipeline pour générer un lot d'images d'exemple et l'enregistrer sous forme de grille :

>>> from diffusers import DDPMPipeline
>>> from diffusers.utils import make_image_grid
>>> import os

>>> def evaluate(config, epoch, pipeline):
...     # Sample some images from random noise (this is the backward diffusion process).
...     # The default pipeline output type is `List[PIL.Image]`
...     images = pipeline(
...         batch_size=config.eval_batch_size,
...         generator=torch.Generator(device='cpu').manual_seed(config.seed), # Use a separate torch generator to avoid rewinding the random state of the main training loop
...     ).images

...     # Make a grid out of the images
...     image_grid = make_image_grid(images, rows=4, cols=4)

...     # Save the images
...     test_dir = os.path.join(config.output_dir, "samples")
...     os.makedirs(test_dir, exist_ok=True)
...     image_grid.save(f"{test_dir}/{epoch:04d}.png")

Vous pouvez maintenant assembler tous ces composants dans une boucle d’entraînement avec 🤗 Accelerate pour bénéficier d’un journalisation TensorBoard simplifiée, d’une accumulation de gradients et d’une formation en précision mixte. Pour téléverser le modèle sur le Hub, écrivez une fonction pour récupérer le nom de votre dépôt et ses informations, puis poussez-le sur le Hub.

💡 La boucle d’entraînement ci-dessous peut sembler intimidante et longue, mais cela en vaudra la peine plus tard lorsque vous lancerez votre entraînement en une seule ligne de code ! Si vous ne pouvez pas attendre et souhaitez commencer à générer des images, n’hésitez pas à copier et exécuter le code ci-dessous. Vous pourrez toujours revenir plus tard pour examiner la boucle d’entraînement plus en détail, par exemple pendant que vous attendez que votre modèle termine son entraînement. 🤗

>>> from accelerate import Accelerator
>>> from huggingface_hub import create_repo, upload_folder
>>> from tqdm.auto import tqdm
>>> from pathlib import Path
>>> import os

>>> def train_loop(config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler):
...     # Initialize accelerator and tensorboard logging
...     accelerator = Accelerator(
...         mixed_precision=config.mixed_precision,
...         gradient_accumulation_steps=config.gradient_accumulation_steps,
...         log_with="tensorboard",
...         project_dir=os.path.join(config.output_dir, "logs"),
...     )
...     if accelerator.is_main_process:
...         if config.output_dir is not None:
...             os.makedirs(config.output_dir, exist_ok=True)
...         if config.push_to_hub:
...             repo_id = create_repo(
...                 repo_id=config.hub_model_id or Path(config.output_dir).name, exist_ok=True
...             ).repo_id
...         accelerator.init_trackers("train_example")

...     # Prepare everything
...     # There is no specific order to remember, you just need to unpack the
...     # objects in the same order you gave them to the prepare method.
...     model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
...         model, optimizer, train_dataloader, lr_scheduler
...     )

...     global_step = 0

...     # Now you train the model
...     for epoch in range(config.num_epochs):
...         progress_bar = tqdm(total=len(train_dataloader), disable=not accelerator.is_local_main_process)
...         progress_bar.set_description(f"Epoch {epoch}")

...         for step, batch in enumerate(train_dataloader):
...             clean_images = batch["images"]
...             # Sample noise to add to the images
...             noise = torch.randn(clean_images.shape, device=clean_images.device)
...             bs = clean_images.shape[0]

...             # Sample a random timestep for each image
...             timesteps = torch.randint(
...                 0, noise_scheduler.config.num_train_timesteps, (bs,), device=clean_images.device,
...                 dtype=torch.int64
...             )

...             # Add noise to the clean images according to the noise magnitude at each timestep
...             # (this is the forward diffusion process)
...             noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)

...             with accelerator.accumulate(model):
...                 # Predict the noise residual
...                 noise_pred = model(noisy_images, timesteps, return_dict=False)[0]
...                 loss = F.mse_loss(noise_pred, noise)
...                 accelerator.backward(loss)

...                 if accelerator.sync_gradients:
...                     accelerator.clip_grad_norm_(model.parameters(), 1.0)
...                 optimizer.step()
...                 lr_scheduler.step()
...                 optimizer.zero_grad()

...             progress_bar.update(1)
...             logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0], "step": global_step}
...             progress_bar.set_postfix(**logs)
...             accelerator.log(logs, step=global_step)
...             global_step += 1

...         # After each epoch you optionally sample some demo images with evaluate() and save the model
...         if accelerator.is_main_process:
...             pipeline = DDPMPipeline(unet=accelerator.unwrap_model(model), scheduler=noise_scheduler)

...             if (epoch + 1) % config.save_image_epochs == 0 or epoch == config.num_epochs - 1:
...                 evaluate(config, epoch, pipeline)

...             if (epoch + 1) % config.save_model_epochs == 0 or epoch == config.num_epochs - 1:
...                 if config.push_to_hub:
...                     upload_folder(
...                         repo_id=repo_id,
...                         folder_path=config.output_dir,
...                         commit_message=f"Epoch {epoch}",
...                         ignore_patterns=["step_*", "epoch_*"],
...                     )
...                 else:
...                     pipeline.save_pretrained(config.output_dir)

Ouf, c’était pas mal de code ! Mais vous êtes enfin prêt à lancer l’entraînement avec la fonction notebook_launcher de 🤗 Accelerate. Transmettez à la fonction la boucle d’entraînement, tous les arguments d’entraînement et le nombre de processus (vous pouvez modifier cette valeur pour correspondre au nombre de GPU disponibles) à utiliser pour l’entraînement :

>>> from accelerate import notebook_launcher

>>> args = (config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler)

>>> notebook_launcher(train_loop, args, num_processes=1)

Une fois l'entraînement terminé, jetez un œil aux images 🦋 finales 🦋 générées par votre modèle de diffusion !

>>> import glob

>>> sample_images = sorted(glob.glob(f"{config.output_dir}/samples/*.png"))
>>> Image.open(sample_images[-1])

Étapes suivantes

La génération d’images non conditionnelle est un exemple de tâche qui peut être entraînée. Vous pouvez explorer d’autres tâches et techniques d’entraînement en consultant la page 🧨 Diffusers Training Examples. Voici quelques exemples de ce que vous pouvez apprendre :

Textual Inversion, un algorithme qui apprend à un modèle un concept visuel spécifique et l’intègre dans l’image générée.
DreamBooth, une technique permettant de générer des images personnalisées d’un sujet à partir de plusieurs images d’entrée de ce sujet.
Guide pour l’ajustement fin d’un modèle Stable Diffusion sur votre propre jeu de données.
Guide pour utiliser LoRA, une technique de fine-tuning économe en mémoire permettant d'entraîner plus rapidement des modèles très volumineux.

Update on GitHub