Comienza a usar LiteRT.js

Esta es una guía integral de LiteRT.js que abarca el proceso de conversión de un modelo de PyTorch para que se ejecute en el navegador con aceleración de WebGPU. En este ejemplo, se usa ResNet18 para el modelo de visión y TensorFlow.js para el preprocesamiento y el posprocesamiento.

En esta guía, se abarcarán los siguientes pasos:

  1. Convierte tu modelo de PyTorch a LiteRT con AI Edge Torch.
    1. Agrega el paquete LiteRT a tu app web.
  2. Carga el modelo.
  3. Escribe la lógica de procesamiento previo y posterior.

Cómo convertir a LiteRT

Usa el notebook de PyTorch Converter para convertir un modelo de PyTorch al formato .tflite adecuado. Para obtener una guía detallada sobre los tipos de errores que puedes encontrar y cómo corregirlos, consulta el README de AI Edge Torch Converter.

Tu modelo debe ser compatible con torch.export.export, lo que significa que debe poder exportarse con TorchDynamo. Por lo tanto, no debe tener ninguna rama condicional de Python que dependa de los valores del tiempo de ejecución dentro de los tensores. Si ves los siguientes errores durante torch.export.export, tu modelo no se puede exportar con torch.export.export. Además, tu modelo no debe tener dimensiones de entrada o salida dinámicas en sus tensores. Esto incluye la dimensión del lote.

También puedes comenzar con un modelo de PyTorch compatible con TensorRT o exportable a ONNX:

  • Una versión de un modelo compatible con TensorRT puede ser un buen punto de partida, ya que algunos tipos de conversiones de TensorRT también requieren que los modelos sean exportables con TorchDynamo. Si usas alguna operación de NVIDIA o CUDA en el modelo, deberás reemplazarla por operaciones estándar de PyTorch.

  • Un modelo de PyTorch exportable a ONNX puede ser un buen punto de partida, aunque algunos modelos de ONNX usan TorchScript en lugar de TorchDynamo para la exportación, en cuyo caso es posible que el modelo no sea exportable a TorchDynamo (aunque es probable que esté más cerca que el código del modelo original).

Para obtener más información, consulta Cómo convertir modelos de PyTorch a LiteRT.

Agrega el paquete de LiteRT

Instala el paquete @litertjs/core desde npm:

npm install @litertjs/core

Importa el paquete y carga sus archivos Wasm:

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

// They are located in node_modules/@litertjs/core/wasm/
// Serve them statically on your server.
await loadLiteRt(`your/path/to/wasm/`);

Carga el modelo

Importa e inicializa LiteRT.js y las utilidades de conversión de LiteRT-TFJS. También debes importar TensorFlow.js para pasar tensores a LiteRT.js.

import {CompileOptions, loadAndCompile, loadLiteRt, setWebGpuDevice} from '@litertjs/core';
import {runWithTfjsTensors} from '@litertjs/tfjs-interop';

// TensorFlow.js imports
import * as tf from '@tensorflow/tfjs';
import '@tensorflow/tfjs-backend-webgpu'; // Only WebGPU is supported
import {WebGPUBackend} from '@tensorflow/tfjs-backend-webgpu';

async function main() {
  // Initialize TensorFlow.js WebGPU backend
  await tf.setBackend('webgpu');

  // Initialize LiteRT.js's Wasm files
  await loadLiteRt('your/path/to/wasm/');

  // Make LiteRt use the same GPU device as TFJS (for tensor conversion)
  const backend = tf.backend() as WebGPUBackend;
  setWebGpuDevice(backend.device);
  // ...
}

main();

Carga el modelo de LiteRT convertido:

const model = await loadAndCompile('path_to_model.tflite', {
  accelerator: 'webgpu', // or 'wasm'
});

Escribe la canalización del modelo

Escribe la lógica de preprocesamiento y posprocesamiento que conecta el modelo a tu app. Se recomienda usar TensorFlow.js para el preprocesamiento y el posprocesamiento, pero, si no está escrito en TensorFlow.js, puedes llamar a await tensor.data para obtener el valor como un ArrayBuffer o await tensor.array para obtener un array de JS estructurado.

A continuación, se muestra un ejemplo de canalización de extremo a extremo para ResNet18:

// Wrap in a tf.tidy call to automatically clean up intermediate TensorFlow.js tensors.
// (Note: tidy only supports synchronous functions).
const top5 = tf.tidy(() => {
  // Get RGB data values from an image element and convert it to range [0, 1).
  const image = tf.browser.fromPixels(dogs, 3).div(255);

  // These preprocessing steps come from https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py#L315
  // The mean and standard deviation for the image normalization come from https://github.com/pytorch/vision/blob/main/torchvision/transforms/_presets.py#L38
  const imageData = image.resizeBilinear([224, 224])
    .sub([0.485, 0.456, 0.406])
    .div([0.229, 0.224, 0.225])
    .reshape([1, 224, 224, 3])
    .transpose([0, 3, 1, 2]);

  // You can pass inputs as a single tensor, an array, or a JS Object
  // where keys are the tensor names in the TFLite model.
  // When passing an Object, the output is also an Object.
  // Here, we're passing a single tensor, so the output is an array.
  const probabilities = runWithTfjsTensors(model, imageData)[0];

  // Get the top five classes.
  return tf.topk(probabilities, 5);
});

