Optimizando la Difusi贸n Estable para CPUs de Intel con NNCF y 馃 Optimum

Optimizing Stable Diffusion for Intel CPUs with NNCF and 馃 Optimum.

Los modelos de difusi贸n latente son revolucionarios cuando se trata de resolver problemas de generaci贸n de texto a imagen. Diffusion Estable es uno de los ejemplos m谩s famosos que ha sido ampliamente adoptado en la comunidad y la industria. La idea detr谩s del modelo de difusi贸n estable es simple y convincente: se genera una imagen a partir de un vector de ruido en m煤ltiples pasos peque帽os refinando el ruido a una representaci贸n de imagen latente.

Sin embargo, este enfoque inevitablemente aumenta el tiempo total de inferencia y provoca una mala experiencia de usuario cuando se implementa en una m谩quina cliente. Uno puede notar que las potentes GPUs pueden ayudar aqu铆 como de costumbre, y esto es cierto, pero el costo de esto aumenta dr谩sticamente. Como referencia, en H1’23, el precio de una instancia potente de CPU r6i.2xlarge con 8 vCPUs y 64 GB de RAM es de $0.504 por hora, mientras que una instancia similar g4dn.2xlarge con NVIDIA T4 con 16 GB de VRAM cuesta $0.75 por hora, lo que es 1.5 veces m谩s ..

Esto hace que los servicios de generaci贸n de im谩genes sean bastante caros tanto para sus propietarios como para los usuarios. El problema es a煤n m谩s agudo en las aplicaciones cliente que se ejecutan en el lado del usuario. 隆Aqu铆 puede que no haya GPU en absoluto! Esto hace que la implementaci贸n del pipeline Stable Diffusion sea un problema desafiante.

A lo largo de los 煤ltimos cinco a帽os, la caja de herramientas OpenVINO ha encapsulado muchas caracter铆sticas para una inferencia de alto rendimiento. Inicialmente dise帽ado para modelos de visi贸n por computadora, todav铆a domina en este dominio mostrando un rendimiento de inferencia de vanguardia para muchos modelos contempor谩neos, incluido Stable Diffusion. Sin embargo, optimizar modelos de difusi贸n estable para aplicaciones con restricciones de recursos requiere ir mucho m谩s all谩 de las simples optimizaciones en tiempo de ejecuci贸n. Y aqu铆 es donde entran en juego las capacidades de optimizaci贸n de modelos del Marco de Compresi贸n de Redes Neuronales (NNCF) de OpenVINO.

En esta publicaci贸n de blog, describiremos los problemas de optimizaci贸n de los modelos de difusi贸n estable y propondremos un flujo de trabajo que reduce sustancialmente la latencia de dichos modelos cuando se ejecutan en un hardware con recursos limitados, como una CPU. En particular, logramos una aceleraci贸n de inferencia de 5.1 veces y una reducci贸n del tama帽o del modelo de 4 veces en comparaci贸n con PyTorch.

Optimizaci贸n de Stable Diffusion

En el pipeline de Stable Diffusion, el modelo UNet es computacionalmente el m谩s costoso de ejecutar. Por lo tanto, optimizar solo un modelo aporta beneficios sustanciales en t茅rminos de velocidad de inferencia.

Sin embargo, resulta que los m茅todos tradicionales de optimizaci贸n de modelos, como la cuantificaci贸n de 8 bits despu茅s del entrenamiento, no funcionan para este modelo. Hay dos razones principales para esto. Primero, los modelos de predicci贸n a nivel de p铆xel, como la segmentaci贸n sem谩ntica, la superresoluci贸n, etc., son uno de los m谩s complicados en t茅rminos de optimizaci贸n de modelos debido a la complejidad de la tarea, por lo que ajustar los par谩metros y la estructura del modelo rompe los resultados de varias formas. La segunda raz贸n es que el modelo tiene un nivel inferior de redundancia porque aloja mucha informaci贸n mientras se entrena con cientos de millones de muestras. Es por eso que los investigadores deben emplear m茅todos de cuantificaci贸n m谩s sofisticados para preservar la precisi贸n despu茅s de la optimizaci贸n. Por ejemplo, Qualcomm utiliz贸 el m茅todo Knowledge Distillation por capas (AdaRound) para cuantificar los modelos de difusi贸n estables. Esto significa que de todos modos se requiere ajustar el modelo despu茅s de la cuantificaci贸n. Si es as铆, 驴por qu茅 no simplemente utilizar el Entrenamiento con Conciencia de Cuantizaci贸n (QAT) que puede ajustar el modelo y los par谩metros de cuantizaci贸n simult谩neamente de la misma manera en que se entrena el modelo fuente? Por lo tanto, probamos este enfoque en nuestro trabajo utilizando NNCF, OpenVINO y Diffusers y lo combinamos con la Fusi贸n de Tokens.

Flujo de trabajo de optimizaci贸n

