Pruebas de modelos de IA web en Google Colab

François Beaufort
François Beaufort

Configurar un entorno de pruebas coherente con GPU puede ser más difícil de lo esperado. Estos son los pasos para probar los modelos de IA del cliente basados en el navegador en entornos de navegador reales, al mismo tiempo que son escalables, automatizables y se encuentran dentro de una configuración de hardware estandarizada conocida.

En este caso, se trata de un navegador Chrome real con compatibilidad de hardware, a diferencia de la emulación de software.

Si eres un desarrollador de IA web, videojuegos web o gráficos, o si te interesan las pruebas de modelos de IA web, esta guía es para ti.

Paso 1: Crea un nuevo notebook de Google Colab

1. Ve a colab.new para crear un notebook de Colab nuevo. Debería verse similar a la figura 1. 2. Sigue las indicaciones para acceder a tu Cuenta de Google.
Captura de pantalla de un nuevo Colab
Figura 1: Un nuevo notebook de Colab

Paso 2: Conéctate a un servidor T4 habilitado para GPU

  1. Haz clic en Conectar cerca de la parte superior derecha del notebook.
  2. Selecciona Cambiar tipo de entorno de ejecución:
    Una captura de pantalla en primer plano que muestra los pasos para cambiar el tiempo de ejecución.
    Figura 2. Cambia el entorno de ejecución en la interfaz de Colab.
  3. En la ventana modal, selecciona GPU T4 como el acelerador de hardware. Cuando te conectes, Colab usará una instancia de Linux con una GPU NVIDIA T4 conectada.
    Captura de pantalla del módulo Cambia el tipo de entorno de ejecución.
    Figura 3: En Acelerador de hardware, selecciona GPU T4.
  4. Haz clic en Guardar.
  5. Haz clic en el botón Conectar para conectarte a tu entorno de ejecución. Después de un tiempo, el botón mostrará una marca de verificación verde, junto con gráficos de uso del disco y la memoria RAM. Esto indica que se creó correctamente un servidor con el hardware requerido.

Buen trabajo. Acabas de crear un servidor con una GPU conectada.

Paso 3: Instala los controladores y las dependencias correctos

  1. Copia y pega las siguientes dos líneas de código en la primera celda de código del notebook. En un entorno de Colab, la ejecución de la línea de comandos está precedida por un signo de exclamación.

    !git clone https://github.com/jasonmayes/headless-chrome-nvidia-t4-gpu-support.git
    !cd headless-chrome-nvidia-t4-gpu-support && chmod +x scriptyMcScriptFace.sh && ./scriptyMcScriptFace.sh
    
    # Update, install correct drivers, and remove the old ones.
    apt-get install -y vulkan-tools libnvidia-gl-525
    
    # Verify NVIDIA drivers can see the T4 GPU and that vulkan is working correctly.
    nvidia-smi
    vulkaninfo --summary
    
    # Now install latest version of Node.js
    npm install -g n
    n lts
    node --version
    npm --version
    
    # Next install Chrome stable
    curl -fsSL https://dl.google.com/linux/linux_signing_key.pub | sudo gpg --dearmor -o /usr/share/keyrings/googlechrom-keyring.gpg
    echo "deb [arch=amd64 signed-by=/usr/share/keyrings/googlechrom-keyring.gpg] http://dl.google.com/linux/chrome/deb/ stable main" | sudo tee /etc/apt/sources.list.d/google-chrome.list
    sudo apt update
    sudo apt install -y google-chrome-stable
    
    # Start dbus to avoid warnings by Chrome later.
    export DBUS_SESSION_BUS_ADDRESS="unix:path=/var/run/dbus/system_bus_socket"
    /etc/init.d/dbus start
    
  2. Haz clic en junto a la celda para ejecutar el código.

    Captura de pantalla de un nuevo Colab
    Figura 4:

  3. Una vez que termine de ejecutarse el código, verifica que nvidia-smi haya impreso algo similar a la siguiente captura de pantalla para confirmar que realmente tienes una GPU conectada y que se reconoce en tu servidor. Es posible que debas desplazarte a una parte anterior de los registros para ver este resultado.

    Figura 5: Busca el resultado que comienza con "NVIDIA-SMI".

Paso 4: Usa y automatiza Chrome sin interfaz gráfica

  1. Haz clic en el botón Code para agregar una nueva celda de código.
  2. Luego, puedes escribir el código personalizado para llamar a un proyecto de Node.js con tus parámetros preferidos (o solo llamar a google-chrome-stable directamente en la línea de comandos). Tenemos ejemplos de ambos casos.

Parte A: Usa Chrome sin interfaz gráfica directamente en la línea de comandos

# Directly call Chrome to dump a PDF of WebGPU testing page
# and store it in /content/gpu.pdf
!google-chrome-stable \
--no-sandbox \
--headless=new \
--use-angle=vulkan \
--enable-features=Vulkan \
--disable-vulkan-surface \
--enable-unsafe-webgpu \
--print-to-pdf=/content/gpu.pdf https://webgpureport.org

