# Entrenar un modelo de difusión

La generación de imágenes incondicional es una aplicación popular de los modelos de difusión que genera imágenes similares a las del conjunto de datos utilizado para el entrenamiento. Por lo general, los mejores resultados se obtienen al realizar un ajuste fino de un modelo preentrenado en un conjunto de datos específico. Puedes encontrar muchos de estos puntos de control en el [Hub](https://huggingface.co/search/full-text?q=unconditional-image-generation&type=model), pero si no encuentras uno que te guste, ¡siempre puedes entrenar el tuyo!

En este tutorial aprenderás a entrenar un [UNet2DModel](/docs/diffusers/main/es/api/models/unet2d#diffusers.UNet2DModel) desde cero en un subconjunto del conjunto de datos [Smithsonian Butterflies](https://huggingface.co/datasets/huggan/smithsonian_butterflies_subset) para generar tus propias mariposas 🦋.

> [!TIP]
> 💡 Este tutorial de entrenamiento se basa en el cuaderno [Training with 🧨 Diffusers](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/training_example.ipynb). Para obtener más detalles y contexto sobre los modelos de difusión, como su funcionamiento, ¡consulta el cuaderno!

Antes de comenzar, asegúrate de tener instalado 🤗 Datasets para cargar y preprocesar conjuntos de datos de imágenes, y 🤗 Accelerate, para simplificar el entrenamiento en cualquier número de GPUs. El siguiente comando también instalará [TensorBoard](https://www.tensorflow.org/tensorboard) para visualizar las métricas de entrenamiento (también puedes usar [Weights & Biases](https://docs.wandb.ai/) para hacer seguimiento de tu entrenamiento).

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

Te animamos a compartir tu modelo con la comunidad. Para hacerlo, necesitas iniciar sesión en tu cuenta de Hugging Face (créala [aquí](https://hf.co/join) si aún no tienes una). Puedes iniciar sesión desde un notebook e introducir tu token cuando se te solicite. Asegúrate de que tu token tenga el rol de escritura.

```py
>>> from huggingface_hub import notebook_login

>>> notebook_login()
```

O inicia sesión desde la terminal:

```bash
hf auth login
```

Dado que los puntos de control del modelo son bastante grandes, instale [Git-LFS](https://git-lfs.com/) para versionar estos archivos grandes:

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

## Configuración de entrenamiento

Para mayor comodidad, crea una clase `TrainingConfig` que contenga los hiperparámetros de entrenamiento (puedes ajustarlos a tu conveniencia):

```py
>>> 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()
```

## Cargar el conjunto de datos

Puedes cargar fácilmente el conjunto de datos [Smithsonian Butterflies](https://huggingface.co/datasets/huggan/smithsonian_butterflies_subset) con la biblioteca 🤗 Datasets:

```py
>>> from datasets import load_dataset

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

> [!TIP]
> 💡 Puedes encontrar conjuntos de datos adicionales del [HugGan Community Event](https://huggingface.co/huggan) o puedes usar tu propio conjunto de datos creando un [`ImageFolder`](https://huggingface.co/docs/datasets/image_dataset#imagefolder) local. Establece `config.dataset_name` en el id del repositorio del conjunto de datos si proviene del HugGan Community Event, o en `imagefolder` si estás usando tus propias imágenes.

🤗 Datasets utiliza la función [Image](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Image) para decodificar automáticamente los datos de la imagen y cargarla como [`PIL.Image`](https://pillow.readthedocs.io/en/stable/reference/Image.html), lo que nos permite visualizarla:

```py
>>> 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()
```

    

Sin embargo, las imágenes tienen todos tamaños diferentes, por lo que primero deberá preprocesarlas:

* `Resize` cambia el tamaño de la imagen al definido en `config.image_size`.
* `RandomHorizontalFlip` aumenta el conjunto de datos mediante el espejado aleatorio de las imágenes.
* `Normalize` es importante para reescalar los valores de píxeles al rango [-1, 1], que es lo que el modelo espera.

```py
>>> 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]),
...     ]
... )
```

Usa el método [set_transform](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset.set_transform) de 🤗 Datasets para aplicar la función `preprocess` sobre la marcha durante el entrenamiento:

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

>>> dataset.set_transform(transform)
```

Puedes visualizar las imágenes de nuevo para confirmar que se han redimensionado. Ahora estás listo para envolver el conjunto de datos en un [DataLoader](https://pytorch.org/docs/stable/data#torch.utils.data.DataLoader) para el entrenamiento.

```py
>>> import torch

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

## Crear un UNet2DModel

Los modelos preentrenados en 🧨 Diffusers se crean fácilmente a partir de su clase de modelo con los parámetros que desee. Por ejemplo, para crear un [UNet2DModel](/docs/diffusers/main/es/api/models/unet2d#diffusers.UNet2DModel):

```py
>>> 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",
...     ),
... )
```

A menudo es buena idea comprobar rápidamente que la forma de la imagen de muestra coincide con la forma de la salida del modelo:

```py
>>> 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])
```

¡Genial! A continuación, necesitarás un programador para añadir algo de ruido a la imagen.

## Crear un programador

El programador se comporta de manera diferente dependiendo de si se utiliza el modelo para entrenamiento o inferencia. Durante la inferencia, el programador genera la imagen a partir del ruido. Durante el entrenamiento, el programador toma una salida del modelo, o una muestra, de un punto específico del proceso de difusión y aplica ruido a la imagen según un *programa de ruido* y una *regla de actualización*.

Echemos un vistazo al [DDPMScheduler](/docs/diffusers/main/es/api/schedulers/ddpm#diffusers.DDPMScheduler) y usemos el método `add_noise` para añadir algo de ruido aleatorio a la `sample_image` de antes:

```py
>>> 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])
```

    

El objetivo de entrenamiento del modelo es predecir el ruido añadido a la imagen. La pérdida en este paso se puede calcular mediante:

```py
>>> import torch.nn.functional as F

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

