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

Jason Mayes
Jason Mayes
Francisco Beaufort
François Beaufort

¡Tenemos buenas noticias! Creaste 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 operar en CPU (WebAssembly) y GPU más potentes (mediante WebGL y WebGPU). La pregunta es: ¿Cómo puedes automatizar sistemáticamente 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 solucionamos para que puedas mejorar el rendimiento de tu aplicación.

Esto no es solo para desarrolladores de IA web. Si estás trabajando en gráficos o videojuegos web, esta publicación también es valiosa para ti.

Qué incluye 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 modernas de GPU 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 modo sin interfaz gráfica, lo que significa que el navegador se ejecuta en un servidor sin una interfaz visible. Estamos usando 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 de manera programática con Puppeteer 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:

Estado de las funciones de gráficos
OpenGL: Inhabilitada
Vulkan: Inhabilitada
WebGL: Solo software. La aceleración de hardware no está disponible.
WebGL2: Solo software. La aceleración de hardware no está disponible.
WebGPU: Inhabilitada

Se detectaron problemas.
Se inhabilitó WebGPU mediante la lista de entidades bloqueadas o la línea de comandos.

No fue un buen comienzo. Está bastante claro que fallaba la detección de hardware. En esencia, WebGL, WebGL2 y WebGPU están inhabilitados o son solo software. No estamos solos en este problema. Hay numerosas conversaciones en línea entre personas que se encuentran en situaciones similares, incluidos 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 habilitarla 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 con permisos de administrador sin esta zona de pruebas.
  • La función experimental --headless=new ejecuta Chrome con el modo sin interfaz gráfica nuevo y mejorado, 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 rasterización en Chrome.
  • La marca --disable-vulkan-surface inhabilita la extensión de la instancia de vulkan VK_KHR_surface. En lugar de usar una cadena de intercambio, se usa Bit blit para el resultado de la renderización actual en la pantalla.
  • La marca --enable-unsafe-webgpu habilita la API experimental de WebGPU en Chrome en Linux y, luego, inhabilita la lista de adaptadores bloqueados de los adaptadores.

Ahora, combinamos todos los cambios que hicimos hasta ahora. Aquí está el guion completo.

/* 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.

Estado de las funciones de gráficos
OpenGL: Inhabilitada
Vulkan: Inhabilitada
WebGL: Solo software. La aceleración de hardware no está disponible.
WebGL2: Solo software. La aceleración de hardware no está disponible.
WebGPU: Solo software. La aceleración de hardware no está disponible.

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 más de cerca 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 causaban problemas con Vulkan, que provocaban que Chrome no pudiera detectar la GPU NVIDIA T4 en el nivel GL_RENDERER, como se muestra en el siguiente resultado. Esto genera problemas con Chrome sin interfaz gráfica.

El resultado predeterminado no detecta la GPU NVIDIA T4.
Información del conductor
GL_RENDERER ANGLE (Google, Vulkan 1.3.0 [dispositivo SwiftShader [Subzero] [0x0000C0DE]), controlador SwiftShader-5.0.0"]

Para solucionar el problema, instalas los controladores correctos que eran compatibles.

Resultado actualizado después de que se instalan los controladores.
Información del conductor
GL_RENDERER ANGLE (NVIDIA Corporation, Tesla T4/PCIe/SSE2, OpenGL ES 3.2 NVIDIA 525.105.17)

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

Vuelve a ejecutar la secuencia de comandos y obtenemos el siguiente resultado. 🎉

Estado de las funciones de gráficos
OpenGL: Habilitado
Vulkan: Habilitado
WebGL: Hardware acelerado, pero con rendimiento reducido.
WebGL2: Hardware acelerado, pero con rendimiento reducido.
WebGPU: Hardware acelerado, pero con rendimiento reducido.

Gracias al uso de controladores y marcas correctos cuando se ejecuta Chrome, ahora contamos con compatibilidad con WebGPU y WebGL que usa 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 que funcionen para el entorno que necesitábamos ejecutar en Google Colab, aunque algunas publicaciones esperanzadoras funcionaron en otros entornos, lo cual era prometedor. En última instancia, no pudimos replicar su éxito en el entorno de NVIDIA T4 de Colab, ya que tuvimos 2 problemas clave:

  1. Algunas combinaciones de marcas permiten la detección de la GPU, pero no te permiten usarla.
  2. En algunos ejemplos de soluciones en funcionamiento de terceros, se usaba la versión anterior sin interfaz gráfica de Chrome, que en algún momento dejará de estar disponible para dar lugar a la versión nueva. Necesitábamos una solución que funcionara con el nuevo Chrome sin interfaz gráfica a fin de estar mejor preparado para el futuro.

Confirmamos el bajo uso 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 “Hello World” del aprendizaje automático).

En una máquina normal, 50 ciclos de entrenamiento (conocidos como ciclos de entrenamiento) deben ejecutarse en menos de 1 segundo cada uno. Si llamamos a Headless Chrome en su estado predeterminado, podemos registrar el resultado de la Consola de JavaScript en la línea de comandos del servidor de Node.js para ver qué tan rápidos eran realmente estos ciclos de entrenamiento.

Como era de esperar, cada ciclo de entrenamiento tardó mucho más de lo esperado (varios segundos), lo que sugiere que Chrome se recurrió a la antigua ejecución de CPU de JS en lugar de utilizar la GPU:

Los ciclos de entrenamiento se mueven a una cadencia más lenta.
Figura 1: Captura en tiempo real que muestra cuánto tiempo tardó en ejecutarse cada ciclo de entrenamiento (segundos).

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 da como resultado ciclos de entrenamiento mucho más rápidos.

Hay un aumento en la velocidad para las épocas...
Figura 2: Captura en tiempo real que muestra la velocidad de los ciclos de entrenamiento.

Resumen

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

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

Un gran éxito en el uso conlleva una gran responsabilidad. En 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 entorno de navegador real, que también sea escalable, automatizable y dentro de una configuración de hardware estandarizada conocida.

Si aprovechas la potencia combinada de los nuevos 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

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

Si esto te resultó útil, envía un saludo en LinkedIn, X (anteriormente Twitter) o en cualquier red social que uses con el hashtag #WebAI. Nos encantaría recibir tus comentarios, para que podamos escribir más cosas como esta en el futuro.

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

Agradecimientos

Agradecemos enormemente a todos los miembros del equipo de Chrome que nos ayudaron a depurar los problemas del controlador y de WebGPU que enfrentamos en esta solución, y queremos agradecerles de manera especial a Jecelyn Yeen y Alexandra White por ayudarlas con esta entrada de blog. Gracias a Yuly Novikov, Andrey Kosyakov y Alex Rudenko, quienes fueron fundamentales en la creación de la solución final y funcional.