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

Jason Mayes

François Beaufort

¡Tenemos buenas noticias! Compilaste una aplicación de IA web genial que ejecuta modelos de aprendizaje automático directamente en el dispositivo de un usuario. Se ejecuta completamente 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 CPU (WebAssembly) y en GPU más potentes (a través de WebGL y WebGPU). La pregunta es: ¿cómo se pueden automatizar de manera coherente las pruebas del navegador con el hardware seleccionado?

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

Configurar un entorno de pruebas coherente con GPU 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.

Esto no es solo para desarrolladores de IA web. Si estás trabajando en videojuegos o gráficos web, 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 es compatible con WebGPU, un potente sucesor 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 manera programática con JavaScript. Con Puppeteer, podemos automatizar Chrome en el 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 comprobar si la aceleración de hardware está activada en Chrome es escribir chrome://gpu en la barra de direcciones. Puedes realizar el equivalente con Puppeteer de manera programática con console.log o guardar el informe completo como PDF para verificarlo 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. Deberías obtener los siguientes resultados:

No fue un gran comienzo. Es bastante claro que la detección de hardware estaba fallando. En esencia, WebGL, WebGL2 y WebGPU están inhabilitados o solo son software. No somos los únicos que nos encargamos de este problema. Existen numerosos debates en línea sobre personas que se encuentran en situaciones similares, por ejemplo, en los canales oficiales de asistencia de Chrome (1) (2).

Habilita la compatibilidad con WebGPU y WebGL

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

• 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 raíz 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 las llamadas a la API de Vulkan.
• La marca --enable-features=Vulkan habilita el backend de gráficos de 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 una cadena de intercambio, se usa Bit blit para el resultado de renderización actual en la pantalla.
• La marca --enable-unsafe-webgpu habilita la API experimental de WebGPU en Chrome en Linux e inhabilita la lista de entidades bloqueadas del adaptador.

Ahora, combinamos todos los cambios que hicimos hasta ahora. Aquí está 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 de software.

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

Instala los controladores de GPU correctos

Investigamos en más detalle el resultado de chrome://gpu, junto con algunos expertos en GPU del equipo de Chrome. Detectamos problemas con los controladores predeterminados instalados en la instancia de Linux Colab, que causan problemas con Vulkan, lo que hace que Chrome no pueda detectar 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 eran compatibles permite solucionar 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 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 obtendremos el siguiente resultado. 🎉

Gracias al uso de los controladores y las marcas correctos cuando se ejecuta Chrome, ahora contamos con compatibilidad con WebGPU y WebGL que utiliza el nuevo modo sin interfaz gráfica.

Detrás de escena: La investigación de nuestro equipo

Después de investigar mucho, no encontramos métodos de trabajo para el entorno que necesitábamos ejecutar en Google Colab, aunque hubo algunas publicaciones esperanzadoras que funcionaron en otros entornos, lo cual era prometedor. En última instancia, no pudimos replicar su éxito en el entorno de Colab NVIDIA T4, ya que hubo 2 problemas clave:

1. Algunas combinaciones de marcas permiten la detección de la GPU, pero no te permiten usar la GPU.
2. Los ejemplos de soluciones de trabajo de terceros usaban la versión anterior sin interfaz gráfica de Chrome, que en algún momento dejará de estar disponible y se usará la nueva. Necesitábamos una solución que funcionara con el nuevo Chrome sin interfaz gráfica a fin de estar preparado para el futuro.

Confirmamos el poco uso de la GPU con la ejecución de 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 “Hello World” del aprendizaje automático).

En una máquina normal, se deben ejecutar 50 ciclos de entrenamiento (conocidos como ciclos de entrenamiento) en menos de 1 segundo cada uno. Si 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 fueron realmente estos ciclos de entrenamiento.

Como se esperaba, cada ciclo de entrenamiento tardó mucho más de lo esperado (varios segundos), lo que sugiere que Chrome volvió a la antigua 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 Headless Chrome, volver a ejecutar el ejemplo de entrenamiento de TensorFlow.js genera ciclos de entrenamiento mucho más rápidos.

Resumen

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

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

El éxito en el uso conlleva una gran responsabilidad. A este nivel de uso en los sistemas de producción, surge la necesidad de probar los modelos de IA del cliente basados en el navegador en un verdadero entorno de navegador, a la vez que son escalables, automatizables y dentro de una configuración de hardware estandarizada conocida.

Si aprovechas la potencia combinada de Headless Chrome y Puppeteer, puedes probar con confianza esas cargas de trabajo en un entorno estandarizado y replicable, lo que garantiza resultados coherentes y confiables.

Conclusión

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

Si esto te resultó útil, menciónalo en LinkedIn, X (antes Twitter) o en cualquier red social que uses con el hashtag #WebAI. Nos encantaría recibir tus comentarios para que sepamos que debemos escribir más artículos como este en el futuro.

Agrega una estrella en el repositorio de GitHub para recibir las actualizaciones futuras.

Agradecimientos

Agradecemos a todos los miembros del equipo de Chrome que ayudaron 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 Wordsmith en esta entrada de blog. Gracias a Yuly Novikov, Andrey Kosyakov y Alex Rudenko, que fueron fundamentales para crear la solución final y funcional.