Pruebas del modelo de IA web Supercharge: WebGPU, WebGL y Headless Chrome

Jason Mayes

François Beaufort

¡Tenemos buenas noticias! Compartiste una aplicación de IA web genial que ejecuta modelos de aprendizaje automático directamente en el dispositivo de un usuario. Se ejecuta por completo en el navegador web del cliente, sin depender de la nube. Este diseño integrado en el dispositivo mejora la privacidad del usuario, aumenta el rendimiento y reduce los costos de manera significativa.

Sin embargo, hay un obstáculo. Tu modelo de TensorFlow.js puede funcionar en CPUs (WebAssembly) y GPUs más potentes (a través de WebGL y WebGPU). La pregunta es: ¿cómo puedes automatizar de forma coherente las pruebas del navegador con el hardware seleccionado?

Mantener la coherencia es fundamental para comparar el rendimiento del modelo de aprendizaje automático con el tiempo a medida que lo iteras y mejoras, antes de la implementación para que los usuarios del mundo real lo usen en sus dispositivos.

La configuración de un entorno de pruebas coherente con GPUs puede ser más difícil de lo esperado. En esta entrada de blog, compartiremos los problemas que enfrentamos y cómo los resolvimos para que puedas mejorar el rendimiento de tu aplicación.

No solo es para desarrolladores de IA web. Si trabajas en juegos web o en gráficos, esta publicación también es valiosa para ti.

Contenido de nuestra caja de herramientas de automatización

Esto es lo que usamos:

• Entorno: Un notebook de Google Colab basado en Linux conectado a una GPU NVIDIA T4 o V100. Si lo prefieres, puedes usar otras plataformas en la nube, como Google Cloud (GCP).
• Navegador: Chrome admite WebGPU, un sucesor potente de WebGL que lleva los avances de las APIs de GPU modernas a la Web.
• Automatización: Puppeteer es una biblioteca de Node.js que te permite controlar los navegadores de forma programática con JavaScript. Con Puppeteer, podemos automatizar Chrome en modo sin interfaz gráfica, lo que significa que el navegador se ejecuta sin una interfaz visible en un servidor. Usamos el nuevo modo sin interfaz gráfica mejorado, no el formulario heredado.

Verifica el entorno

La mejor manera de verificar si la aceleración de hardware está activada en Chrome es escribir chrome://gpu en la barra de direcciones. Puedes realizar el equivalente de forma programática con Puppeteer con console.log o guardar el informe completo como PDF para revisarlo de forma manual:

/* Incomplete example.js */
import puppeteer from 'puppeteer';

// Configure launch parameters: Expands later
const browser = await puppeteer.launch({
  headless: 'new',
  args:  ['--no-sandbox']
});

const page = await browser.newPage();
await page.goto('chrome://gpu');

// Verify: log the WebGPU status or save the GPU report as PDF
const txt = await page.waitForSelector('text/WebGPU');
const status = await txt.evaluate(g => g.parentElement.textContent);
console.log(status);
await page.pdf({ path: './gpu.pdf' });

await browser.close();

Abre chrome://gpu y deberías ver los siguientes resultados:

No es un buen comienzo. Está bastante claro que la detección de hardware fallaba. WebGL, WebGL2 y WebGPU están esencialmente inhabilitados o solo son software. No estamos solos con este problema. Hay muchas conversaciones en línea de personas en una situación similar, incluso en los canales oficiales de asistencia de Chrome (1) y (2).

Habilita la compatibilidad con WebGPU y WebGL

De forma predeterminada, Chrome sin interfaz gráfica inhabilita la GPU. Para habilitarlo en Linux, aplica todas las siguientes marcas cuando inicies Chrome sin cabeza:

• La marca --no-sandbox inhabilita la Zona de pruebas de seguridad de Chrome, que aísla el proceso del navegador del resto del sistema. No se admite la ejecución de Chrome como root sin esta zona de pruebas.
• La marca --headless=new ejecuta Chrome con el nuevo y mejorado modo sin interfaz gráfica, sin ninguna IU visible.
• La marca --use-angle=vulkan le indica a Chrome que use el backend de Vulkan para ANGLE, que traduce las llamadas de OpenGL ES 2/3 a llamadas a la API de Vulkan.
• La marca --enable-features=Vulkan habilita el backend de gráficos Vulkan para la composición y la rasterización en Chrome.
• La marca --disable-vulkan-surface inhabilita la extensión de instancia de vulkan VK_KHR_surface. En lugar de usar un grupo de intercambio, se usa Bit blit para mostrar el resultado de la renderización en pantalla.
• La marca --enable-unsafe-webgpu habilita la API experimental de WebGPU en Chrome en Linux y, además, inhabilita la lista de bloqueo de adaptadores.

Ahora, combinamos todos los cambios que hicimos hasta ahora. Esta es la secuencia de comandos completa.

/* Complete example.js */
import puppeteer from 'puppeteer';

// Configure launch parameters
const browser = await puppeteer.launch({
  headless: 'new',
  args: [
    '--no-sandbox',
    '--headless=new',
    '--use-angle=vulkan',
    '--enable-features=Vulkan',
    '--disable-vulkan-surface',
    '--enable-unsafe-webgpu',
  ]
});

const page = await browser.newPage();
await page.goto('chrome://gpu');

// Verify: log the WebGPU status or save the GPU report as PDF
const txt = await page.waitForSelector('text/WebGPU');
const status = await txt.evaluate(g => g.parentElement.textContent);
console.log(status);
await page.pdf({path: './gpu.pdf'});

