Premiers pas avec LiteRT.js

Ce guide LiteRT.js de bout en bout couvre le processus de conversion d'un modèle PyTorch pour qu'il s'exécute dans le navigateur avec l'accélération WebGPU. Cet exemple utilise ResNet18 pour le modèle de vision et TensorFlow.js pour le prétraitement et le post-traitement.

Ce guide décrit les opérations suivantes :

  1. Convertissez votre modèle PyTorch en LiteRT à l'aide d'AI Edge Torch.
    1. Ajoutez le package LiteRT à votre application Web.
  2. Chargez le modèle.
  3. Écrivez la logique de prétraitement et de post-traitement.

Convertir en LiteRT

Utilisez le notebook PyTorch Converter pour convertir un modèle PyTorch au format .tflite approprié. Pour obtenir un guide détaillé sur les types d'erreurs que vous pouvez rencontrer et sur la manière de les corriger, consultez le fichier README du convertisseur Torch AI Edge.

Votre modèle doit être compatible avec torch.export.export, ce qui signifie qu'il doit pouvoir être exporté avec TorchDynamo. Par conséquent, il ne doit pas comporter de branches conditionnelles Python qui dépendent des valeurs d'exécution dans les Tensors. Si les erreurs suivantes s'affichent lors de torch.export.export, votre modèle ne peut pas être exporté avec torch.export.export. Votre modèle ne doit pas non plus comporter de dimensions d'entrée ou de sortie dynamiques sur ses Tensors. Cela inclut la dimension de lot.

Vous pouvez également commencer avec un modèle PyTorch compatible avec TensorRT ou exportable au format ONNX :

  • Une version d'un modèle compatible avec TensorRT peut être un bon point de départ, car certains types de conversions TensorRT nécessitent également que les modèles soient exportables avec TorchDynamo. Si vous utilisez des opérations NVIDIA / CUDA dans le modèle, vous devrez les remplacer par des opérations PyTorch standards.

  • Un modèle PyTorch exportable au format ONNX peut être un bon point de départ, bien que certains modèles ONNX utilisent TorchScript au lieu de TorchDynamo pour l'exportation. Dans ce cas, il est possible que le modèle ne soit pas exportable avec TorchDynamo (bien qu'il soit probablement plus proche du code du modèle d'origine).

Pour en savoir plus, consultez Convertir des modèles PyTorch en LiteRT.

Ajouter le package LiteRT

Installez le package @litertjs/core depuis npm :

npm install @litertjs/core

Importez le package et chargez ses fichiers 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/`);

Charger le modèle

Importez et initialisez LiteRT.js et les utilitaires de conversion LiteRT-TFJS. Vous devez également importer TensorFlow.js pour transmettre des Tensors à 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();

Chargez le modèle LiteRT converti :

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

Écrire le pipeline du modèle

Écrivez la logique de prétraitement et de post-traitement qui relie le modèle à votre application. Il est recommandé d'utiliser TensorFlow.js pour le prétraitement et le post-traitement, mais si ce n'est pas écrit dans TensorFlow.js, vous pouvez appeler await tensor.data pour obtenir la valeur sous forme d'ArrayBuffer ou await tensor.array pour obtenir un tableau JS structuré.

Voici un exemple de pipeline de bout en bout pour 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);
}

Test et dépannage

Consultez les sections suivantes pour découvrir comment tester votre application et gérer les erreurs.

Tester avec de fausses entrées

Après avoir chargé un modèle, il est conseillé de le tester d'abord avec de fausses entrées. Cela permettra de détecter les erreurs d'exécution avant que vous ne passiez du temps à écrire la logique de pré et post-traitement pour votre pipeline de modèle. Pour vérifier cela, vous pouvez utiliser l'outil de test de modèle LiteRT.js ou effectuer un test manuel.

Testeur de modèle LiteRT.js

Le testeur de modèle LiteRT.js exécute votre modèle sur GPU et CPU à l'aide d'entrées aléatoires pour vérifier que le modèle s'exécute correctement sur GPU. Il vérifie les points suivants :

  • Indique si les types de données d'entrée et de sortie sont acceptés.
  • Indique si toutes les opérations sont disponibles sur le GPU.
  • Dans quelle mesure les sorties GPU correspondent aux sorties CPU de référence.
  • Performances de l'inférence GPU.

Pour exécuter le testeur de modèle LiteRT.js, exécutez npm i @litertjs/model-tester, puis npx model-tester. Un onglet de navigateur s'ouvre pour vous permettre d'exécuter votre modèle.

Test manuel du modèle

Si vous préférez tester manuellement le modèle au lieu d'utiliser le testeur de modèle LiteRT.js (@litertjs/model-tester), vous pouvez générer de fausses entrées et exécuter le modèle avec runWithTfjsTensors.

Pour générer des entrées factices, vous devez connaître les noms et les formes des Tensors d'entrée. Vous pouvez les trouver avec LiteRT.js en appelant model.getInputDetails ou model.getOutputDetails. Pour les trouver facilement, vous pouvez définir un point d'arrêt après la création du modèle. Vous pouvez également utiliser l'explorateur de modèles.

Une fois que vous connaissez les formes et les noms d'entrée et de sortie, vous pouvez tester le modèle avec une fausse entrée. Cela vous donne une certaine confiance quant à l'exécution du modèle avant d'écrire le reste du pipeline de machine learning. Cela permet de vérifier que toutes les opérations du modèle sont prises en charge. Exemple :

// 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);

Types d'erreurs

Il est possible que LiteRT.js ne soit pas compatible avec certains modèles LiteRT. Les erreurs appartiennent généralement aux catégories suivantes :

  • Incompatibilité de forme : bug connu qui n'affecte que le GPU.
  • Opération non acceptée : l'environnement d'exécution n'accepte pas une opération dans le modèle. Le backend WebGPU a une couverture plus limitée que le CPU. Si vous voyez cette erreur sur le GPU, vous pourrez peut-être exécuter le modèle sur le CPU à la place.
  • Type de Tensor non compatible : LiteRT.js n'accepte que les Tensors int32 et float32 pour les entrées et sorties de modèle.
  • Modèle trop volumineux : LiteRT.js est limité en termes de taille des modèles qu'il peut charger.

Opération non acceptée

Cela indique que le backend utilisé n'est pas compatible avec l'une des opérations du modèle. Vous devrez réécrire le modèle PyTorch d'origine pour éviter cette opération et le reconvertir, ou vous pourrez peut-être exécuter le modèle sur le processeur.

Dans le cas de BROADCAST_TO, ce problème peut être résolu en définissant la même dimension de lot pour chaque Tensor d'entrée du modèle. D'autres cas peuvent être plus complexes.

Type de Tensor non accepté

LiteRT.js n'accepte que les Tensors int32 et float32 pour les entrées et les sorties du modèle.

Modèle trop volumineux

Cela se présente généralement sous la forme d'un appel à Aborted() ou d'un échec d'allocation de mémoire lors du chargement du modèle. LiteRT.js est limité en termes de taille des modèles qu'il peut charger. Si vous voyez ce message, il est possible que votre modèle soit trop volumineux. Vous pouvez essayer de quantifier les poids avec ai-edge-quantizer, mais conservez les calculs en float32 ou float16, et les entrées et sorties du modèle en float32 ou int32.