LiteRT.js を使用したウェブ用 LiteRT

LiteRT.js は、本番環境のウェブ アプリケーションを対象とした Google の高性能 WebAI ランタイムです。LiteRT スタックの継続であり、マルチフレームワークのサポートを保証し、すべてのプラットフォームでコア ランタイムを統合します。

LiteRT.js は、次のコア機能をサポートしています。

  1. LiteRT モデルのブラウザ内サポート: WebAssembly（Wasm）の XNNPack を介して高速化された CPU、および WebGPU API を使用する GPU で、クラス最高のパフォーマンスでモデルを実行します。
  2. マルチ フレームワークの互換性: PyTorch、Jax、TensorFlow など、お好みの ML フレームワークを使用できます。
  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. モデル パイプラインが期待どおりの結果を出力することをテストします。

runWithTfjsTensors で LiteRT.js を使用する場合は、モデル入力に次の変更が必要になることもあります。

  1. 入力の順序を変更する: コンバータがモデルの入力と出力をどのように順序付けたかに応じて、入力と出力を渡すときに順序を変更する必要がある場合があります。
  2. 入力の転置: コンバータが TensorFlow.js で使用されるものと比較して、モデルの入出力のレイアウトを変更した可能性もあります。モデルに合わせて入力を転置し、パイプラインの残りの部分に合わせて出力を転置する必要がある場合があります。
  3. 入力の名前を変更する: 名前付き入力を使用している場合は、名前も変更されている可能性があります。

モデルの入力と出力の詳細については、model.getInputDetails() と model.getOutputDetails() をご覧ください。