LiteRT para la Web con LiteRT.js

LiteRT.js es el entorno de ejecución de WebAI de alto rendimiento de Google, diseñado para aplicaciones web de producción. Es una continuación de la pila de LiteRT, que garantiza la compatibilidad con varios frameworks y unifica nuestro tiempo de ejecución principal en todas las plataformas.

LiteRT.js admite las siguientes funciones principales:

  1. Compatibilidad con los modelos de LiteRT en el navegador: Ejecuta modelos con el mejor rendimiento de su clase en la CPU, acelerados a través de XNNPack en WebAssembly (Wasm) y la GPU con la API de WebGPU.
  2. Compatibilidad con varios frameworks: Usa tu framework de AA preferido: PyTorch, Jax o TensorFlow.
  3. Compila sobre canalizaciones existentes: Integra canalizaciones existentes de TensorFlow.js admitiendo tensores de TensorFlow.js como entradas y salidas.

Instalación

Instala el paquete @litertjs/core desde npm:

npm install @litertjs/core

Los archivos Wasm se encuentran en node_modules/@litertjs/core/wasm/. Para mayor comodidad, copia y publica toda la carpeta wasm/. Luego, importa el paquete y carga los archivos Wasm:

import {loadLiteRt} from '@litertjs/core;

// Host LiteRT's Wasm files on your server.
await loadLiteRt(`your/path/to/wasm/`);

Conversión de modelos

LiteRT.js usa el mismo formato .tflite que Android y iOS, y admite modelos existentes en Kaggle y Huggingface. Si tienes un modelo nuevo de PyTorch, deberás convertirlo.

Cómo convertir un modelo de PyTorch a LiteRT

Para convertir un modelo de PyTorch a LiteRT, usa el conversor ai-edge-torch.

import ai_edge_torch

# Load your torch model. We're using resnet for this example.
resnet18 = torchvision.models.resnet18(torchvision.models.ResNet18_Weights.IMAGENET1K_V1)

sample_inputs = (torch.randn(1, 3, 224, 224),)

# Convert the model to LiteRT.
edge_model = ai_edge_torch.convert(resnet18.eval(), sample_inputs)

# Export the model.
edge_model.export('resnet.tflite')

Ejecuta el modelo convertido

Después de convertir el modelo en un archivo .tflite, puedes ejecutarlo en el navegador.

import {loadAndCompile} from '@litertjs/core';

// Load the model hosted from your server. This makes an http(s) request.
const model = await loadAndCompile('/path/to/model.tflite', {
    accelerator: 'webgpu', // or 'wasm' for XNNPack CPU inference
});
// The model can also be loaded from a Uint8Array if you want to fetch it yourself.

// Create image input data
const image = new Float32Array(224 * 224 * 3).fill(0);
const inputTensor =
    await new Tensor(image, /* shape */ [1, 3, 224, 224]).moveTo('webgpu');

// Run the model
const outputs = model(inputTensor);
// You can also use model([inputTensor])
// or model({'input_tensor_name': inputTensor})

// Clean up and get outputs
inputTensor.delete();
const outputTensorCpu = await outputs[0].moveTo('wasm');
const outputData = outputTensorCpu.toTypedArray();
outputTensorCpu.delete();

Integración en canalizaciones existentes de TensorFlow.js

Debes considerar integrar LiteRT.js en tus canalizaciones de TensorFlow.js por los siguientes motivos:

  1. Rendimiento de WebGPU de primer nivel: Los modelos convertidos que se ejecutan en WebGPU de LiteRT.js están optimizados para el rendimiento del navegador y son especialmente rápidos en los navegadores basados en Chromium.
  2. Ruta de conversión de modelos más sencilla: La ruta de conversión de LiteRT.js va directamente de PyTorch a LiteRT. La ruta de conversión de PyTorch a TensorFlow.js es mucho más complicada, ya que requiere que pases de PyTorch -> ONNX -> TensorFlow -> TensorFlow.js.
  3. Herramientas de depuración: La ruta de conversión de LiteRT.js incluye herramientas de depuración.

LiteRT.js está diseñado para funcionar dentro de las canalizaciones de TensorFlow.js y es compatible con el preprocesamiento y el posprocesamiento de TensorFlow.js, por lo que lo único que debes migrar es el modelo en sí.

Sigue estos pasos para integrar LiteRT.js en las canalizaciones de TensorFlow.js:

  1. Convierte tu modelo original de TensorFlow, JAX o PyTorch a .tflite. Para obtener más detalles, consulta la sección Conversión de modelos.
  2. Instala los paquetes @litertjs/core y @litertjs/tfjs-interop de NPM.
  3. Importa y usa el backend de WebGPU de TensorFlow.js. Esto es necesario para que LiteRT.js interopere con TensorFlow.js.
  4. Reemplaza cargar el modelo de TensorFlow.js por cargar el modelo de LiteRT.js.
  5. Sustituye model.predict(entradas) o model.execute(entradas) de TensorFlow.js por runWithTfjsTensors(liteRtModel, entradas). runWithTfjsTensors toma los mismos tensores de entrada que usan los modelos de TensorFlow.js y genera tensores de TensorFlow.js.
  6. Prueba que la canalización del modelo genere los resultados que esperas.

Usar LiteRT.js con runWithTfjsTensors también puede requerir los siguientes cambios en las entradas del modelo:

  1. Reordenar entradas: Según cómo el conversor ordenó las entradas y salidas del modelo, es posible que debas cambiar su orden a medida que las pasas.
  2. Transponer entradas: También es posible que el conversor haya cambiado el diseño de las entradas y salidas del modelo en comparación con lo que usa TensorFlow.js. Es posible que debas transponer tus entradas para que coincidan con el modelo y las salidas para que coincidan con el resto de la canalización.
  3. Cambiar el nombre de las entradas: Si usas entradas con nombre, es posible que los nombres también hayan cambiado.

Puedes obtener más información sobre las entradas y salidas del modelo con model.getInputDetails() y model.getOutputDetails().