Génération d'images non conditionnelle

Les modèles de génération d'images non conditionnés ne sont pas conditionnés par du texte ou des images lors de l'entraînement. Ils ne génèrent que des images qui ressemblent à la distribution de leurs données d'entraînement.

Ce guide examine le script d’entraînement train_unconditional.py afin de vous familiariser avec celui-ci et de vous montrer comment l’adapter à votre propre cas d’usage.

Avant d'exécuter le script, assurez-vous d'installer la bibliothèque à partir du code source :

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

Puis accédez au dossier d'exemple contenant le script d'entraînement et installez les dépendances requises :

cd examples/unconditional_image_generation
pip install -r requirements.txt

🤗 Accelerate est une bibliothèque qui vous aide à entraîner sur plusieurs GPU/TPU ou avec une précision mixte. Elle configure automatiquement votre configuration d’entraînement en fonction de votre matériel et de votre environnement. Consultez le tour rapide de 🤗 Accelerate pour en savoir plus.

Initialiser un environnement 🤗 Accelerate :

accelerate config

Pour configurer un environnement 🤗 Accelerate par défaut sans choisir de configuration :

accelerate config default

Sinon, si votre environnement ne prend pas en charge un shell interactif comme un notebook, vous pouvez utiliser :

from accelerate.utils import write_basic_config

write_basic_config()

Enfin, si vous souhaitez entraîner un modèle sur votre propre jeu de données, consultez le guide Create a dataset for training pour apprendre à créer un jeu de données compatible avec le script d'entraînement.

Paramètres du script

Les sections suivantes mettent en évidence les parties du script d’entraînement importantes pour comprendre comment le modifier, mais elles ne couvrent pas tous les aspects du script en détail. Si vous souhaitez en savoir plus, n’hésitez pas à lire le script et à nous faire part de vos questions ou préoccupations.

Le script d’entraînement fournit de nombreux paramètres pour vous aider à personnaliser votre session d’entraînement. Tous les paramètres et leurs descriptions se trouvent dans la fonction parse_args(). Elle fournit des valeurs par défaut pour chaque paramètre, telles que la taille du lot d’entraînement et le taux d’apprentissage, mais vous pouvez également définir vos propres valeurs dans la commande d’entraînement si vous le souhaitez.

Par exemple, pour accélérer l’entraînement avec la précision mixte en utilisant le format bf16, ajoutez le paramètre --mixed_precision à la commande d’entraînement :

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

Parmi les paramètres de base et importants à spécifier, on trouve :

• --dataset_name : le nom du jeu de données sur le Hub ou un chemin local vers le jeu de données sur lequel entraîner le modèle
• --output_dir : l’endroit où enregistrer le modèle entraîné
• --push_to_hub : indique s'il faut pousser le modèle entraîné sur le Hub
• --checkpointing_steps : fréquence de sauvegarde d'un point de contrôle pendant l'entraînement du modèle ; cela est utile si l'entraînement est interrompu, vous pouvez reprendre l'entraînement à partir de ce point de contrôle en ajoutant --resume_from_checkpoint à votre commande d'entraînement

Apportez votre jeu de données et laissez le script d'entraînement s'occuper de tout le reste !

Script d'entraînement

Le code de prétraitement du jeu de données et de la boucle d’entraînement se trouve dans la fonction main(). Si vous devez adapter le script d’entraînement, c’est ici que vous devrez effectuer vos modifications.

Le script train_unconditional initialise un UNet2DModel si vous ne fournissez pas de configuration de modèle. Vous pouvez configurer le UNet ici si vous le souhaitez :

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",
    ),
)

Ensuite, le script initialise un scheduler et un 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,
)

Ensuite, il charge un jeu de données et vous pouvez spécifier comment le prétraiter :

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]),
    ]
)

Enfin, la boucle d’entraînement gère tout le reste, comme l’ajout de bruit aux images, la prédiction du résidu de bruit, le calcul de la perte, l’enregistrement des points de contrôle aux étapes spécifiées, ainsi que l’enregistrement et la publication du modèle sur le Hub. Si vous souhaitez en savoir plus sur le fonctionnement de la boucle d’entraînement, consultez le tutoriel Comprendre les pipelines, les modèles et les planificateurs, qui détaille le schéma de base du processus de débruitage.

Lancer le script

Une fois que vous avez apporté toutes vos modifications ou que vous êtes satisfait de la configuration par défaut, vous êtes prêt à lancer le script d’entraînement ! 🚀

Une session d'entraînement complète prend 2 heures sur 4 GPU V100.

accelerate launch train_unconditional.py \
  --dataset_name="huggan/flowers-102-categories" \
  --output_dir="ddpm-ema-flowers-64" \
  --mixed_precision="fp16" \
  --push_to_hub

accelerate launch --multi_gpu train_unconditional.py \
  --dataset_name="huggan/flowers-102-categories" \
  --output_dir="ddpm-ema-flowers-64" \
  --mixed_precision="fp16" \
  --push_to_hub

Le script d’entraînement crée et enregistre un fichier de point de contrôle dans votre dépôt. Vous pouvez maintenant charger et utiliser votre modèle entraîné pour l’inférence :

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]