Начните работу с 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';

// 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 и утилиты преобразования LiteRT-TFJS. Вам также необходимо импортировать TensorFlow.js для передачи тензоров в LiteRT.js.

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

// TensorFlow.js imports
import * as tf from '@tensorflow/tfjs';
import '@tensorflow/tfjs-backend-webgpu';
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 TFJS use the same GPU device as LiteRT.js (for tensor conversion)
  const device = await getWebGpuDevice();
  tf.removeBackend('webgpu');
  tf.registerBackend('webgpu', () => new WebGPUBackend(device, device.adapterInfo));
  await tf.setBackend('webgpu');
  // ...
}

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 , чтобы получить структурированный массив JavaScript.

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

// Wrap in a tf.tidy call to automatically clean up intermediate TensorFlow.js tensors.
// (Note: tidy only supports synchronous functions).
const imageData = 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
  return 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]);
});

// Run the model
const outputs = await runWithTfjsTensors(model, [imageData]);
const probabilities = outputs[0];

// Get the top five classes.
const top5 = tf.topk(probabilities, 5);

const values = await top5.values.data();
const indices = await top5.indices.data();

// Clean up TFJS tensors
tf.dispose(outputs);
tf.dispose(top5);
tf.dispose(imageData);

// 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 Model Tester или протестировать модель вручную.

LiteRT.js Model Tester

LiteRT.js Model Tester запускает вашу модель на GPU и CPU, используя случайные входные данные, чтобы убедиться в корректной работе модели на 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 = await 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.