Normalmente comenzamos la optimizaci贸n de un modelo despu茅s de que se haya entrenado. Aqu铆, comenzamos con un modelo ajustado a los datos del conjunto de datos de Pokemons, que contiene im谩genes de Pokemons y sus descripciones de texto.

Utilizamos el ejemplo de ajuste fino de texto a imagen para Stable Diffusion de Diffusers e integramos QAT de NNCF en el siguiente script de entrenamiento. Tambi茅n cambiamos la funci贸n de p茅rdida para incorporar destilaci贸n de conocimiento del modelo fuente que act煤a como profesor en este proceso, mientras que el modelo actual que se entrena act煤a como estudiante. Este enfoque es diferente al m茅todo cl谩sico de destilaci贸n de conocimiento, donde el modelo profesor entrenado se destila en un modelo estudiante m谩s peque帽o. En nuestro caso, la destilaci贸n de conocimiento se utiliza como un m茅todo auxiliar que ayuda a mejorar la precisi贸n final del modelo optimizado. Tambi茅n utilizamos el m茅todo de Promedio M贸vil Exponencial (EMA) para los par谩metros del modelo, excluyendo los cuantizadores, lo que nos permite hacer que el proceso de entrenamiento sea m谩s estable. Ajustamos el modelo durante solo 4096 iteraciones.

Con algunos trucos, como la checkpointing de gradientes y mantener el modelo EMA en RAM en lugar de VRAM, podemos ejecutar el proceso de optimizaci贸n utilizando una sola GPU con 24 GB de VRAM. 隆Toda la optimizaci贸n lleva menos de un d铆a usando una GPU!

Trascendiendo el entrenamiento consciente de la cuantizaci贸n

La cuantizaci贸n por s铆 sola puede aportar mejoras significativas al reducir el tama帽o del modelo, el tiempo de carga, el consumo de memoria y la latencia de inferencia. Pero lo maravilloso de la cuantizaci贸n es que se puede aplicar junto con otros m茅todos de optimizaci贸n, lo que lleva a una aceleraci贸n acumulativa.

Recientemente, Facebook Research introdujo un m茅todo de fusi贸n de tokens para modelos de Vision Transformer. La esencia del m茅todo es que fusiona tokens redundantes con tokens importantes utilizando una de las estrategias disponibles (promedio, toma de valores m谩ximos, etc.). Esto se hace antes del bloque de autoatenci贸n, que es la parte m谩s computacionalmente exigente de los modelos Transformer. Por lo tanto, reducir la dimensi贸n del token reduce el tiempo de c谩lculo general en los bloques de autoatenci贸n. Este m茅todo tambi茅n se ha adaptado para modelos de Difusi贸n Estable y ha mostrado resultados prometedores al optimizar pipelines de Difusi贸n Estable para la s铆ntesis de im谩genes de alta resoluci贸n en ejecuci贸n en GPUs.

Modificamos el m茅todo de fusi贸n de tokens para que sea compatible con OpenVINO y lo combinamos con la cuantizaci贸n de 8 bits cuando se aplica al modelo Attention UNet. Esto tambi茅n incluye todas las t茅cnicas mencionadas, incluida la destilaci贸n de conocimiento, etc. En cuanto a la cuantizaci贸n, se requiere un ajuste fino para restaurar la precisi贸n. Tambi茅n comenzamos la optimizaci贸n y el ajuste fino a partir del modelo entrenado en el conjunto de datos de Pokemons. La figura a continuaci贸n muestra un flujo de trabajo de optimizaci贸n general.

El modelo resultante es altamente beneficioso cuando se ejecuta la inferencia en dispositivos con recursos computacionales limitados, como CPUs de cliente o de borde. Como se mencion贸, combinar la fusi贸n de tokens con la cuantizaci贸n conduce a una reducci贸n adicional en la latencia de inferencia.

Resultados de la demostraci贸n de generaci贸n de im谩genes utilizando diferentes modelos optimizados. La entrada es “p谩jaro de dibujos animados”, el n煤mero aleatorio es 42. Los modelos utilizan OpenVINO 2022.3 en Hugging Face Spaces utilizando una instancia de “actualizaci贸n de CPU” que utiliza procesadores escalables Intel庐 Xeon庐 de tercera generaci贸n con tecnolog铆a Intel庐 Deep Learning Boost.

Resultados

Utilizamos los flujos de trabajo de optimizaci贸n divulgados para obtener dos tipos de modelos optimizados, cuantizados de 8 bits y cuantizados con fusi贸n de tokens, y los comparamos con la l铆nea de base de PyTorch. Tambi茅n convertimos la l铆nea de base a un modelo de punto flotante (FP32) de OpenVINO para una comparaci贸n exhaustiva.

