LiteRT для веба с LiteRT.js

LiteRT.js — это высокопроизводительная среда выполнения WebAI от Google, ориентированная на разработку веб-приложений. Она является продолжением стека LiteRT, обеспечивая поддержку различных фреймворков и унификацию нашей базовой среды выполнения на всех платформах.

LiteRT.js поддерживает следующие основные функции:

  1. Поддержка моделей LiteRT в браузере : запуск моделей с лучшей в своем классе производительностью на ЦП, ускоренных с помощью XNNPack на WebAssembly (Wasm) и GPU с использованием API WebGPU .
  2. Совместимость с несколькими фреймворками : используйте предпочитаемую вами инфраструктуру машинного обучения: PyTorch, Jax или TensorFlow.
  3. Разработка на основе существующих конвейеров : интеграция с существующими конвейерами TensorFlow.js путем поддержки тензоров TensorFlow.js в качестве входов и выходов.

Установка

Установите пакет @litertjs/core из npm:

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 использует тот же формат .tflite , что и Android и iOS, и поддерживает существующие модели на 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. Установите пакеты NPM @litertjs/core и @litertjs/tfjs-interop .
  3. Импортируйте и используйте бэкенд TensorFlow.js WebGPU . Это необходимо для взаимодействия LiteRT.js с TensorFlow.js.
  4. Замените загрузку модели TensorFlow.js загрузкой модели LiteRT.js .
  5. Замените model.predict (входные данные) или model.execute (входные данные) TensorFlow.js на runWithTfjsTensors (liteRtModel, входные данные). runWithTfjsTensors принимает те же входные тензоры, что и модели TensorFlow.js, и выводит тензоры TensorFlow.js.
  6. Проверьте, выдает ли конвейер модели ожидаемые вами результаты.

Использование LiteRT.js с runWithTfjsTensors также может потребовать внесения следующих изменений во входные данные модели:

  1. Изменение порядка входных данных : в зависимости от того, как преобразователь упорядочил входные и выходные данные модели, вам может потребоваться изменить их порядок при передаче.
  2. Транспонирование входных данных : также возможно, что конвертер изменил расположение входных и выходных данных модели по сравнению с тем, что использует TensorFlow.js. Возможно, вам потребуется транспонировать входные данные для соответствия модели, а выходные — для соответствия остальной части конвейера.
  3. Переименование входов : если вы используете именованные входы, их имена также могли измениться.

Дополнительную информацию о входах и выходах модели можно получить с помощью model.getInputDetails() и model.getOutputDetails() .