Diffusers documentation

LoRA

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

LoRA

Esto es experimental y la API puede cambiar en el futuro.

LoRA (Low-Rank Adaptation of Large Language Models) es una técnica de entrenamiento popular y ligera que reduce significativamente el número de parámetros entrenables. Funciona insertando un número menor de nuevos pesos en el modelo, y solo estos se entrenan. Esto hace que el entrenamiento con LoRA sea mucho más rápido, eficiente en memoria y produzca pesos de modelo más pequeños (unos pocos cientos de MB), que son más fáciles de almacenar y compartir. LoRA también se puede combinar con otras técnicas de entrenamiento como DreamBooth para acelerar el entrenamiento.

LoRA es muy versátil y se admite para DreamBooth, Kandinsky 2.2, Stable Diffusion XL y text-to-image.

Esta guía explorará el script train_text_to_image_lora.py para que te familiarices más con él y para que sepas cómo adaptarlo a tu 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 .

Navegue hasta la carpeta de ejemplo con el script de entrenamiento e instale las dependencias requeridas para el script que esté utilizando:

cd examples/text_to_image
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.

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.

Parámetros del script

El script de entrenamiento tiene 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(). Se proporcionan valores predeterminados para la mayoría de los parámetros que funcionan bastante bien, pero también puedes establecer tus propios valores en el comando de entrenamiento si lo deseas.

Por ejemplo, para aumentar el número de épocas de entrenamiento:

accelerate launch train_text_to_image_lora.py \
  --num_train_epochs=150 \

Muchos de los parámetros básicos e importantes se describen en la guía de entrenamiento de Texto a imagen, por lo que esta guía se centra únicamente en los parámetros relevantes para LoRA:

  • --rank: la dimensión interna de las matrices de bajo rango que se entrenan; un rango más alto significa más parámetros entrenables
  • --learning_rate: la tasa de aprendizaje predeterminada es 1e-4, pero con LoRA puedes usar una tasa de aprendizaje más alta

Script de entrenamiento

El código de preprocesamiento del conjunto de datos y el bucle de entrenamiento se encuentran en la función main(), y si necesitas adaptar el script de entrenamiento, este es el lugar donde realizarás tus cambios.

Al igual que con los parámetros del script, se proporciona un recorrido por el script de entrenamiento en la guía de entrenamiento Texto a imagen. En su lugar, esta guía examina las partes del script relevantes para LoRA.

UNet
codificador de texto

Diffusers utiliza LoraConfig de la biblioteca PEFT para configurar los parámetros del adaptador LoRA, como el rango, alpha y en qué módulos insertar los pesos de LoRA. El adaptador se añade a la UNet y solo las capas LoRA se filtran para la optimización en lora_layers.

unet_lora_config = LoraConfig(
    r=args.rank,
    lora_alpha=args.rank,
    init_lora_weights="gaussian",
    target_modules=["to_k", "to_q", "to_v", "to_out.0"],
)

unet.add_adapter(unet_lora_config)
lora_layers = filter(lambda p: p.requires_grad, unet.parameters())

El optimizer se inicializa con lora_layers porque estos son los únicos pesos que se optimizarán:

optimizer = optimizer_cls(
    lora_layers,
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

Además de configurar las capas LoRA, el script de entrenamiento es más o menos el mismo que train_text_to_image.py!

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! 🚀

Entrenemos con el conjunto de datos Naruto BLIP captions para generar tus propios personajes de Naruto. Establece las variables de entorno MODEL_NAME y DATASET_NAME para el modelo y el conjunto de datos, respectivamente. También debes especificar dónde guardar el modelo en OUTPUT_DIR y el nombre del modelo a guardar en el Hub con HUB_MODEL_ID. El script crea y guarda los siguientes archivos en tu repositorio:

  • puntos de control de modelos guardados
  • pytorch_lora_weights.safetensors (los pesos de LoRA entrenados)

Si estás entrenando en más de una GPU, añade el parámetro --multi_gpu al comando accelerate launch.

Una ejecución completa de entrenamiento toma ~5 horas en una GPU 2080 Ti con 11 GB de VRAM.

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="/sddata/finetune/lora/naruto"
export HUB_MODEL_ID="naruto-lora"
export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$DATASET_NAME \
  --dataloader_num_workers=8 \
  --resolution=512 \
  --center_crop \
  --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=15000 \
  --learning_rate=1e-04 \
  --max_grad_norm=1 \
  --lr_scheduler="cosine" \
  --lr_warmup_steps=0 \
  --output_dir=${OUTPUT_DIR} \
  --push_to_hub \
  --hub_model_id=${HUB_MODEL_ID} \
  --report_to=wandb \
  --checkpointing_steps=500 \
  --validation_prompt="A naruto with blue eyes." \
  --seed=1337

Una vez completado el entrenamiento, puedes usar tu modelo para inferencia:

from diffusers import AutoPipelineForText2Image
import torch

pipeline = AutoPipelineForText2Image.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.load_lora_weights("path/to/lora/model", weight_name="pytorch_lora_weights.safetensors")
image = pipeline("A naruto with blue eyes").images[0]

Próximos pasos

¡Felicidades por entrenar un nuevo modelo con LoRA! Para aprender más sobre cómo usar tu nuevo modelo, las siguientes guías pueden ser útiles:

Update on GitHub