适用于 Web 的 LiteRT (LiteRT.js)

LiteRT.js 是 Google 的高性能 WebAI 运行时，面向生产 Web 应用。它是 LiteRT 堆栈的延续，可确保多框架支持，并统一所有平台上的核心运行时。

LiteRT.js 支持以下核心功能：

  1. 支持在浏览器中使用 LiteRT 模型：在 CPU 上以出色的性能运行模型，通过 WebAssembly (Wasm) 上的 XNNPack 和使用 WebGPU API 的 GPU 进行加速。
  2. 多框架兼容性：使用您偏好的机器学习框架：PyTorch、Jax 或 TensorFlow。
  3. 基于现有流水线构建：通过支持 TensorFlow.js 张量作为输入和输出，与现有 TensorFlow.js 流水线集成。

安装

从 npm 安装 @litertjs/core 软件包：

npm install @litertjs/core

Wasm 文件位于 node_modules/@litertjs/core/wasm/ 中。为方便起见，请复制并提供整个 wasm/ 文件夹。然后，导入软件包并加载 Wasm 文件：

import {loadLiteRt} from '@litertjs/core;

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

模型转换

LiteRT.js 使用与 Android 和 iOS 相同的 .tflite 格式，并且支持 Kaggle 和 Huggingface 上的现有模型。如果您有新的 PyTorch 模型，则需要对其进行转换。

将 PyTorch 模型转换为 LiteRT

如需将 PyTorch 模型转换为 LiteRT，请使用 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')

运行转换后的模型

将模型转换为 .tflite 文件后，您可以在浏览器中运行该模型。

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

集成到现有的 TensorFlow.js 流水线中

您应考虑将 LiteRT.js 集成到 TensorFlow.js 流水线中，原因如下：

  1. 出色的 WebGPU 性能：在 LiteRT.js WebGPU 上运行的转换模型针对浏览器性能进行了优化，在基于 Chromium 的浏览器上尤其快速。
  2. 更轻松的模型转换路径：LiteRT.js 转换路径可直接从 PyTorch 转换到 LiteRT。PyTorch 到 TensorFlow.js 的转换路径要复杂得多，需要您从 PyTorch -> ONNX -> TensorFlow -> TensorFlow.js。
  3. 调试工具：LiteRT.js 转化路径附带调试工具。

LiteRT.js 旨在在 TensorFlow.js 流水线中运行，并且与 TensorFlow.js 预处理和后处理兼容，因此您只需迁移模型本身。

按照以下步骤将 LiteRT.js 集成到 TensorFlow.js 流水线中：

  1. 将原始 TensorFlow、JAX 或 PyTorch 模型转换为 .tflite。如需了解详情，请参阅模型转换部分。
  2. 安装 @litertjs/core 和 @litertjs/tfjs-interop NPM 软件包。
  3. 导入并使用 TensorFlow.js WebGPU 后端。 这是 LiteRT.js 与 TensorFlow.js 互操作所必需的。
  4. 将加载 TensorFlow.js 模型替换为加载 LiteRT.js 模型。
  5. 将 TensorFlow.js model.predict(inputs) 或 model.execute(inputs) 替换为 runWithTfjsTensors(liteRtModel, inputs)。runWithTfjsTensors 采用与 TensorFlow.js 模型相同的输入张量，并输出 TensorFlow.js 张量。
  6. 测试模型流水线是否输出了您预期的结果。

将 LiteRT.js 与 runWithTfjsTensors 搭配使用时，可能还需要对模型输入进行以下更改：

  1. 重新排序输入：根据转换器对模型输入和输出的排序方式，您可能需要在传递输入时更改它们的顺序。
  2. 转置输入：与 TensorFlow.js 使用的布局相比，转换器也可能会更改模型的输入和输出布局。您可能需要转置输入以匹配模型，并转置输出以匹配流水线的其余部分。
  3. 重命名输入：如果您使用的是命名输入，名称可能也发生了变化。

您可以使用 model.getInputDetails() 和 model.getOutputDetails() 详细了解模型的输入和输出。