Erste Schritte mit LiteRT.js

Dies ist eine End-to-End-Anleitung für LiteRT.js, die den Prozess der Konvertierung eines PyTorch-Modells für die Ausführung im Browser mit WebGPU-Beschleunigung abdeckt. In diesem Beispiel wird ResNet18 für das Vision-Modell und TensorFlow.js für die Vor- und Nachbearbeitung verwendet.

In diesem Leitfaden werden die folgenden Schritte behandelt:

  1. Konvertieren Sie Ihr PyTorch-Modell mit AI Edge Torch in LiteRT.
    1. Fügen Sie Ihrer Web-App das LiteRT-Paket hinzu.
  2. Laden Sie das Modell.
  3. Vor- und Nachbearbeitungslogik schreiben.

In LiteRT konvertieren

Verwenden Sie das PyTorch Converter-Notebook, um ein PyTorch-Modell in das entsprechende .tflite-Format zu konvertieren. Einen ausführlichen Leitfaden zu den Fehlertypen, die auftreten können, und zur Behebung dieser Fehler finden Sie in der README-Datei für den AI Edge Torch Converter.

Ihr Modell muss mit torch.export.export kompatibel sein. Das bedeutet, dass es mit TorchDynamo exportiert werden muss. Daher darf sie keine Python-Bedingungszweige enthalten, die von den Laufzeitwerten in Tensoren abhängen. Wenn Sie die folgenden Fehler während torch.export.export sehen, kann Ihr Modell nicht mit torch.export.export exportiert werden. Die Tensoren Ihres Modells dürfen außerdem keine dynamischen Ein- oder Ausgabedimensionen haben. Dazu gehört auch die Batch-Dimension.

Sie können auch mit einem TensorRT-kompatiblen oder ONNX-exportierbaren PyTorch-Modell beginnen:

  • Eine TensorRT-kompatible Version eines Modells kann ein guter Ausgangspunkt sein, da für einige Arten von TensorRT-Konvertierungen auch erforderlich ist, dass Modelle TorchDynamo-exportierbar sind. Wenn Sie NVIDIA-/CUDA-Vorgänge im Modell verwenden, müssen Sie diese durch standardmäßige PyTorch-Vorgänge ersetzen.

  • Ein ONNX-exportierbares PyTorch-Modell kann ein guter Ausgangspunkt sein. Einige ONNX-Modelle verwenden jedoch TorchScript anstelle von TorchDynamo zum Exportieren. In diesem Fall ist das Modell möglicherweise nicht TorchDynamo-exportierbar (obwohl es wahrscheinlich näher am ursprünglichen Modellcode liegt).

Weitere Informationen finden Sie unter PyTorch-Modelle in LiteRT konvertieren.

LiteRT-Paket hinzufügen

Installieren Sie das Paket @litertjs/core von npm:

npm install @litertjs/core

Importieren Sie das Paket und laden Sie die Wasm-Dateien:

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

Modell laden

Importieren und initialisieren Sie LiteRT.js und die LiteRT-TFJS-Konvertierungstools. Außerdem müssen Sie TensorFlow.js importieren, um Tensoren an LiteRT.js zu übergeben.

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

Laden Sie das konvertierte LiteRT-Modell:

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

Modellpipeline schreiben

Schreiben Sie die Logik für die Vor- und Nachbearbeitung, die das Modell mit Ihrer App verbindet. Die Verwendung von TensorFlow.js für die Vor- und Nachbearbeitung wird empfohlen. Wenn sie nicht in TensorFlow.js geschrieben ist, können Sie await tensor.data aufrufen, um den Wert als ArrayBuffer zu erhalten, oder await tensor.array, um ein strukturiertes JS-Array zu erhalten.

Im Folgenden sehen Sie ein Beispiel für eine End-to-End-Pipeline für 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);
}

Testen und Fehlerbehebung

In den folgenden Abschnitten finden Sie Informationen dazu, wie Sie Ihre Anwendung testen und Fehler beheben können.

Mit gefälschten Eingaben testen

Nach dem Laden eines Modells ist es ratsam, es zuerst mit gefälschten Eingaben zu testen. So können Sie Laufzeitfehler abfangen, bevor Sie Zeit mit dem Schreiben der Vor- und Nachbearbeitungslogik für Ihre Modellpipeline verbringen. Sie können dies mit dem LiteRT.js Model Tester prüfen oder manuell testen.

