Начните работу с LiteRT.js

Это полное руководство по LiteRT.js, описывающее процесс конвертации модели PyTorch для работы в браузере с ускорением WebGPU. В этом примере для модели машинного зрения используется ResNet18 , а для пред- и постобработки — TensorFlow.js .

В руководстве будут рассмотрены следующие этапы:

  1. Конвертируйте модель PyTorch в LiteRT с помощью AI Edge Torch .
    1. Добавьте пакет LiteRT в ваше веб-приложение.
  2. Загрузите модель.
  3. Напишите логику предварительной и постобработки.

Конвертировать в LiteRT

Используйте блокнот PyTorch Converter для преобразования модели PyTorch в соответствующий формат .tflite . Подробное руководство по типам ошибок, с которыми вы можете столкнуться, и способам их устранения см. в файле README AI Edge Torch Converter .

Ваша модель должна быть совместима с torch.export.export , то есть её необходимо экспортировать с помощью TorchDynamo. Следовательно, в ней не должно быть условных ветвлений Python, зависящих от значений времени выполнения в тензорах. Если при выполнении torch.export.export возникают следующие ошибки , ваша модель не может быть экспортирована с помощью torch.export.export . В вашей модели также не должно быть динамических входных или выходных измерений в тензорах, включая измерение пакета.

Вы также можете начать с модели PyTorch, совместимой с TensorRT или экспортируемой в ONNX:

  • Совместимая с TensorRT версия модели может стать хорошей отправной точкой, поскольку некоторые типы преобразований TensorRT также требуют, чтобы модели были экспортируемы в TorchDynamo . Если в модели используются какие-либо операции NVIDIA/CUDA, их необходимо заменить стандартными операциями PyTorch.

  • Модель PyTorch, экспортируемая в ONNX, может стать хорошей отправной точкой, хотя некоторые модели ONNX используют для экспорта TorchScript вместо TorchDynamo, и в этом случае модель может оказаться неэкспортируемой в TorchDynamo (хотя она, вероятно, ближе, чем исходный код модели).

Для получения дополнительной информации см. раздел Преобразование моделей PyTorch в LiteRT .

Добавьте пакет LiteRT

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

npm install @litertjs/core

Импортируйте пакет и загрузите его файлы Wasm:

import {loadLiteRt} from '@litertjs/core';

// They are located in node_modules/@litertjs/core/wasm/
// Serve them statically on your server.
await loadLiteRt(`your/path/to/wasm/`);

Загрузить модель

Импортируйте и инициализируйте LiteRT.js и утилиты конвертации LiteRT-TFJS. Также необходимо импортировать TensorFlow.js для передачи тензоров в LiteRT.js.

import {CompileOptions, loadAndCompile, loadLiteRt, setWebGpuDevice} from '@litertjs/core';
import {runWithTfjsTensors} from '@litertjs/tfjs-interop';

// TensorFlow.js imports
import * as tf from '@tensorflow/tfjs';
import '@tensorflow/tfjs-backend-webgpu'; // Only WebGPU is supported
import {WebGPUBackend} from '@tensorflow/tfjs-backend-webgpu';

async function main() {
  // Initialize TensorFlow.js WebGPU backend
  await tf.setBackend('webgpu');

  // Initialize LiteRT.js's Wasm files
  await loadLiteRt('your/path/to/wasm/');

  // Make LiteRt use the same GPU device as TFJS (for tensor conversion)
  const backend = tf.backend() as WebGPUBackend;
  setWebGpuDevice(backend.device);
  // ...
}

main();

Загрузите преобразованную модель LiteRT:

const model = await loadAndCompile('path_to_model.tflite', {
  accelerator: 'webgpu', // or 'wasm'
});

Напишите модель конвейера

Напишите логику предварительной и постобработки, которая связывает модель с вашим приложением. Рекомендуется использовать TensorFlow.js для предварительной и постобработки, но если TensorFlow.js не написан, вы можете вызвать await tensor.data для получения значения в виде ArrayBuffer или await tensor.array для получения структурированного JS-массива.

Ниже приведен пример сквозного конвейера для ResNet18:

// Wrap in a tf.tidy call to automatically clean up intermediate TensorFlow.js tensors.
// (Note: tidy only supports synchronous functions).
const top5 = tf.tidy(() => {
  // Get RGB data values from an image element and convert it to range [0, 1).
  const image = tf.browser.fromPixels(dogs, 3).div(255);

  // These preprocessing steps come from https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py#L315
  // The mean and standard deviation for the image normalization come from https://github.com/pytorch/vision/blob/main/torchvision/transforms/_presets.py#L38
  const imageData = image.resizeBilinear([224, 224])
    .sub([0.485, 0.456, 0.406])
    .div([0.229, 0.224, 0.225])
    .reshape([1, 224, 224, 3])
    .transpose([0, 3, 1, 2]);

  // You can pass inputs as a single tensor, an array, or a JS Object
  // where keys are the tensor names in the TFLite model.
  // When passing an Object, the output is also an Object.
  // Here, we're passing a single tensor, so the output is an array.
  const probabilities = runWithTfjsTensors(model, imageData)[0];

  // Get the top five classes.
  return tf.topk(probabilities, 5);
});

