Generación de imágenes incondicional

Los modelos de generación de imágenes incondicionales no se condicionan con texto ni con imágenes durante el entrenamiento. Solo generan imágenes que se asemejan a la distribución de sus datos de entrenamiento.

Esta guía explorará el script de entrenamiento train_unconditional.py para que se familiarice con él y para que pueda adaptarlo a su propio caso de uso.

Antes de ejecutar el script, asegúrate de instalar la biblioteca desde el código fuente:

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

Luego, navegue a la carpeta de ejemplo que contiene el script de entrenamiento e instale las dependencias necesarias:

cd examples/unconditional_image_generation
pip install -r requirements.txt

🤗 Accelerate es una biblioteca para ayudarte a entrenar en varias GPUs/TPUs o con precisión mixta. Configura automáticamente tu configuración de entrenamiento en función de tu hardware y entorno. Echa un vistazo al Quick tour de 🤗 Accelerate para saber más.

Inicializa un entorno de 🤗 Accelerate:

accelerate config

Para configurar un entorno predeterminado de 🤗 Accelerate sin elegir ninguna configuración:

accelerate config default

O si su entorno no admite una shell interactiva como un notebook, puede usar:

from accelerate.utils import write_basic_config

write_basic_config()

Por último, si deseas entrenar un modelo con tu propio conjunto de datos, consulta la guía Crear un conjunto de datos para el entrenamiento para aprender a crear un conjunto de datos que funcione con el script de entrenamiento.

Parámetros del script

Las siguientes secciones resaltan las partes del script de entrenamiento que son importantes para entender cómo modificarlo, pero no abordan todos los aspectos del script en detalle. Si te interesa aprender más, no dudes en leer el script e informarnos si tienes alguna pregunta o inquietud.

El script de entrenamiento ofrece muchos parámetros para ayudarte a personalizar tu ejecución de entrenamiento. Todos los parámetros y sus descripciones se encuentran en la función parse_args(). Proporciona valores predeterminados para cada parámetro, como el tamaño del lote de entrenamiento y la tasa de aprendizaje, pero también puedes establecer tus propios valores en el comando de entrenamiento si lo deseas.

Por ejemplo, para acelerar el entrenamiento con precisión mixta usando el formato bf16, añade el parámetro --mixed_precision al comando de entrenamiento:

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

Entre los parámetros básicos e importantes que se deben especificar se incluyen:

• --dataset_name: el nombre del conjunto de datos en el Hub o una ruta local al conjunto de datos en el que se va a entrenar
• --output_dir: dónde guardar el modelo entrenado
• --push_to_hub: si se debe subir el modelo entrenado al Hub
• --checkpointing_steps: frecuencia de guardado de un punto de control durante el entrenamiento del modelo; esto es útil si el entrenamiento se interrumpe, ya que puedes continuar el entrenamiento desde ese punto de control añadiendo --resume_from_checkpoint a tu comando de entrenamiento

Trae tu conjunto de datos y deja que el script de entrenamiento se encargue de todo lo demás.

Script de entrenamiento

El código para preprocesar el conjunto de datos y el bucle de entrenamiento se encuentra en la función main(). Si necesitas adaptar el script de entrenamiento, este es el lugar donde debes realizar los cambios.

El script train_unconditional inicializa un UNet2DModel si no proporcionas una configuración de modelo. Puedes configurar el UNet aquí si lo deseas:

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

A continuación, el script inicializa un scheduler y 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,
)

Luego carga un conjunto de datos y puedes especificar cómo preprocesarlo:

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

Por último, el bucle de entrenamiento se encarga de todo lo demás, como añadir ruido a las imágenes, predecir el residuo de ruido, calcular la pérdida, guardar puntos de control en los pasos especificados y guardar y publicar el modelo en el Hub. Si quieres aprender más sobre cómo funciona el bucle de entrenamiento, consulta el tutorial Understanding pipelines, models and schedulers, que desglosa el patrón básico del proceso de denoising.

Ejecutar el script

Una vez que hayas realizado todos tus cambios o estés conforme con la configuración predeterminada, ¡estás listo para lanzar el script de entrenamiento! 🚀

Una ejecución completa de entrenamiento toma 2 horas en 4xV100 GPUs.

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

El script de entrenamiento crea y guarda un archivo de punto de control en tu repositorio. Ahora puedes cargar y usar tu modelo entrenado para inferencia:

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]