const values = await top5.values.data();
const indices = await top5.indices.data();
top5.values.dispose(); // Clean up the tfjs tensors.
top5.indices.dispose();

// Print the top five classes.
const classes = ... // Class names are loaded from a JSON file in the demo.
for (let i = 0; i < 5; ++i) {
  const text = `${classes[indices[i]]}: ${values[i]}`;
  console.log(text);
}

Prueba y solución de problemas

Consulta las siguientes secciones sobre las formas de probar tu aplicación y controlar los errores.

Pruebas con entradas simuladas

Después de cargar un modelo, es una buena idea probarlo primero con entradas falsas. Esto detectará cualquier error de tiempo de ejecución antes de que dediques tiempo a escribir la lógica de procesamiento previo y posterior para tu canalización de modelos. Para verificarlo, puedes usar el verificador de modelos de LiteRT.js o probarlo de forma manual.

Probador de modelos de LiteRT.js

El verificador de modelos de LiteRT.js ejecuta tu modelo en la GPU y la CPU con entradas aleatorias para verificar que el modelo se ejecute correctamente en la GPU. Verifica lo siguiente:

  • Indica si se admiten los tipos de datos de entrada y salida.
  • Indica si todas las operaciones están disponibles en la GPU.
  • Indica qué tan cerca están los resultados de la GPU de los resultados de referencia de la CPU.
  • El rendimiento de la inferencia en la GPU

Para ejecutar el verificador de modelos de LiteRT.js, ejecuta npm i @litertjs/model-tester y, luego, npx model-tester. Se abrirá una pestaña del navegador para que ejecutes tu modelo.

Pruebas manuales de modelos

Si prefieres probar el modelo de forma manual en lugar de usar el verificador de modelos de LiteRT.js (@litertjs/model-tester), puedes generar entradas falsas y ejecutar el modelo con runWithTfjsTensors.

Para generar entradas falsas, debes conocer los nombres y las formas de los tensores de entrada. Se pueden encontrar con LiteRT.js llamando a model.getInputDetails o model.getOutputDetails. Una forma sencilla de encontrarlos es establecer un punto de interrupción después de que se crea el modelo. También puedes usar el Explorador de modelos.

Una vez que conozcas las formas y los nombres de entrada y salida, podrás probar el modelo con una entrada falsa. Esto te da cierta confianza de que el modelo se ejecutará antes de que escribas el resto de la canalización de aprendizaje automático. Esto probaría que se admiten todas las operaciones del modelo. Por ejemplo:

// Imports, initialization, and model loading...
// Create fake inputs for the model
const fakeInputs = model.getInputDetails().map(
    ({shape, dtype}) => tf.ones(shape, dtype));

// Run the model
const outputs = runWithTfjsTensors(model, fakeInputs);
console.log(outputs);

Tipos de errores

Es posible que LiteRT.js no admita algunos modelos de LiteRT. Por lo general, los errores se clasifican en las siguientes categorías:

  • Shape Mismatch: Es un error conocido que solo afecta a la GPU.
  • Operation Not Supported: El tiempo de ejecución no admite una operación en el modelo. El backend de WebGPU tiene una cobertura más limitada que la CPU, por lo que, si ves este error en la GPU, es posible que puedas ejecutar el modelo en la CPU.
  • Tipo de tensor no admitido: LiteRT.js solo admite tensores int32 y float32 para las entradas y salidas del modelo.
  • Modelo demasiado grande: LiteRT.js tiene un límite en el tamaño de los modelos que puede cargar.

No se admite la operación

Esto indica que el backend que se usa no admite una de las operaciones del modelo. Deberás volver a escribir el modelo original de PyTorch para evitar esta operación y volver a convertirlo, o bien podrás ejecutar el modelo en la CPU.

En el caso de BROADCAST_TO, esto se puede solucionar haciendo que la dimensión del lote sea la misma para cada tensor de entrada del modelo. Otros casos pueden ser más complicados.

Tipo de tensor no admitido

LiteRT.js solo admite tensores int32 y float32 para las entradas y salidas del modelo.

Modelo demasiado grande

Por lo general, esto aparece como una llamada a Aborted() o una falla en la asignación de memoria durante la carga del modelo. LiteRT.js tiene un límite en el tamaño de los modelos que puede cargar, por lo que, si ves este mensaje, es posible que tu modelo sea demasiado grande. Puedes intentar cuantificar los pesos con ai-edge-quantizer, pero mantén los cálculos en float32 o float16, y las entradas y salidas del modelo como float32 o int32.