LiteRT pour le Web avec LiteRT.js

LiteRT.js est l'environnement d'exécution WebAI hautes performances de Google, destiné aux applications Web de production. Il s'agit d'une continuation de la pile LiteRT, qui assure la prise en charge multifragment et unifie notre cœur d'exécution sur toutes les plates-formes.

LiteRT.js est compatible avec les fonctionnalités de base suivantes :

  1. Compatibilité dans le navigateur pour les modèles LiteRT : exécutez des modèles avec les meilleures performances sur le processeur, accélérés via XNNPack sur WebAssembly (Wasm) et le GPU à l'aide de l'API WebGPU.
  2. Compatibilité multi-framework : utilisez le framework de ML de votre choix (PyTorch, Jax ou TensorFlow).
  3. Exploitez les pipelines existants : intégrez-les aux pipelines TensorFlow.js existants en acceptant les Tensors TensorFlow.js comme entrées et sorties.

Installation

Installez le package @litertjs/core depuis npm :

npm install @litertjs/core

Les fichiers Wasm se trouvent dans node_modules/@litertjs/core/wasm/. Pour plus de commodité, copiez et diffusez l'intégralité du dossier wasm/. Importez ensuite le package et chargez les fichiers Wasm :

import {loadLiteRt} from '@litertjs/core;

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

Conversion de modèles

LiteRT.js utilise le même format .tflite qu'Android et iOS, et est compatible avec les modèles existants sur Kaggle et Huggingface. Si vous disposez d'un nouveau modèle PyTorch, vous devrez le convertir.

Convertir un modèle PyTorch en LiteRT

Pour convertir un modèle PyTorch en LiteRT, utilisez le convertisseur 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')

Exécuter le modèle converti

Après avoir converti le modèle en fichier .tflite, vous pouvez l'exécuter dans le navigateur.

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

Intégration aux pipelines TensorFlow.js existants

Vous devriez envisager d'intégrer LiteRT.js à vos pipelines TensorFlow.js pour les raisons suivantes :

  1. Performances WebGPU de premier ordre : les modèles convertis s'exécutant sur LiteRT.js WebGPU sont optimisés pour les performances du navigateur et sont particulièrement rapides sur les navigateurs basés sur Chromium.
  2. Chemin de conversion de modèle plus simple : le chemin de conversion LiteRT.js passe directement de PyTorch à LiteRT. Le chemin de conversion de PyTorch vers TensorFlow.js est beaucoup plus complexe. Il vous oblige à passer de PyTorch à ONNX, puis à TensorFlow et enfin à TensorFlow.js.
  3. Outils de débogage : le chemin de conversion LiteRT.js est fourni avec des outils de débogage.

LiteRT.js est conçu pour fonctionner dans les pipelines TensorFlow.js et est compatible avec le prétraitement et le post-traitement TensorFlow.js. Vous n'avez donc qu'à migrer le modèle lui-même.

Pour intégrer LiteRT.js dans les pipelines TensorFlow.js, procédez comme suit :

  1. Convertissez votre modèle TensorFlow, JAX ou PyTorch d'origine en .tflite. Pour en savoir plus, consultez la section Conversion de modèles.
  2. Installez les packages NPM @litertjs/core et @litertjs/tfjs-interop.
  3. Importez et utilisez le backend WebGPU TensorFlow.js. Cela est nécessaire pour que LiteRT.js puisse interagir avec TensorFlow.js.
  4. Remplacez loading the TensorFlow.js model (chargement du modèle TensorFlow.js) par loading the LiteRT.js model (chargement du modèle LiteRT.js).
  5. Remplacez TensorFlow.js model.predict(inputs) ou model.execute(inputs) par runWithTfjsTensors(liteRtModel, inputs). runWithTfjsTensors utilise les mêmes Tensors d'entrée que les modèles TensorFlow.js et génère des Tensors TensorFlow.js.
  6. Vérifiez que le pipeline du modèle génère les résultats attendus.

L'utilisation de LiteRT.js avec runWithTfjsTensors peut également nécessiter les modifications suivantes des entrées du modèle :

  1. Réorganiser les entrées : selon la façon dont le convertisseur a ordonné les entrées et les sorties du modèle, vous devrez peut-être modifier leur ordre lorsque vous les transmettez.
  2. Transposer les entrées : il est également possible que le convertisseur ait modifié la mise en page des entrées et des sorties du modèle par rapport à ce que TensorFlow.js utilise. Vous devrez peut-être transposer vos entrées pour qu'elles correspondent au modèle et vos sorties pour qu'elles correspondent au reste du pipeline.
  3. Renommer les entrées : si vous utilisez des entrées nommées, il est possible que leurs noms aient également changé.

Vous pouvez obtenir plus d'informations sur les entrées et les sorties du modèle avec model.getInputDetails() et model.getOutputDetails().