## Entrenar el modelo

A estas alturas, ya tienes la mayoría de las piezas para empezar a entrenar el modelo y lo único que queda es ponerlo todo en marcha.

Primero, necesitarás un optimizador y un programador de tasa de aprendizaje:

```py
>>> 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),
... )
```

Luego, necesitarás una forma de evaluar el modelo. Para la evaluación, puedes usar [DDPMPipeline](/docs/diffusers/main/es/api/pipelines/ddpm#diffusers.DDPMPipeline) para generar un lote de imágenes de muestra y guardarlo como una cuadrícula:

```py
>>> 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")
```

Ahora puedes agrupar todos estos componentes en un bucle de entrenamiento con 🤗 Accelerate para un registro sencillo en TensorBoard, acumulación de gradientes y entrenamiento con precisión mixta. Para subir el modelo al Hub, escribe una función para obtener el nombre y la información de tu repositorio y luego súbelo al Hub.

> [!TIP]
> 💡 El bucle de entrenamiento de abajo puede parecer intimidante y largo, pero valdrá la pena más adelante cuando inicies tu entrenamiento con solo una línea de código. Si no puedes esperar y quieres empezar a generar imágenes, no dudes en copiar y ejecutar el código de abajo. Siempre puedes volver a examinar el bucle de entrenamiento con más detalle más tarde, por ejemplo, mientras esperas a que tu modelo termine de entrenar. 🤗

```py
>>> 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)
```

¡Uf, ha sido bastante código! Pero por fin estás listo para lanzar el entrenamiento con la función [notebook_launcher](https://huggingface.co/docs/accelerate/main/en/package_reference/launchers#accelerate.notebook_launcher) de 🤗 Accelerate. Pasa a la función el bucle de entrenamiento, todos los argumentos de entrenamiento y el número de procesos (puedes cambiar este valor por el número de GPUs disponibles) que se usarán para el entrenamiento:

```py
>>> from accelerate import notebook_launcher

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

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

Una vez que el entrenamiento esté completo, ¡echa un vistazo a las imágenes 🦋 finales 🦋 generadas por tu modelo de difusión!

```py
>>> import glob

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

    

## Próximos pasos

La generación de imágenes incondicional es un ejemplo de una tarea que se puede entrenar. Puede explorar otras tareas y técnicas de entrenamiento visitando la página [🧨 Diffusers Training Examples](../training/overview). A continuación se muestran algunos ejemplos de lo que puede aprender:

* [Textual Inversion](../training/text_inversion), un algoritmo que enseña a un modelo un concepto visual específico y lo integra en la imagen generada.
* [DreamBooth](../training/dreambooth), una técnica para generar imágenes personalizadas de un sujeto a partir de varias imágenes de entrada del sujeto.
* [Guía](../training/text2image) para ajustar un modelo de Stable Diffusion con tu propio conjunto de datos.
* [Guía](../training/lora) para usar LoRA, una técnica eficiente en memoria para ajustar modelos muy grandes más rápido.