const values = await top5.values.data();
const indices = await top5.indices.data();
top5.values.dispose(); // Clean up the tfjs tensors.
top5.indices.dispose();

// Print the top five classes.
const classes = ... // Class names are loaded from a JSON file in the demo.
for (let i = 0; i < 5; ++i) {
  const text = `${classes[indices[i]]}: ${values[i]}`;
  console.log(text);
}

Тестирование и устранение неполадок

Информацию о способах тестирования приложения и обработки ошибок см. в следующих разделах.

Тестирование с поддельными входными данными

После загрузки модели рекомендуется сначала протестировать её с помощью фиктивных входных данных. Это позволит выявить любые ошибки выполнения, прежде чем тратить время на написание логики предварительной и последующей обработки для конвейера модели. Для проверки можно использовать тестер моделей LiteRT.js или протестировать вручную.

Тестер моделей LiteRT.js

Тестировщик моделей LiteRT.js запускает вашу модель на GPU и CPU, используя случайные входные данные, чтобы убедиться в корректной работе модели на GPU. Он проверяет следующее:

  • Поддерживаются ли типы входных и выходных данных.
  • Все ли операции доступны на GPU.
  • Насколько близко выходные данные графического процессора соответствуют выходным данным эталонного центрального процессора.
  • Производительность вывода GPU.

Чтобы запустить LiteRT.js Model Tester, выполните npm i @litertjs/model-tester , а затем npx model-tester . Откроется вкладка браузера для запуска модели.

Ручное тестирование модели

Если вы предпочитаете вручную тестировать модель вместо использования тестера моделей LiteRT.js ( @litertjs/model-tester ), вы можете сгенерировать поддельные входные данные и запустить модель с помощью runWithTfjsTensors .

Для генерации фиктивных входных данных необходимо знать имена и формы входных тензоров. Их можно найти в LiteRT.js, вызвав model.getInputDetails или model.getOutputDetails . Простой способ узнать их — установить точку останова после создания модели. Также можно использовать Model Explorer .

Узнав формы и имена входных и выходных данных, вы можете протестировать модель с помощью фиктивных входных данных. Это даст некоторую уверенность в том, что модель будет работать, прежде чем вы напишете оставшуюся часть конвейера машинного обучения. Это позволит проверить, поддерживаются ли все операции модели. Например:

// Imports, initialization, and model loading...
// Create fake inputs for the model
const fakeInputs = model.getInputDetails().map(
    ({shape, dtype}) => tf.ones(shape, dtype));

// Run the model
const outputs = runWithTfjsTensors(model, fakeInputs);
console.log(outputs);

Типы ошибок

Некоторые модели LiteRT могут не поддерживаться LiteRT.js. Ошибки обычно относятся к следующим категориям:

  • Несоответствие формы : известная ошибка, которая влияет только на графический процессор.
  • Операция не поддерживается : среда выполнения не поддерживает операцию в модели. Серверная часть WebGPU имеет более ограниченное покрытие, чем CPU, поэтому, если вы видите эту ошибку на GPU, возможно, вы можете запустить модель на CPU.
  • Неподдерживаемый тип тензора : LiteRT.js поддерживает только тензоры int32 и float32 для входных и выходных данных модели.
  • Слишком большая модель : LiteRT.js ограничен в размере загружаемых им моделей.

Операция не поддерживается

Это означает, что используемый бэкенд не поддерживает одну из операций в модели. Вам потребуется переписать исходную модель PyTorch, чтобы избежать этой операции, и повторно преобразовать её, или же вы сможете запустить модель на центральном процессоре.

В случае BROADCAST_TO эту проблему можно решить, сделав размерность пакета одинаковой для каждого входного тензора модели. Другие случаи могут быть более сложными.

Неподдерживаемый тип тензора

LiteRT.js поддерживает только тензоры int32 и float32 для входных и выходных данных модели.

Модель слишком большая

Обычно это проявляется в виде вызова Aborted() или ошибки выделения памяти во время загрузки модели. LiteRT.js ограничен в размере загружаемых моделей, поэтому, если вы видите это, ваша модель может быть слишком большой. Вы можете попробовать квантовать веса с помощью ai-edge-quantizer , но сохраняйте вычисления в формате float32 или float16, а входные и выходные данные модели — в формате float32 или int32.