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

Jason Mayes

François Beaufort

¡Tenemos buenas noticias! Creaste una genial aplicación de IA web 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. En el dispositivo Diseño que mejora la privacidad del usuario, aumenta el rendimiento y reduce los costos de manera significativa.

Sin embargo, hay un obstáculo. Tu El modelo de TensorFlow.js puede funcionar en con CPU (WebAssembly) y GPU más potentes (mediante WebGL y WebGPU). La pregunta es la siguiente: ¿Cómo puedes automatizar de manera coherente las pruebas del navegador con el hardware seleccionado?

Mantener la coherencia es crucial para comparar el modelo de aprendizaje automático rendimiento con el tiempo a medida que los itera y mejora, antes de la implementación para usuarios reales en sus dispositivos.

Configurar un entorno de pruebas coherente con GPU puede ser más difícil que 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 trabajas en videojuegos web o 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 entorno de Google Colab basado en Linux notebook que esté conectado a una consola GPU T4 o V100. Puedes usar otras plataformas en la nube, como Google Cloud (GCP), si así lo prefieres.
• Navegador: Chrome es compatible con WebGPU, un poderoso 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 permite controlas 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 un visible en un servidor. Usamos la función nuevo modo sin interfaz gráfica, no 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. 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. Mié no están solos en este problema: hay numerosos debates en línea de personas en una situación similar, incluso en los canales oficiales de asistencia de Chrome (1), (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 marcas siguientes cuando inicies Headless. Chrome:

• La marca --no-sandbox inhabilita la zona de pruebas de seguridad de Chrome, que aísla los de navegación del resto del sistema. No se admite la ejecución de Chrome como root sin esta zona de pruebas.
• La función experimental --headless=new ejecuta Chrome con la versión nueva y mejorada 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 el vulkan VK_KHR_surface extensión de instancia. 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 e inhabilita la lista de entidades bloqueadas de adaptadores.

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 detectaron problemas de WebGPU y el valor cambia de inhabilitado solo para software.

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

Instala los controladores de GPU correctos

Investigamos en más detalle el resultado de chrome://gpu, con algunos expertos en GPU del equipo de Chrome. Detectamos problemas con los controladores predeterminados instalados en el Colab de Linux lo que causaba problemas con Vulkan, lo que provocaba que Chrome no pudiera detectar la GPU NVIDIA T4 en el nivel GL_RENDERER, como se muestra en el siguiente resultado. Esta 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. El las dos líneas anteriores 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 obtendremos el siguiente resultado. 🎉

Ahora que usamos los controladores y marcas correctos cuando se ejecuta Chrome, tenemos WebGPU. y la compatibilidad con WebGL mediante el nuevo modo sin interfaz gráfica.

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

Después de muchas investigaciones, no encontramos métodos que funcionaran para el medio ambiente que necesarios para ejecutarse en Google Colab, aunque había algunos publicaciones esperanzadoras que funcionaba en otros entornos, lo que era prometedor. 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. Algunos ejemplos de soluciones de trabajo de terceros usan la interfaz sin interfaz gráfica de Chrome que en algún momento dejará de estar disponible y se usará nueva versión. Necesitábamos una solución que funcionaba con el nuevo Chrome sin interfaz gráfica para estar mejor preparado para el futuro.

Confirmamos el poco uso de la GPU ejecutando un Ejemplo de página web de TensorFlow.js para reconocimiento de imágenes, donde entrenamos un modelo para reconocer muestras de ropa (algo así como en todo el mundo" 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 una. Llamando a Headless Chrome en su estado predeterminado, podríamos registrar el resultado de la consola de JavaScript a la línea de comandos del servidor de Node.js para ver cómo la velocidad de 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 navegador como WebGPU, WebGL y WebAssembly, un modelo de aprendizaje automático las operaciones matemáticas se pueden acelerar aún más en el 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 éxito en el uso conlleva una gran responsabilidad. En este nivel de uso En los sistemas de producción, surge la necesidad de probar la IA del cliente basada en el navegador en un entorno de navegador real, a la vez que es escalable, automatizable y en una configuración de hardware estandarizada conocida.

Al aprovechar el poder combinado de los nuevos Headless Chrome y Puppeteer, pueden probar con confianza esas cargas de trabajo en un entorno estandarizado y replicable en el entorno de nube, lo que garantiza resultados coherentes y confiables.

Conclusión

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

Si esto te resultó útil, menciónalo en LinkedIn, X (anteriormente Twitter) o cualquier red social que usas con el hashtag #WebAI. Sería genial escuchar comentarios para que sepamos que debemos escribir más contenido como este en el futuro.

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

Agradecimientos

Muchas gracias a todos los miembros del equipo de Chrome que ayudaron a depurar el conductor y WebGPU que tuvimos en esta solución, con un agradecimiento especial a Jecelyn Yeen y Alexandra White por ayudar a {i>wordsmith<i} 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.