Começar a usar o LiteRT.js

Este é um guia completo do LiteRT.js que aborda o processo de conversão de um modelo do PyTorch para ser executado no navegador com aceleração do WebGPU. Este exemplo usa o ResNet18 para o modelo de visão e o TensorFlow.js para pré e pós-processamento.

O guia inclui as seguintes etapas:

  1. Converta seu modelo PyTorch para LiteRT usando o AI Edge Torch.
    1. Adicione o pacote LiteRT ao seu app da Web.
  2. Carregue o modelo.
  3. Escreva a lógica de pré e pós-processamento.

Converter para LiteRT

Use o notebook do conversor do PyTorch para converter um modelo do PyTorch no formato .tflite adequado. Para um guia detalhado sobre os tipos de erros que podem ocorrer e como corrigi-los, consulte o README do conversor do AI Edge Torch.

Seu modelo precisa ser compatível com torch.export.export, o que significa que ele precisa ser exportável com o TorchDynamo. Portanto, não pode ter ramificações condicionais do Python que dependam dos valores de tempo de execução em tensores. Se você encontrar os seguintes erros durante torch.export.export, não será possível exportar seu modelo com torch.export.export. Além disso, o modelo não pode ter dimensões de entrada ou saída dinâmicas nos tensores. Isso inclui a dimensão de lote.

Você também pode começar com um modelo PyTorch compatível com TensorRT ou exportável para ONNX:

  • Uma versão de um modelo compatível com TensorRT pode ser um bom ponto de partida, já que alguns tipos de conversões do TensorRT também exigem que os modelos sejam exportáveis pelo TorchDynamo. Se você usar operações NVIDIA / CUDA no modelo, será necessário substituí-las por operações padrão do PyTorch.

  • Um modelo PyTorch exportável para ONNX pode ser um bom ponto de partida, embora alguns modelos ONNX usem TorchScript em vez de TorchDynamo para exportar. Nesse caso, o modelo pode não ser exportável para TorchDynamo, mas provavelmente é mais próximo do que o código do modelo original.

Para mais informações, consulte Converter modelos do PyTorch para LiteRT.

Adicionar o pacote LiteRT

Instale o pacote @litertjs/core do npm:

npm install @litertjs/core

Importe o pacote e carregue os arquivos Wasm dele:

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/`);

Carregar o modelo

Importe e inicialize o LiteRT.js e os utilitários de conversão do LiteRT-TFJS. Você também precisa importar o TensorFlow.js para transmitir tensores ao 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();

Carregue o modelo LiteRT convertido:

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

Escrever o pipeline de modelo

Escreva a lógica de pré e pós-processamento que conecta o modelo ao app. Recomendamos usar o TensorFlow.js para pré e pós-processamento, mas se ele não estiver escrito em TensorFlow.js, chame await tensor.data para receber o valor como um ArrayBuffer ou await tensor.array para receber uma matriz JS estruturada.

Confira um exemplo de pipeline completo para 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);
}

Testar e resolver problemas.

Consulte as seções a seguir sobre como testar seu aplicativo e processar erros.

Testar com entradas falsas

Depois de carregar um modelo, é recomendável testá-lo primeiro com entradas falsas. Isso vai detectar erros de execução antes de você gastar tempo escrevendo a lógica de pré e pós-processamento para seu pipeline de modelo. Para verificar isso, use o testador de modelo LiteRT.js ou faça o teste manualmente.

Testador de modelos LiteRT.js

O testador de modelos LiteRT.js executa seu modelo na GPU e na CPU usando entradas aleatórias para verificar se o modelo é executado corretamente na GPU. Ele verifica o seguinte:

  • Se os tipos de dados de entrada e saída são compatíveis.
  • Se todas as operações estão disponíveis na GPU.
  • O nível de correspondência entre as saídas da GPU e da CPU de referência.
  • O desempenho da inferência de GPU.

Para executar o testador de modelos LiteRT.js, execute npm i @litertjs/model-tester e depois npx model-tester. Uma guia do navegador será aberta para você executar o modelo.

Teste manual de modelos

Se você preferir testar o modelo manualmente em vez de usar o testador de modelo LiteRT.js (@litertjs/model-tester), gere entradas falsas e execute o modelo com runWithTfjsTensors.

Para gerar entradas falsas, é necessário saber os nomes e as formas dos tensores de entrada. Eles podem ser encontrados com o LiteRT.js chamando model.getInputDetails ou model.getOutputDetails. Uma maneira simples de encontrá-los é definir um ponto de interrupção depois que o modelo é criado. Se preferir, use o Simulador de modelos.

Depois de saber as formas e os nomes de entrada e saída, é possível testar o modelo com uma entrada falsa. Isso dá alguma confiança de que o modelo será executado antes de você escrever o restante do pipeline de machine learning. Isso testaria se todas as operações de modelo são compatíveis. Exemplo:

// 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);

Tipos de erro

Alguns modelos do LiteRT podem não ser compatíveis com o LiteRT.js. Os erros geralmente se enquadram nestas categorias:

  • Incompatibilidade de forma: um bug conhecido que afeta apenas a GPU.
  • Operação sem suporte: o ambiente de execução não é compatível com uma operação no modelo. O back-end WebGPU tem uma cobertura mais limitada do que a CPU. Portanto, se você estiver recebendo esse erro na GPU, talvez seja possível executar o modelo na CPU.
  • Tipo de tensor não compatível: o LiteRT.js só aceita tensores int32 e float32 para entradas e saídas de modelo.
  • Modelo muito grande: o LiteRT.js tem um limite de tamanho para os modelos que pode carregar.

Operação não compatível

Isso indica que o back-end usado não é compatível com uma das operações no modelo. Você precisará reescrever o modelo original do PyTorch para evitar essa operação e reconvertê-lo ou poderá executar o modelo na CPU.

No caso de BROADCAST_TO, isso pode ser resolvido fazendo com que a dimensão do lote seja a mesma para todos os tensores de entrada do modelo. Outros casos podem ser mais complicados.

Tipo de tensor sem suporte

O LiteRT.js só é compatível com tensores int32 e float32 para as entradas e saídas do modelo.

Modelo muito grande

Isso geralmente aparece como uma chamada para Aborted() ou uma falha na alocação de memória durante o carregamento do modelo. O LiteRT.js tem um limite no tamanho dos modelos que pode carregar. Se você estiver vendo isso, talvez seu modelo seja muito grande. Você pode tentar quantizar os pesos com o ai-edge-quantizer, mas mantenha os cálculos em float32 ou float16 e as entradas e saídas do modelo como float32 ou int32.