Difusión estable
La pintura de IA ha roto recientemente con éxito el círculo y se ha convertido en un tema candente. DALLE, GLIDE, Stable Diffusion y otros modelos generativos basados en la difusión han provocado un cambio cualitativo en el mapeo de la IA, lo que permite a las personas ver el amanecer de la "IA convirtiéndose en productividad".
Entre estos modelos de difusión, Stable Diffusion se ha convertido en uno de los representantes con sus excelentes efectos y pesos de código abierto, y ha recibido una amplia atención y experiencia. Se basa en el conjunto de datos de pares de "texto-imagen" de ultra gran escala Laion5B, y Stable AI afirma usar 5,000 A100 durante varios meses para entrenar. Magic Square AI reprodujo recientemente el entrenamiento de Stable Diffusion utilizando el conjunto de datos de Google Caption en Firefly II y lo optimizó. A través del complemento hfai.pl desarrollado por Magic Square, el marco Pytorch Lightning del código fuente se puede integrar fácilmente con las características del clúster Firefly, y el entrenamiento del modelo se puede acelerar a través de herramientas de optimización como 3FS , hfreduce y operadores .
Este artículo compartirá nuestra experiencia en la optimización del entrenamiento de difusión estable para ayudar a los investigadores y desarrolladores a reducir el umbral de investigación.
Título del artículo: Síntesis de imágenes de alta resolución con modelos de difusión latente
Dirección original: https://arxiv.org/abs/2112.10752
Dirección de origen: https://github.com/CompVis/stable-diffusion
Repositorio de modelos: https://github.com/HFAiLab/stable-diffusion
Introducción al modelo
Stable Diffusion amplió la capacitación basada en Latent Diffusion, que reemplazó a Text Encoder de BERT a CLIP Text Encoder. Primero comprendamos el diseño del modelo de Latent Diffusion.
En el pasado, aunque la capacidad de generación del modelo de difusión era muy poderosa y alcanzaba el nivel SOTA en muchos tipos diferentes de tareas de generación, debido a sus características de generación iterativa, a menudo requería muchos recursos de GPU durante el entrenamiento y el razonamiento. La difusión latente ha mejorado en este punto. Al cambiar el proceso de difusión del dominio de píxeles de la imagen al espacio latente codificado, este método reduce en gran medida la complejidad del modelo de difusión cuando se está ejecutando, al tiempo que conserva mejores efectos de generación de imágenes y detalles. . La estructura general se muestra en la siguiente figura:
Codificación latente
Latent Diffusion realiza la conversión entre el espacio de píxeles y el espacio latente al agregar un codificador automático variacional al modelo de difusión general. Antes de entrenar el modelo DDPM, primero se entrena un modelo VAE en ImageNet, cuyo objetivo es aprender un codificador y decodificador para comprimir imágenes en códigos latentes. VAE comprimirá la imagen al tamaño original de 1/4 ~ 1/8, por lo que puede reducir en gran medida la complejidad computacional cuando se opera en el espacio latente. En este caso, la imagen después de la compresión y restauración puede retener mejor la información de la imagen original y básicamente no habrá una pérdida excesiva debido a la compresión.
atención cruzada
La difusión latente utiliza de manera innovadora el mecanismo de atención para fusionar la información de control condicional con el proceso de generación de imágenes. En términos de generación de control condicional, el autor introduce información de control condicional en cada capa del modelo U-Net para controlar la dirección de generación de imágenes. En la fusión de la información de control y el proceso de generación de imágenes, el autor introdujo la atención cruzada. El mecanismo de fusión basado en la atención facilita que el modelo use diferentes tipos de información de control condicional, como texto para generar imágenes, imágenes para generar imágenes o gráficos semánticos para generar imágenes, etc.
Generación basada en texto
A diferencia de Latent Diffusion, Stable Diffusion se enfoca en generar imágenes a partir de texto. Stable Diffusion se entrenó utilizando 2500 millones de pares de imagen y texto en el conjunto de datos LAION-5B, que es mucho más grande que Latent Diffusion. Además, inspirado en el trabajo de otra generación como Imagen, el codificador de texto BERT utilizado en Latent Diffusion fue reemplazado por un mejor codificador de texto entrenado previamente por CLIP/ViT-L-14.
practica modelo
conjunto de datos de entrenamiento
Para verificar el rendimiento del entrenamiento del modelo Stable Diffusion, reproducimos el entrenamiento Stable Diffusion utilizando el conjunto de datos de subtítulos conceptuales de Google. Google Conceptual Caption es un conjunto de datos multimodal de escala relativamente pequeña con 2,85 millones de pares de imagen y texto. El conjunto de datos se integró en el almacén de conjuntos de datos de Magic Square AI, se convirtió a un formato de datos de entrenamiento frecord y se almacenó en almacenamiento de alta velocidad 3FS. Los usuarios pueden obtener lectura de datos de entrenamiento de alta velocidad a través de los siguientes métodos:
from hfai.datasets import GoogleConceptualCaption
dataset = GoogleConceptualCaption(split="train", transform=transform)
dataloader = dataset.loader(**args)hfai.pl
Pytorch Lightning (PL) es un contenedor sobre PyTorch con su exclusiva interfaz de entrenamiento en paralelo. El código fuente de Stable Diffusion está construido en base a PL Para que aproveche varias funciones de optimización del clúster firefly, utilizamos el complemento hfai.pl desarrollado por Magic Square para la adaptación, que incluye:
hfai.pl.HFAIEnvironment, se adapta automáticamente al entorno paralelo de múltiples tarjetas del clúster firefly y se puede usar normalmente después de agregar un complemento durante el entrenamiento;
hfreduce_bind_numa, use hfreduce para acelerar la comunicación, enlace buma para evitar una sobrecarga de red adicional entre varias tarjetas;
hfai.pl.nn_to_hfai, use el operador de optimización hfai para reemplazar el operador básico en el modelo para acelerar el entrenamiento
El funcionamiento específico es el siguiente:
1. Especifique la estrategia del entrenador como hfreduce_bind_numa en el archivo de configuración:
trainer:
max_epochs: 300
strategy: hfreduce_bind_numa
...2. En el código de entrenamiento, use la aceleración del operador nn_to_hfai y la función de configuración del entorno de HFAIEnvironment:
diffusionModelModule = nn_to_hfai(diffusionModelModule)
...
trainer = Trainer.from_argparse_args(trainer_opt, **trainer_kwargs)
trainer.plugins = [HFAIEnvironment()]A través de los sencillos pasos anteriores, las funciones de aceleración del clúster Firefly se pueden integrar en el código de entrenamiento de Stable Diffusion. Pruebe el cambio de velocidad de entrenamiento del modelo antes y después de usar hfai.pl Se puede encontrar que el tiempo de avance único del modelo se ha acelerado de 0,787 segundos a 0,758 segundos, y la velocidad ha aumentado en un 3,8 % .
optimización del entrenamiento
Entrenamos en Google Conceptual Caption a una resolución de 256x256, usando pesos previamente entrenados en ImageNet para inicializar el modelo VQVAE para el mapeo del espacio latente.
Durante el entrenamiento, intentamos usar 4, 8, 16 y 32 nodos para el entrenamiento de difusión estable para probar la sensibilidad en diferentes condiciones paralelas. En el proceso de expansión gradual de la escala del entrenamiento de Stable Diffusion, descubrimos que Stable Diffusion es muy sensible a la tasa de aprendizaje, pero la tasa de aprendizaje no puede aumentar fácilmente con el aumento del tamaño del lote, lo que puede conducir fácilmente a una explosión de gradiente. Por lo tanto, usamos los métodos de Calentamiento y Gradient Clipping durante el entrenamiento para ayudar al modelo a acelerar la convergencia y evitar la situación en la que no puede converger.
Las curvas naranja y azul de la figura anterior muestran la pérdida cuando se entrena con o sin Calentamiento y Gradient Clipping, respectivamente. Como se muestra en la curva naranja, cuando no se usan las estrategias Clipping y Warmup, el modelo tiene un fenómeno de explosión de gradiente cuando el modelo se entrena al paso 400 y no puede continuar convergiendo normalmente. Después de usarlo, el crecimiento de la tasa de aprendizaje del entrenamiento del modelo se vuelve más gradual, lo que puede evitar efectivamente la ocurrencia de una explosión de gradiente.
resultado de entrenamiento
我们将模型在 Googlecc 数据集上训练了 240K Steps 后进行测试。在 COCO FID-30K（一个 COCO Caption 数据集中随机挑选的子集，由 3w 张图片组成）上，模型取得了 16.5 的 FID 指标，说明模型生成的图片能比较好的体现出文本中的内容。
以下是一些样例文本和将他们输入模型后生成的图片：
(a) A photo of a woman skiing on a white mountain.
(b) A painting of a squirrel eating a burger.
(c) A photo of a red train being operated on a train track.
(d) A photo of a dog playing in a green field next to a lake.
可见虽然训练数据集规模较小，但模型仍然达到了较好的生成效果。
体验总结
Stable Diffusion 作为 AI 作图领域的旗舰模型，受到了广泛的关注，在小范围数据上训练也可以实现惊艳的生成效果。我们借助幻方萤火集群，通过简单几步改造，能比较轻松地实现 Stable Diffusion 的训练加速，证明了萤火集群的易用性和实力。
综合体验打分如下：
01：研究新颖度 ★★★★
作者提出了一种在隐空间上进行扩散的生成模型结构，降低了扩散模型运行开销的同时保证了生成质量。模型还应用了交叉注意力机制来辅助条件控制生成，并且支持多种不同模态条件下的图像生成。
02：开源指数 ★★★★★
作为首个完全开源代码、训练数据和预训练权重的 AI 绘画预训练大模型，stable-diffusion 在学术界和其他相关领域都产生了极大的影响力。
03：算力门槛 ★★
由于模型对资源占用有所优化，且开源工作完善，因此单个普通 GPU 即可运行模型推理。但训练开销较高。
04：通用指数 ★★★★
作者提出的在隐空间上进行扩散的方法对一般的扩散模型都能够适用，并且基于交叉注意力的条件控制方法也能将模型应用于许多不同任务类型，对生成领域研究工作有广泛的借鉴意义。
05：模型适配度 ★★★
该项目依赖 pytorch-lightning，需要对萤火集群进行一定适配，但通过 hfai.pl 工具也能比较容易的在幻方 AI 环境运行并获得加速效果。
我们希望让更多“想象力”和“创造力”生长。期待与各方科学家及开发者们一同共建AI时代。