LiteRT für Web mit LiteRT.js

LiteRT.js ist die leistungsstarke WebAI-Laufzeit von Google für Webanwendungen in der Produktion. Es ist eine Fortsetzung des LiteRT-Stacks, der Multi-Framework-Unterstützung bietet und unsere Core-Laufzeitumgebung auf allen Plattformen vereinheitlicht.

LiteRT.js unterstützt die folgenden Kernfunktionen:

  1. In-Browser-Unterstützung für LiteRT-Modelle: Modelle mit erstklassiger Leistung auf der CPU ausführen, beschleunigt über XNNPack auf WebAssembly (Wasm) und GPU über die WebGPU API.
  2. Kompatibilität mit mehreren Frameworks: Sie können Ihr bevorzugtes ML-Framework verwenden: PyTorch, Jax oder TensorFlow.
  3. Auf vorhandenen Pipelines aufbauen: Durch die Unterstützung von TensorFlow.js-Tensoren als Ein- und Ausgaben können Sie die Integration in vorhandene TensorFlow.js-Pipelines vornehmen.

Installation

Installieren Sie das Paket @litertjs/core von npm:

npm install @litertjs/core

Die Wasm-Dateien befinden sich in node_modules/@litertjs/core/wasm/. Kopieren Sie den gesamten wasm/-Ordner und stellen Sie ihn bereit. Importieren Sie dann das Paket und laden Sie die Wasm-Dateien:

import {loadLiteRt} from '@litertjs/core;

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

Modellkonvertierung

LiteRT.js verwendet dasselbe .tflite-Format wie Android und iOS und unterstützt vorhandene Modelle auf Kaggle und Huggingface. Wenn Sie ein neues PyTorch-Modell haben, müssen Sie es konvertieren.

PyTorch-Modell in LiteRT konvertieren

Wenn Sie ein PyTorch-Modell in LiteRT konvertieren möchten, verwenden Sie den ai-edge-torch-Converter.

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

Konvertiertes Modell ausführen

Nachdem Sie das Modell in eine .tflite-Datei konvertiert haben, können Sie es im Browser ausführen.

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.run(inputTensor);
// You can also use model.run([inputTensor])
// or model.run({'input_tensor_name': inputTensor})

// Clean up and get outputs
inputTensor.delete();
const outputTensorCpu = await outputs[0].moveTo('wasm');
const outputData = outputTensorCpu.toTypedArray();
outputTensorCpu.delete();

In vorhandene TensorFlow.js-Pipelines einbinden

Aus den folgenden Gründen sollten Sie LiteRT.js in Ihre TensorFlow.js-Pipelines einbinden:

  1. Erstklassige WebGPU-Leistung: Konvertierte Modelle, die auf LiteRT.js WebGPU ausgeführt werden, sind für die Browserleistung optimiert und besonders schnell in Chromium-basierten Browsern.
  2. Einfacherer Modellkonvertierungspfad: Der LiteRT.js-Konvertierungspfad führt direkt von PyTorch zu LiteRT. Die Konvertierung von PyTorch zu TensorFlow.js ist wesentlich komplizierter, da Sie den folgenden Pfad durchlaufen müssen: PyTorch –> ONNX –> TensorFlow –> TensorFlow.js.
  3. Debugging-Tools: Der LiteRT.js-Conversion-Pfad enthält Debugging-Tools.

LiteRT.js ist für die Verwendung in TensorFlow.js-Pipelines konzipiert und mit der Vor- und Nachbearbeitung von TensorFlow.js kompatibel. Sie müssen also nur das Modell selbst migrieren.

So binden Sie LiteRT.js in TensorFlow.js-Pipelines ein:

  1. Konvertieren Sie Ihr ursprüngliches TensorFlow-, JAX- oder PyTorch-Modell in .tflite. Weitere Informationen finden Sie im Abschnitt Modellkonvertierung.
  2. Installieren Sie die NPM-Pakete @litertjs/core und @litertjs/tfjs-interop.
  3. Importieren und verwenden Sie das TensorFlow.js WebGPU-Backend. Dies ist erforderlich, damit LiteRT.js mit TensorFlow.js zusammenarbeiten kann.
  4. Ersetzen Sie TensorFlow.js-Modell laden durch LiteRT.js-Modell laden.
  5. Ersetzen Sie model.predict(inputs) oder model.execute(inputs) von TensorFlow.js durch runWithTfjsTensors(liteRtModel, inputs). runWithTfjsTensors verwendet dieselben Eingabetensoren wie TensorFlow.js-Modelle und gibt TensorFlow.js-Tensoren aus.
  6. Testen Sie, ob die Modellpipeline die erwarteten Ergebnisse ausgibt.

Wenn Sie LiteRT.js mit runWithTfjsTensors verwenden, sind möglicherweise auch die folgenden Änderungen an den Modelleingaben erforderlich:

  1. Eingaben neu anordnen: Je nachdem, wie der Converter die Ein- und Ausgaben des Modells angeordnet hat, müssen Sie ihre Reihenfolge möglicherweise ändern, wenn Sie sie übergeben.
  2. Eingaben transponieren: Möglicherweise hat der Konverter das Layout der Ein- und Ausgaben des Modells im Vergleich zu TensorFlow.js geändert. Möglicherweise müssen Sie Ihre Eingaben transponieren, damit sie dem Modell entsprechen, und Ihre Ausgaben, damit sie dem Rest der Pipeline entsprechen.
  3. Eingaben umbenennen: Wenn Sie benannte Eingaben verwenden, haben sich möglicherweise auch die Namen geändert.

Mit model.getInputDetails() und model.getOutputDetails() können Sie weitere Informationen zu den Ein- und Ausgaben des Modells abrufen.