En el ejemplo, almacenamos la captura de PDF resultante en /content/gpu.pdf. Para ver ese archivo, expande el contenido . Luego, haz clic en para descargar el archivo PDF en tu máquina local.

Captura de pantalla de un nuevo Colab
Figura 6: Se muestran los pasos para descargar el PDF en esta captura de pantalla de la interfaz de Colab.

Parte B: Comandos de Chrome con Puppeteer

Proporcionamos un ejemplo minimalista en el que se usa Puppeteer para controlar Chrome sin interfaz gráfica que se puede ejecutar de la siguiente manera:

# Call example node.js project to perform any task you want by passing
# a URL as a parameter
!node headless-chrome-nvidia-t4-gpu-support/examples/puppeteer/jPuppet.js chrome://gpu

En el ejemplo de jPuppet, podemos llamar a una secuencia de comandos de Node.js para crear una captura de pantalla. Pero ¿cómo funciona? Consulta esta explicación del código de Node.js en jPuppet.js.

Desglose del código de Node de jPuppet.js

Primero, importa Puppeteer. Esto te permite controlar Chrome de forma remota con Node.js:

import puppeteer from 'puppeteer';

A continuación, verifica qué argumentos de la línea de comandos se pasaron a la aplicación de Node. Asegúrate de que esté configurado el tercer argumento, que representa una URL a la que navegar. Debes inspeccionar el tercer argumento aquí porque los dos primeros se llaman a Node en sí y a la secuencia de comandos que estamos ejecutando. En realidad, el tercer elemento contiene el primer parámetro que se pasó al programa de Node:

const url = process.argv[2];
if (!url) {
  throw "Please provide a URL as the first argument";
}

Ahora define una función asíncrona llamada runWebpage(). De esta manera, se crea un objeto de navegador configurado con los argumentos de línea de comandos para ejecutar el objeto binario de Chrome de la manera en que necesitamos que WebGL y WebGPU funcionen, como se describe en Cómo habilitar la compatibilidad con WebGPU y WebGL.

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

Crea un nuevo objeto de página del navegador que luego puedas usar para visitar cualquier URL:

const page = await browser.newPage();

Luego, agrega un objeto de escucha de eventos para escuchar los eventos console.log cuando la página web ejecute JavaScript. Esto te permite registrar mensajes en la línea de comandos de Node y también inspeccionar el texto de la consola en busca de una frase especial (en este caso, captureAndEnd) que activa una captura de pantalla y, luego, finaliza el proceso del navegador en Node. Esto es útil para las páginas web que necesitan realizar una cantidad de trabajo antes de tomar una captura de pantalla y tienen un tiempo de ejecución no determinista.

page.on('console', async function(msg) {
  console.log(msg.text());
  if (msg.text() === 'captureAndEnd') {
    await page.screenshot({ path: '/content/screenshotEnd.png' });
    await browser.close();
  }
});

Por último, solicita a la página que visite la URL especificada y toma una captura de pantalla inicial cuando se cargue la página.

Si decides tomar una captura de pantalla de chrome://gpu, puedes cerrar la sesión del navegador de inmediato en lugar de esperar a que aparezca algún resultado de la consola, ya que esta página no se controla con tu propio código.

  await page.goto(url,  { waitUntil: 'networkidle2' });
  await page.screenshot({path: '/content/screenshot.png'});
  if (url === 'chrome://gpu') {
    await browser.close();
  }
}
runWebpage();

Modifica package.json

Tal vez notaste que usamos una sentencia de importación al comienzo del archivo jPuppet.js. Tu package.json debe establecer los valores de tipo como module, de lo contrario, recibirás un error que indica que el módulo no es válido.

 {
    "dependencies":  {
      "puppeteer": "*"
    },
    "name": "content",
    "version": "1.0.0",
    "main": "jPuppet.js",
    "devDependencies": {},
    "keywords": [],
    "type": "module",
    "description": "Node.js Puppeteer application to interface with headless Chrome with GPU support to capture screenshots and get console output from target webpage"
}

¡Eso es todo! Puppeteer facilita la interfaz con Chrome de manera programática.

Completada correctamente

Ahora, podemos verificar que el clasificador de Fashion MNIST de TensorFlow.js pueda reconocer correctamente un par de pantalones en una imagen, con procesamiento del cliente en el navegador mediante la GPU.

Puedes usarlo para cualquier carga de trabajo basada en GPU del cliente, desde modelos de aprendizaje automático hasta pruebas de gráficos y juegos.

Captura de pantalla de un nuevo Colab
Figura 7: Una captura exitosa de un modelo de TensorFlow.js acelerado por GPU capaz de reconocer en tiempo real el lado del cliente de indumentaria en el navegador

Recursos

Destaca el repositorio de GitHub para recibir actualizaciones futuras.