LiteRT.js-Modelltester

Der LiteRT.js-Modelltester führt Ihr Modell auf der GPU und CPU mit zufälligen Eingaben aus, um zu prüfen, ob das Modell auf der GPU richtig ausgeführt wird. Dabei wird Folgendes überprüft:

  • Ob die Ein- und Ausgabedatentypen unterstützt werden.
  • Gibt an, ob alle Vorgänge auf der GPU verfügbar sind.
  • Wie genau die GPU-Ausgaben mit den CPU-Referenzausgaben übereinstimmen.
  • Die Leistung der GPU-Inferenz.

Führen Sie npm i @litertjs/model-tester und dann npx model-tester aus, um den LiteRT.js-Modelltester auszuführen. Daraufhin wird ein Browsertab geöffnet, in dem Sie Ihr Modell ausführen können.

Manuelle Modelltests

Wenn Sie das Modell lieber manuell testen möchten, anstatt den LiteRT.js-Modelltester (@litertjs/model-tester) zu verwenden, können Sie gefälschte Eingaben generieren und das Modell mit runWithTfjsTensors ausführen.

Um gefälschte Eingaben zu generieren, müssen Sie die Namen und Formen der Eingabetensoren kennen. Sie können mit LiteRT.js durch Aufrufen von model.getInputDetails oder model.getOutputDetails abgerufen werden. Am einfachsten finden Sie sie, indem Sie nach der Erstellung des Modells einen Haltepunkt festlegen. Alternativ können Sie auch den Modell-Explorer verwenden.

Sobald Sie die Eingabe- und Ausgabeformen und -namen kennen, können Sie das Modell mit einer gefälschten Eingabe testen. So können Sie sich darauf verlassen, dass das Modell ausgeführt wird, bevor Sie den Rest der Pipeline für maschinelles Lernen schreiben. So wird getestet, ob alle Modellvorgänge unterstützt werden. Beispiel:

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

Fehlertypen

Einige LiteRT-Modelle werden möglicherweise nicht von LiteRT.js unterstützt. Fehler lassen sich in der Regel in die folgenden Kategorien einteilen:

  • Form stimmt nicht überein: Ein bekannter Fehler, der nur GPUs betrifft.
  • Operation Not Supported (Vorgang nicht unterstützt): Die Laufzeitumgebung unterstützt einen Vorgang im Modell nicht. Das WebGPU-Backend hat eine geringere Abdeckung als die CPU. Wenn Sie diesen Fehler also auf der GPU sehen, können Sie das Modell möglicherweise stattdessen auf der CPU ausführen.
  • Nicht unterstützter Tensortyp: LiteRT.js unterstützt nur int32- und float32-Tensoren für Modelleingaben und -ausgaben.
  • Modell zu groß: LiteRT.js kann nur Modelle bis zu einer bestimmten Größe laden.

Vorgang nicht unterstützt

Das bedeutet, dass das verwendete Backend einen der Vorgänge im Modell nicht unterstützt. Sie müssen das ursprüngliche PyTorch-Modell neu schreiben, um diesen Vorgang zu vermeiden, und es neu konvertieren. Alternativ können Sie das Modell möglicherweise auf der CPU ausführen.

Bei BROADCAST_TO kann das Problem behoben werden, indem die Batchdimension für jeden Eingabetensor des Modells gleich ist. Andere Fälle können komplizierter sein.

Nicht unterstützter Tensortyp

LiteRT.js unterstützt nur int32- und float32-Tensoren für die Ein- und Ausgaben des Modells.

Modell zu groß

Dies äußert sich in der Regel als Aufruf von Aborted() oder als Fehler bei der Speicherzuweisung beim Laden des Modells. LiteRT.js kann nur Modelle bis zu einer bestimmten Größe laden. Wenn Sie diese Meldung sehen, ist Ihr Modell möglicherweise zu groß. Sie können versuchen, die Gewichte mit dem ai-edge-quantizer zu quantisieren, aber die Berechnungen in float32 oder float16 und die Modelleingaben und -ausgaben als float32 oder int32 beibehalten.