LiteRT per il web con LiteRT.js

LiteRT.js è il runtime WebAI ad alte prestazioni di Google, destinato alle applicazioni web di produzione. È una continuazione dello stack LiteRT, che garantisce il supporto multi-framework e unifica il nostro runtime di base su tutte le piattaforme.

LiteRT.js supporta le seguenti funzionalità di base:

  1. Supporto in-browser per i modelli LiteRT: esegui modelli con prestazioni ottimali su CPU, accelerati tramite XNNPack su WebAssembly (Wasm) e GPU utilizzando l'API WebGPU.
  2. Compatibilità con più framework: utilizza il framework ML che preferisci: PyTorch, Jax o TensorFlow.
  3. Crea pipeline esistenti: esegui l'integrazione con le pipeline TensorFlow.js esistenti supportando i tensori TensorFlow.js come input e output.

Installazione

Installa il pacchetto @litertjs/core da npm:

npm install @litertjs/core

I file Wasm si trovano in node_modules/@litertjs/core/wasm/. Per comodità, copia e pubblica l'intera cartella wasm/. Quindi, importa il pacchetto e carica i file Wasm:

import {loadLiteRt} from '@litertjs/core;

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

Conversione del modello

LiteRT.js utilizza lo stesso formato .tflite di Android e iOS e supporta i modelli esistenti su Kaggle e Huggingface. Se hai un nuovo modello PyTorch, dovrai convertirlo.

Converti un modello PyTorch in LiteRT

Per convertire un modello PyTorch in LiteRT, utilizza lo strumento di conversione 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')

Esegui il modello convertito

Dopo aver convertito il modello in un file .tflite, puoi eseguirlo nel browser.

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

Eseguire l'integrazione nelle pipeline TensorFlow.js esistenti

Ti consigliamo di integrare LiteRT.js nelle pipeline TensorFlow.js per i seguenti motivi:

  1. Prestazioni WebGPU migliori della categoria: i modelli convertiti in esecuzione su LiteRT.js WebGPU sono ottimizzati per le prestazioni del browser e sono particolarmente veloci sui browser basati su Chromium.
  2. Percorso di conversione del modello più semplice: il percorso di conversione LiteRT.js va direttamente da PyTorch a LiteRT. Il percorso di conversione da PyTorch a TensorFlow.js è molto più complicato e richiede di passare da PyTorch -> ONNX -> TensorFlow -> TensorFlow.js.
  3. Strumenti di debug: il percorso di conversione LiteRT.js include strumenti di debug.

LiteRT.js è progettato per funzionare all'interno delle pipeline TensorFlow.js ed è compatibile con la pre-elaborazione e la post-elaborazione di TensorFlow.js, quindi l'unica cosa che devi migrare è il modello stesso.

Integra LiteRT.js nelle pipeline TensorFlow.js seguendo questi passaggi:

  1. Converti il modello TensorFlow, JAX o PyTorch originale in .tflite. Per maggiori dettagli, consulta la sezione Conversione del modello.
  2. Installa i pacchetti NPM @litertjs/core e @litertjs/tfjs-interop.
  3. Importa e utilizza il backend TensorFlow.js WebGPU. Questo è necessario per l'interoperabilità di LiteRT.js con TensorFlow.js.
  4. Sostituisci il caricamento del modello TensorFlow.js con il caricamento del modello LiteRT.js.
  5. Sostituisci TensorFlow.js model.predict(input) o model.execute(input) con runWithTfjsTensors(liteRtModel, input). runWithTfjsTensors accetta gli stessi tensori di input utilizzati dai modelli TensorFlow.js e restituisce tensori TensorFlow.js.
  6. Verifica che la pipeline del modello restituisca i risultati previsti.

L'utilizzo di LiteRT.js con runWithTfjsTensors potrebbe richiedere anche le seguenti modifiche agli input del modello:

  1. Riordina gli input: a seconda di come il convertitore ha ordinato gli input e gli output del modello, potrebbe essere necessario modificarne l'ordine durante il passaggio.
  2. Trasposizione degli input: è anche possibile che il convertitore abbia modificato il layout degli input e degli output del modello rispetto a quello utilizzato da TensorFlow.js. Potrebbe essere necessario trasporre gli input in modo che corrispondano al modello e gli output in modo che corrispondano al resto della pipeline.
  3. Rinomina input: se utilizzi input denominati, anche i nomi potrebbero essere cambiati.

Puoi ottenere maggiori informazioni sugli input e sugli output del modello con model.getInputDetails() e model.getOutputDetails().