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

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

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

  1. Поддержка моделей LiteRT в браузере : запуск моделей с лучшей в своем классе производительностью на ЦП, ускоренной XNNPack на WebAssembly (Wasm), и на графическом процессоре с использованием 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';

// 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 , что и 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 = 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. Лучшая в своем классе производительность 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 (inputs)` или model.execute (inputs)` в TensorFlow.js на ` runWithTfjsTensors (liteRtModel, inputs)`. runWithTfjsTensors принимает те же входные тензоры, что и модели TensorFlow.js, и выдает на выходе тензоры TensorFlow.js.
  6. Проверьте, что конвейер обработки модели выдает ожидаемые результаты.

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

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

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