Entrenamiento de difusión estable con Dreambooth utilizando difusores

Entrenamiento de difusión estable con Dreambooth y difusores.

Dreambooth es una técnica para enseñar nuevos conceptos a la Difusión Estable utilizando una forma especializada de ajuste fino. Algunas personas la han estado utilizando con algunas de sus fotos para situarse en situaciones fantásticas, mientras que otros la están utilizando para incorporar nuevos estilos. 🧨 Diffusers proporciona un script de entrenamiento de Dreambooth. No lleva mucho tiempo entrenar, pero es difícil seleccionar el conjunto correcto de hiperparámetros y es fácil sobreajustar.

Realizamos una gran cantidad de experimentos para analizar el efecto de diferentes configuraciones en Dreambooth. Esta publicación presenta nuestros hallazgos y algunos consejos para mejorar tus resultados al ajustar finamente la Difusión Estable con Dreambooth.

Antes de comenzar, ten en cuenta que este método nunca debe usarse con fines maliciosos, para generar daño de alguna manera o para suplantar personas sin su conocimiento. Los modelos entrenados con él siguen estando limitados por la licencia CreativeML Open RAIL-M que rige la distribución de los modelos de Difusión Estable.

Nota: una versión anterior de esta publicación se publicó como un informe de W&B.

TL;DR: Configuración Recomendada

• Dreambooth tiende a sobreajustarse rápidamente. Para obtener imágenes de buena calidad, debemos encontrar un punto óptimo entre el número de pasos de entrenamiento y la tasa de aprendizaje. Recomendamos utilizar una tasa de aprendizaje baja y aumentar progresivamente el número de pasos hasta que los resultados sean satisfactorios.
• Dreambooth necesita más pasos de entrenamiento para caras. En nuestros experimentos, 800-1200 pasos funcionaron bien cuando se utilizaba un tamaño de lote de 2 y una tasa de aprendizaje de 1e-6.
• La preservación previa es importante para evitar el sobreajuste al entrenar caras. Para otros temas, no parece marcar una gran diferencia.
• Si ves que las imágenes generadas tienen ruido o la calidad se degrada, probablemente signifique sobreajuste. Primero, prueba los pasos anteriores para evitarlo. Si las imágenes generadas siguen teniendo ruido, utiliza el planificador DDIM o realiza más pasos de inferencia (~100 funcionaron bien en nuestros experimentos).
• Entrenar el codificador de texto además de la UNet tiene un gran impacto en la calidad. Nuestros mejores resultados se obtuvieron utilizando una combinación de ajuste fino del codificador de texto, una baja tasa de aprendizaje y un número adecuado de pasos. Sin embargo, el ajuste fino del codificador de texto requiere más memoria, por lo que una GPU con al menos 24 GB de RAM es ideal. Utilizando técnicas como Adam de 8 bits, entrenamiento fp16 o acumulación de gradientes, es posible entrenar en GPU de 16 GB como las proporcionadas por Google Colab o Kaggle.
• El ajuste fino con o sin EMA produjo resultados similares.
• No es necesario utilizar la palabra sks para entrenar Dreambooth. Una de las primeras implementaciones la utilizó porque era un token raro en el vocabulario, pero en realidad es una especie de rifle. Nuestros experimentos, y los de, por ejemplo, @nitrosocke, muestran que está bien seleccionar términos que naturalmente usarías para describir tu objetivo.

Impacto de la Tasa de Aprendizaje

Dreambooth se sobreajusta muy rápidamente. Para obtener buenos resultados, ajusta la tasa de aprendizaje y el número de pasos de entrenamiento de una manera que tenga sentido para tu conjunto de datos. En nuestros experimentos (detallados a continuación), ajustamos finamente en cuatro conjuntos de datos diferentes con tasas de aprendizaje altas y bajas. En todos los casos, obtuvimos mejores resultados con una tasa de aprendizaje baja.

• Generando texto a nivel humano con búsqueda contrastiva en Transfor...
• Director de Perspicacias de Aprendizaje Automático [Parte 4]
• Una visión general de las soluciones de inferencia en Hugging Face

Configuración de los Experimentos

Todos nuestros experimentos se realizaron utilizando el script train_dreambooth.py con el optimizador AdamW en 2x 40GB A100s. Utilizamos la misma semilla y mantuvimos todos los hiperparámetros iguales en todas las ejecuciones, excepto la tasa de aprendizaje, el número de pasos de entrenamiento y el uso de preservación previa.

Para los primeros 3 ejemplos (varios objetos), ajustamos finamente el modelo con un tamaño de lote de 4 (2 por GPU) durante 400 pasos. Utilizamos una tasa de aprendizaje alta de 5e-6 y una tasa de aprendizaje baja de 2e-6. No se utilizó preservación previa.

El último experimento intenta agregar un sujeto humano al modelo. Utilizamos preservación previa con un tamaño de lote de 2 (1 por GPU), 800 y 1200 pasos en este caso. Utilizamos una tasa de aprendizaje alta de 5e-6 y una tasa de aprendizaje baja de 2e-6.

Ten en cuenta que puedes utilizar Adam de 8 bits, entrenamiento fp16 o acumulación de gradientes para reducir los requisitos de memoria y realizar experimentos similares en GPU con 16 GB de memoria.

Juguete para gatos

Tasa de aprendizaje alta ( 5e-6 )

Tasa de aprendizaje baja ( 2e-6 )

Cabeza de cerdo

Tasa de aprendizaje alta ( 5e-6 ). Tenga en cuenta que los artefactos de color son restos de ruido: ejecutar más pasos de inferencia podría ayudar a resolver algunos de esos detalles.

