LiteRT para la Web con LiteRT.js

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

LiteRT.js admite las siguientes funciones principales:

  1. Inferencia acelerada por hardware en el navegador: Ejecuta modelos con un rendimiento excepcional de la CPU, acelerado por XNNPack asignado a WebAssembly (Wasm) ligero. Para el escalamiento de GPU y hardware dedicado (como las NPU), LiteRT.js muestra de forma nativa la API de WebGPU y la API WebNN emergente, lo que permite una optimización detallada específica de la plataforma.
  2. Compatibilidad con varios frameworks: Agiliza la semántica de desarrollo compilando desde tu framework de AA preferido de forma nativa: PyTorch, JAX o TensorFlow.
  3. Itera en las canalizaciones existentes: Integración lista para usar con las arquitecturas existentes de TensorFlow.js mediante el análisis de tensores de TensorFlow.js compatibles de forma nativa como entradas y salidas de límite directo.

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';

// Load the LiteRT.js Wasm files from a CDN.
await loadLiteRt('https://cdn.jsdelivr.net/npm/@litertjs/core/wasm/')
// Alternatively, host them from your server.
// They are located in node_modules/@litertjs/core/wasm/
await loadLiteRt(`your/path/to/wasm/`);

Conversión de modelos

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

Convierte un modelo de PyTorch a LiteRT

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

import litert_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 = litert_torch.convert(resnet18.eval(), sample_inputs)

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

Ejecuta el modelo convertido

Después de convertir el modelo a 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',
    // Can select from 'webnn', 'webgpu', & 'wasm'.
    // Additionally, you can pass an array of accelerators e.g. ['webnn', 'wasm']
    // if you would like to fallback to CPU execution,
    // Note that ONLY cpu fallback is supported for now
    // (i.e. specifying ['webnn', 'webgpu']) will lead to compilation errors
});
// 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 = new Tensor(image, /* shape */ [1, 3, 224, 224]);

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

// Clean up and get outputs
inputTensor.delete();
const output = outputs[0];
const outputData = await output.data();
output.delete();

Integración en canalizaciones existentes de TensorFlow.js

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

  1. Rendimiento excepcional de la GPU y el hardware: Los modelos de LiteRT.js aprovechan la aceleración de WebGPU para optimizar el rendimiento en las arquitecturas del navegador. Con compatibilidad para WebGPU y la próxima WebNN, LiteRT.js ofrece una aceleración de hardware flexible en una variedad de dispositivos perimetrales.
  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í.

Integra LiteRT.js en las canalizaciones de TensorFlow.js con los siguientes pasos:

  1. Convierte tu modelo original de TensorFlow, JAX o PyTorch a .tflite. Para obtener más detalles, consulta la sección de conversión de modelos.
  2. Instala los paquetes NPM @litertjs/core y @litertjs/tfjs-interop.
  3. Importa y usa el backend de WebGPU de TensorFlow.js. Esto es necesario para que LiteRT.js interoperar con TensorFlow.js.
  4. Reemplaza la carga del modelo de TensorFlow.js por la carga del modelo de LiteRT.js.
  5. Sustituye model.predict(inputs) o model.execute(inputs) de TensorFlow.js por runWithTfjsTensors(liteRtModel, inputs). 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.

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

  1. Reordena las 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. Cambia el nombre de las entradas: Si usas entradas con nombre, es posible que también hayan cambiado los nombres.

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