使用 LiteRT.js 的網頁版 LiteRT

LiteRT.js 是 Google 的高效能 WebAI 執行階段,適用於正式版 Web 應用程式。這是 LiteRT 堆疊的延續,可確保多框架支援,並統一所有平台的核心執行階段。

LiteRT.js 支援下列核心功能:

  1. 瀏覽器內建硬體加速推論:執行模型時,可透過對應至輕量型 WebAssembly (Wasm) 的 XNNPack 加速,享有卓越的 CPU 效能。對於 GPU 和專用硬體縮放 (例如 NPU),LiteRT.js 會以原生方式顯示 WebGPU API 和新興的 WebNN API,進而實現精細的平台專屬最佳化。
  2. 多架構相容性:從偏好的 ML 架構 (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';

// Load the LiteRT.js Wasm files from a CDN.
await loadLiteRt('https://cdn.jsdelivr.net/npm/@litertjs/core/wasm/')
// Alternatively, host them from your server.
// They are located in node_modules/@litertjs/core/wasm/
await loadLiteRt(`your/path/to/wasm/`);

模型轉換

LiteRT.js 使用的 .tflite 格式與 LiteRT 生態系統的其餘部分相同,並支援 KaggleHuggingface 上的現有模型。如果您有新的 PyTorch 模型,就必須轉換該模型。

將 PyTorch 模型轉換為 LiteRT

如要將 PyTorch 模型轉換為 LiteRT,請使用 litert-torch 轉換工具

import litert_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 = litert_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',
    // Can select from 'webnn', 'webgpu', & 'wasm'.
    // Additionally, you can pass an array of accelerators e.g. ['webnn', 'wasm']
    // if you would like to fallback to CPU execution,
    // Note that ONLY cpu fallback is supported for now
    // (i.e. specifying ['webnn', 'webgpu']) will lead to compilation errors
});
// 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 = new Tensor(image, /* shape */ [1, 3, 224, 224]);

// Run the model
const outputs = await model.run(inputTensor);
// You can also use `await model.run([inputTensor]);`
// or `await model.run({'input_tensor_name': inputTensor});`

// Clean up and get outputs
inputTensor.delete();
const output = outputs[0];
const outputData = await output.data();
output.delete();

整合至現有的 TensorFlow.js 管道

建議您將 LiteRT.js 整合至 TensorFlow.js 管道,原因如下:

  1. 優異的 GPU 和硬體效能:LiteRT.js 模型會運用 WebGPU 加速功能,在各種瀏覽器架構中發揮最佳效能。LiteRT.js 支援 WebGPU 和即將推出的 WebNN,可為各種邊緣裝置提供彈性的硬體加速功能。
  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. 測試模型管道是否輸出預期結果。

搭配 runWithTfjsTensors 使用 LiteRT.js 時,可能也需要對模型輸入內容進行下列變更:

  1. 重新排序輸入內容:視轉換工具排序模型輸入和輸出的方式而定,您可能需要變更輸入內容的順序。
  2. 轉置輸入內容:與 TensorFlow.js 使用的內容相比,轉換器也可能變更模型的輸入和輸出配置。您可能需要轉置輸入內容,以符合模型,並轉置輸出內容,以符合其餘的管道。
  3. 重新命名輸入內容:如果您使用具名輸入內容,名稱可能也會變更。

如要進一步瞭解模型的輸入和輸出內容,請使用 model.getInputDetails()model.getOutputDetails()