Tasa de aprendizaje baja ( 2e-6 )

Sr. Potato Head

Tasa de aprendizaje alta ( 5e-6 ). Tenga en cuenta que los artefactos de color son restos de ruido: ejecutar más pasos de inferencia podría ayudar a resolver algunos de esos detalles.

Tasa de aprendizaje baja ( 2e-6 )

Rostro humano

Intentamos incorporar el personaje de Kramer de Seinfeld en Stable Diffusion. Como se mencionó anteriormente, entrenamos durante más pasos con un tamaño de lote más pequeño. Aun así, los resultados no fueron excelentes. Por brevedad, hemos omitido estas imágenes de muestra y remitimos al lector a las siguientes secciones, donde el entrenamiento de rostros se convirtió en el foco de nuestros esfuerzos.

Resumen de los resultados iniciales

Para obtener buenos resultados entrenando Stable Diffusion con Dreambooth, es importante ajustar la tasa de aprendizaje y los pasos de entrenamiento para su conjunto de datos.

• Tasas de aprendizaje altas y demasiados pasos de entrenamiento conducirán al sobreajuste. El modelo generará principalmente imágenes de sus datos de entrenamiento, sin importar el texto de entrada utilizado.
• Tasas de aprendizaje bajas y muy pocos pasos conducirán al subajuste: el modelo no podrá generar el concepto que intentamos incorporar.

Los rostros son más difíciles de entrenar. En nuestros experimentos, una tasa de aprendizaje de 2e-6 con 400 pasos de entrenamiento funciona bien para objetos, pero los rostros requieren 1e-6 (o 2e-6 ) con ~1200 pasos.

La calidad de la imagen se degrada mucho si el modelo se sobreajusta, y esto sucede si:

• La tasa de aprendizaje es demasiado alta.
• Ejecutamos demasiados pasos de entrenamiento.
• En el caso de los rostros, cuando no se utiliza ninguna preservación previa, como se muestra en la siguiente sección.

Uso de la preservación previa al entrenar rostros

La preservación previa es una técnica que utiliza imágenes adicionales de la misma clase que estamos intentando entrenar como parte del proceso de ajuste fino. Por ejemplo, si intentamos incorporar a una nueva persona en el modelo, la clase que queremos preservar podría ser “persona”. La preservación previa intenta reducir el sobreajuste utilizando fotos de la nueva persona combinadas con fotos de otras personas. Lo bueno es que podemos generar esas imágenes adicionales de clase utilizando el propio modelo Stable Diffusion. El script de entrenamiento se encarga de eso automáticamente si lo desea, pero también puede proporcionar una carpeta con sus propias imágenes de preservación previa.

Preservación previa, 1200 pasos, lr= 2e-6 .

Sin preservación previa, 1200 pasos, lr= 2e-6 .

Como puedes ver, los resultados son mejores cuando se utiliza la preservación previa, pero aún hay manchas de ruido. ¡Es hora de algunos trucos adicionales!

Efecto de los planificadores

En los ejemplos anteriores, utilizamos el planificador PNDM para muestrear imágenes durante el proceso de inferencia. Observamos que cuando el modelo se sobreajusta, DDIM suele funcionar mucho mejor que PNDM y LMSDiscrete. Además, la calidad puede mejorarse ejecutando la inferencia durante más pasos: 100 parece ser una buena elección. Los pasos adicionales ayudan a resolver algunas de las manchas de ruido en detalles de la imagen.

PNDM, cara de Kramer

LMSDiscrete, cara de Kramer. ¡Los resultados son terribles!

DDIM, cara de Kramer. Mucho mejor

Se puede observar un comportamiento similar para otros sujetos, aunque en menor medida.

PNDM, Sr. Cabeza de Papa

LMSDiscrete, Sr. Cabeza de Papa

DDIM, Sr. Cabeza de Papa

Ajuste fino del codificador de texto

El artículo original de Dreambooth describe un método para ajustar finamente el componente UNet del modelo pero mantiene el codificador de texto congelado. Sin embargo, observamos que ajustar finamente el codificador produce mejores resultados. Experimentamos con este enfoque después de verlo utilizado en otras implementaciones de Dreambooth, ¡y los resultados son sorprendentes!

Codificador de texto congelado

Codificador de texto ajustado finamente

Ajustar finamente el codificador de texto produce los mejores resultados, especialmente con caras. Genera imágenes más realistas, tiene menos tendencia al sobreajuste y también logra una mejor interpretabilidad de las instrucciones, siendo capaz de manejar instrucciones más complejas.

Epílogo: Inversión textual + Dreambooth

También realizamos un experimento final en el que combinamos la Inversión textual con Dreambooth. Ambas técnicas tienen un objetivo similar, pero sus enfoques son diferentes.

En este experimento, primero ejecutamos la inversión textual durante 2000 pasos. A partir de ese modelo, luego ejecutamos Dreambooth durante 500 pasos adicionales utilizando una tasa de aprendizaje de 1e-6. Estos son los resultados:

Pensamos que los resultados son mucho mejores que hacer simplemente Dreambooth pero no tan buenos como cuando ajustamos finamente todo el codificador de texto. Parece copiar más el estilo de las imágenes de entrenamiento, por lo que podría estar sobreajustándose a ellas. No exploramos más esta combinación, pero podría ser una alternativa interesante para mejorar Dreambooth y aún así ajustar el proceso en una GPU de 16GB. ¡Siéntete libre de explorar y contarnos tus resultados!