La imagen de arriba muestra los resultados de la generaci贸n de im谩genes y algunas caracter铆sticas del modelo. Como se puede ver, simplemente la conversi贸n a OpenVINO produce una disminuci贸n significativa en la latencia de inferencia ( 1.9x ). La aplicaci贸n de la cuantizaci贸n de 8 bits aumenta a煤n m谩s la velocidad de inferencia, lo que lleva a una aceleraci贸n de 3.9x en comparaci贸n con PyTorch. Otro beneficio de la cuantizaci贸n es una reducci贸n significativa del tama帽o del modelo, 0.25x del punto de control de PyTorch, lo que tambi茅n mejora el tiempo de carga del modelo. La aplicaci贸n de la fusi贸n de tokens (ToME) (con una tasa de fusi贸n de 0.4 ) sobre la cuantizaci贸n aporta una aceleraci贸n de rendimiento de 5.1x manteniendo el tama帽o del modelo en el mismo nivel. No proporcionamos un an谩lisis exhaustivo de la calidad visual de los modelos optimizados, pero, como se puede ver, los resultados son bastante s贸lidos.

A continuaci贸n, mostramos c贸mo realizar inferencias con el pipeline final optimizado para ejecutarse en CPUs de Intel:

from optimum.intel import OVStableDiffusionPipeline

# Cargar y compilar el pipeline para obtener un rendimiento 贸ptimo.
nombre = "OpenVINO/stable-diffusion-pokemons-tome-quantized-aggressive"
pipe = OVStableDiffusionPipeline.from_pretrained(nombre, compile=False)
pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
pipe.compile()

# Generar una imagen.
prompt = "un dibujo de un pokemon verde con ojos rojos"
output = pipe(prompt, num_inference_steps=50, output_type="pil").images[0]
output.save("imagen.png")

Puede encontrar el c贸digo de entrenamiento y cuantizaci贸n en la biblioteca Hugging Face Optimum Intel. El cuaderno que muestra la diferencia entre los modelos optimizados y los originales est谩 disponible aqu铆 . Tambi茅n puede encontrar muchos modelos en el Hugging Face Hub en la organizaci贸n de OpenVINO . Adem谩s, hemos creado una demostraci贸n en Hugging Face Spaces que se ejecuta en una instancia r6id.2xlarge con procesador escalable Intel Xeon de tercera generaci贸n.

驴Qu茅 pasa con el modelo de Difusi贸n Estable de prop贸sito general?

Como mostramos con la tarea de generaci贸n de im谩genes de Pokemon, es posible lograr un alto nivel de optimizaci贸n del pipeline de Difusi贸n Estable utilizando una cantidad relativamente peque帽a de recursos de entrenamiento. Al mismo tiempo, es bien sabido que entrenar un modelo de Difusi贸n Estable de prop贸sito general es una tarea costosa . Sin embargo, con suficiente presupuesto y recursos de hardware, es posible optimizar el modelo de prop贸sito general utilizando el enfoque descrito y ajustarlo para producir im谩genes de alta calidad. La 煤nica precauci贸n que tenemos est谩 relacionada con el m茅todo de fusi贸n de tokens que reduce sustancialmente la capacidad del modelo. La regla general aqu铆 es que cuanto m谩s complicado sea el conjunto de datos que tenga para el entrenamiento, menor deber铆a ser la tasa de fusi贸n que debe utilizar durante la optimizaci贸n.

Si disfrutaste leyendo esta publicaci贸n, es posible que tambi茅n est茅s interesado en revisar esta publicaci贸n que discute otros enfoques complementarios para optimizar el rendimiento de Stable Diffusion en CPUs Intel Xeon de 4ta generaci贸n.

We will continue to update Zepes; if you have any questions or suggestions, please contact us!

Share:

Was this article helpful?

93 out of 132 found this helpful

Discover more

Ciencia de Datos

C贸mo acceder a futuras versiones de Python como la 3.12 antes que las masas.

Un tutorial sobre la instalaci贸n y prueba de futuras versiones de Python, como la 3.12, antes que la mayor铆a, para ex...

Inteligencia Artificial

Un 'mapa' m谩s preciso de las luces que ves cuando cierras los ojos puede mejorar los resultados de los 'ojos bi贸nicos'.

Un m茅todo de mapeo de fosfenos desarrollado por investigadores de la Universidad de Monash en Australia podr铆a ayudar...

Inteligencia Artificial

C贸mo los LLM basados en Transformer extraen conocimiento de sus par谩metros

En los 煤ltimos a帽os, los modelos de lenguaje basados en transformadores (LLMs, por sus siglas en ingl茅s) se han vuelt...

Inteligencia Artificial

Investigadores desarrollaron un novedoso m茅todo de IA sin marcadores para rastrear las posturas de las aves en 3D utilizando grabaciones de video.

Rastrear el comportamiento, la mirada y los movimientos detallados de animales y aves ha sido una tarea desafiante pa...

Inteligencia Artificial

LLMs y An谩lisis de Datos C贸mo la IA est谩 dando sentido a los grandes datos para obtener informaci贸n empresarial

Los Modelos de Lenguaje Grandes (LLMs, por sus siglas en ingl茅s) tienen la capacidad de analizar extensos conjuntos d...