Pruebas de modelos de IA web en Google Colab

Jason Mayes

François Beaufort

Configurar un entorno de prueba coherente con GPUs puede ser más difícil de lo esperado. Estos son los pasos para probar modelos de IA del cliente basados en el navegador en entornos de navegador reales, a la vez que son escalables, automatizables y dentro de una configuración de hardware estandarizada conocida.

En este caso, el navegador es un navegador Chrome real con compatibilidad de hardware, en lugar de una emulación de software.

Ya sea que seas desarrollador de IA web, juegos web o gráficos, o te sientas interesado en las pruebas de modelos de IA web, esta guía es para ti.

Paso 1: Crea un nuevo notebook de Google Colab

Paso 2: Conéctate a un servidor compatible con GPU T4

1. Haz clic en Conectar arrow_drop_down cerca de la esquina superior derecha de la notebook.
2. Selecciona Cambiar el tipo de entorno de ejecución:

   

3. En la ventana modal, selecciona GPU T4 como acelerador de hardware. Cuando te conectes, Colab usará una instancia de Linux con una GPU NVIDIA T4 conectada.

   

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 de RAM y disco. 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 se antepone con 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
   

   • Puedes inspeccionar la secuencia de comandos en GitHub para ver el código de línea de comandos sin procesar que ejecuta.

   # 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 play_circle junto a la celda para ejecutar el código.

   

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

   

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

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

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 folder. Luego, haz clic en more_vert para descargar el archivo PDF en tu máquina local.

Parte B: Comandos de Chrome con Puppeteer

Proporcionamos un ejemplo minimalista que usa Puppeteer para controlar Headless Chrome 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 nodo 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 línea de comandos se pasaron a la aplicación de Node. Asegúrate de que el tercer argumento esté configurado, lo que representa una URL a la que navegar. Debes inspeccionar el tercer argumento aquí porque los dos primeros llaman a Node y a la secuencia de comandos que ejecutamos. En realidad, el tercer elemento contiene el primer parámetro que se pasa al programa 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(). Esto crea un objeto de navegador que se configura 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 podrás usar más adelante para visitar cualquier URL:

const page = await browser.newPage();

Luego, agrega un objeto de escucha de eventos para detectar 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 active una captura de pantalla y, luego, finalice el proceso del navegador en Node. Esto es útil para las páginas web que deben realizar una cantidad de trabajo antes de que se pueda tomar una captura de pantalla y tienen una cantidad no determinista de tiempo de ejecución.

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, ordena 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 se muestre algún resultado de la consola, ya que esta página no está controlada por 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

Es posible que hayas notado que usamos una sentencia de importación al comienzo del archivo jPuppet.js. Tu package.json debe establecer los valores de tipo como module, o 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! El uso de Puppeteer facilita la interacción con Chrome de manera programática.

Listo

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

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

Recursos

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