await browser.close();

Vuelve a ejecutar la secuencia de comandos. No se detectan problemas de WebGPU, y el valor cambia de inhabilitado a solo software.

Sin embargo, la aceleración de hardware aún no está disponible, y no se detecta la GPU NVIDIA T4.

Instala los controladores de GPU correctos

Investigamos con mayor detalle el resultado de chrome://gpu con algunos expertos en GPU del equipo de Chrome. Detectamos problemas con los controladores predeterminados instalados en la instancia de Linux Colab, lo que causaba problemas con Vulkan, lo que impedía que Chrome detectara la GPU NVIDIA T4 en el nivel GL_RENDERER, como se muestra en el siguiente resultado. Esto causa problemas con Headless Chrome.

Por lo tanto, instalar los controladores correctos que sean compatibles soluciona el problema.

Para instalar los controladores correctos, ejecuta los siguientes comandos durante la configuración. Las últimas dos líneas te ayudan a registrar los resultados de lo que detectan los controladores de NVIDIA junto con vulkaninfo.

apt-get install -y vulkan-tools libnvidia-gl-525

// Verify the NVIDIA drivers detects along with vulkaninfo
nvidia-smi
vulkaninfo --summary

Ahora, vuelve a ejecutar la secuencia de comandos y obtén el siguiente resultado. 🎉

Cuando usamos los controladores y las marcas correctos cuando ejecutamos Chrome, ahora tenemos compatibilidad con WebGPU y WebGL con el nuevo modo sin cabeza.

Tras bambalinas: La investigación de nuestro equipo

Después de mucha investigación, no encontramos métodos que funcionaran para el entorno que necesitábamos ejecutar en Google Colab, aunque había algunas publicaciones prometedoras que funcionaban en otros entornos. En última instancia, no pudimos replicar su éxito en el entorno de Colab NVIDIA T4, ya que tuvimos 2 problemas clave:

1. Algunas combinaciones de marcas permiten detectar la GPU, pero no te permiten usarla.
2. En los ejemplos de soluciones funcionales de terceros, se usó la versión anterior de Chrome sin interfaz gráfica, que en algún momento dejará de estar disponible a favor de la versión nueva. Necesitábamos una solución que funcionara con el nuevo Headless Chrome para estar mejor preparados para el futuro.

Confirmamos el uso insuficiente de la GPU ejecutando una página web de ejemplo de TensorFlow.js para el reconocimiento de imágenes, en la que entrenamos un modelo para reconocer muestras de ropa (algo así como un "hola mundo" del aprendizaje automático).

En una máquina normal, 50 ciclos de entrenamiento (conocidos como épocas) deberían ejecutarse en menos de 1 segundo cada uno. Cuando llamamos a Headless Chrome en su estado predeterminado, pudimos registrar el resultado de la consola de JavaScript en la línea de comandos del servidor de Node.js para ver qué tan rápido se estaban realizando estos ciclos de entrenamiento.

Como se esperaba, cada época de entrenamiento tardó mucho más de lo esperado (varios segundos), lo que sugiere que Chrome volvió a la ejecución de CPU de JS en lugar de usar la GPU:

Después de corregir los controladores y usar la combinación correcta de marcas para Chrome sin cabeza, volver a ejecutar el ejemplo de entrenamiento de TensorFlow.js genera épocas de entrenamiento mucho más rápidas.

Resumen

La IA web creció de forma exponencial desde su creación en 2017. Con tecnologías de navegadores como WebGPU, WebGL y WebAssembly, las operaciones matemáticas de un modelo de aprendizaje automático se pueden acelerar aún más del lado del cliente.

A partir de 2023, TensorFlow.js y MediaPipe Web superaron los 1,000 millones de descargas de modelos y bibliotecas, un hito histórico y una señal de cómo los desarrolladores y los ingenieros web están cambiando para adoptar la IA en sus apps web de nueva generación para crear soluciones realmente increíbles.

El gran éxito en el uso conlleva una gran responsabilidad. En este nivel de uso en sistemas de producción, surge la necesidad de probar modelos de IA basados en el navegador y del cliente en un entorno de navegador real, a la vez que se pueden escalar, automatizar y ejecutar dentro de una configuración de hardware estandarizada conocida.

Si aprovechas la potencia combinada de la nueva versión de Chrome sin interfaz gráfica y Puppeteer, puedes probar con confianza esas cargas de trabajo en un entorno estandarizado y reproducible, lo que garantiza resultados coherentes y confiables.

Conclusión

En nuestra documentación, hay disponible una guía paso a paso para que puedas probar la configuración completa por tu cuenta.

Si te resultó útil, menciónanos en LinkedIn, X (anteriormente Twitter) o en cualquier red social que uses con el hashtag #WebAI. Sería genial conocer tus comentarios para saber qué temas escribir en el futuro.

Agrega una estrella al repositorio de GitHub para recibir actualizaciones futuras.

Agradecimientos

Muchas gracias a todo el equipo de Chrome que ayudó a depurar el controlador y los problemas de WebGPU que enfrentamos en esta solución, con un agradecimiento especial a Jecelyn Yeen y Alexandra White por ayudar a redactar esta entrada de blog. Gracias a Yuly Novikov, Andrey Kosyakov y Alex Rudenko, que fueron fundamentales en la creación de la solución final y funcional.