Pruebas de modelos de IA web en Google Colab

Jason Mayes
François Beaufort
François Beaufort
GitHub

Configurar un entorno de pruebas coherente con GPU 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 compatible con el hardware, a diferencia de la emulación de software.

Esta guía es para ti si eres desarrollador de IA web, juegos web o gráficos, o si te interesa probar modelos de IA web.

Paso 1: Crea un nuevo notebook de Google Colab

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

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

  1. Haz clic en Conectar cerca de la esquina superior derecha de la notebook.
  2. Selecciona Cambiar el tipo de entorno de ejecución:
    Una captura de pantalla en primer plano que muestra los pasos para cambiar el entorno 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 acelerador de hardware. Cuando te conectes, Colab usará una instancia de Linux con una GPU NVIDIA T4 conectada.
    Captura de pantalla del módulo Cambiar el tipo de tiempo 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 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

    # 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.

    Una captura de pantalla de un Colab nuevo
    Figura 4.

  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.

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

Una captura de pantalla de un Colab nuevo
Figura 6: Mira 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 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 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 línea de comandos se pasaron a la aplicación de Node. Asegúrate de que el tercer argumento esté configurado, 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 necesitan realizar una cierta cantidad de trabajo antes de que se pueda tomar una captura de pantalla y tiene 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, 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; 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! El uso de Puppeteer facilita la interacción con Chrome de manera programática.

Listo

Ahora podemos verificar que el clasificador 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 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.

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

Recursos

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