LiteRT para Web com LiteRT.js

O LiteRT.js é o ambiente de execução WebAI de alto desempenho do Google, voltado para aplicativos da Web de produção. É uma continuação da pilha LiteRT, garantindo suporte multiframework e unificando nosso tempo de execução principal em todas as plataformas.

O LiteRT.js é compatível com os seguintes recursos principais:

  1. Suporte no navegador para modelos LiteRT: execute modelos com a melhor performance da categoria na CPU, acelerada pelo XNNPack no WebAssembly (Wasm) e na GPU usando a API WebGPU.
  2. Compatibilidade com vários frameworks: use o framework de ML de sua preferência: PyTorch, Jax ou TensorFlow.
  3. Crie pipelines atuais: integre-se aos pipelines do TensorFlow.js existentes oferecendo suporte a tensores do TensorFlow.js como entradas e saídas.

Instalação

Instale o pacote @litertjs/core do npm:

npm install @litertjs/core

Os arquivos Wasm estão localizados em node_modules/@litertjs/core/wasm/. Para facilitar, copie e disponibilize toda a pasta wasm/. Em seguida, importe o pacote e carregue os arquivos 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/`);

Conversão de modelos

O LiteRT.js usa o mesmo formato .tflite do Android e do iOS e é compatível com modelos atuais no Kaggle e no Huggingface. Se você tiver um novo modelo do PyTorch, será necessário convertê-lo.

Converter um modelo do PyTorch para LiteRT

Para converter um modelo do PyTorch em LiteRT, use o conversor 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')

Executar o modelo convertido

Depois de converter o modelo em um arquivo .tflite, é possível executá-lo no navegador.

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

Integrar aos pipelines do TensorFlow.js

Considere integrar o LiteRT.js aos seus pipelines do TensorFlow.js pelos seguintes motivos:

  1. Melhor performance da WebGPU da categoria: modelos convertidos executados no LiteRT.js WebGPU são otimizados para a performance do navegador e são especialmente rápidos em navegadores baseados no Chromium.
  2. Caminho de conversão de modelo mais fácil: o caminho de conversão do LiteRT.js vai diretamente do PyTorch para o LiteRT. O caminho de conversão do PyTorch para o TensorFlow.js é muito mais complicado, exigindo que você vá de PyTorch -> ONNX -> TensorFlow -> TensorFlow.js.
  3. Ferramentas de depuração: o caminho de conversão LiteRT.js vem com ferramentas de depuração.

O LiteRT.js foi projetado para funcionar em pipelines do TensorFlow.js e é compatível com o pré e pós-processamento do TensorFlow.js. Portanto, a única coisa que você precisa migrar é o modelo em si.

Integre o LiteRT.js aos pipelines do TensorFlow.js seguindo estas etapas:

  1. Converta seu modelo original do TensorFlow, JAX ou PyTorch para .tflite. Para detalhes, consulte a seção conversão de modelo.
  2. Instale os pacotes @litertjs/core e @litertjs/tfjs-interop do NPM.
  3. Importe e use o backend WebGPU do TensorFlow.js. Isso é necessário para que o LiteRT.js funcione com o TensorFlow.js.
  4. Substitua carregar o modelo do TensorFlow.js por carregar o modelo do LiteRT.js.
  5. Substitua o TensorFlow.js model.predict(inputs) ou model.execute(inputs) por runWithTfjsTensors(liteRtModel, inputs). O runWithTfjsTensors usa os mesmos tensores de entrada que os modelos do TensorFlow.js e gera tensores do TensorFlow.js.
  6. Teste se o pipeline do modelo gera os resultados esperados.

Usar o LiteRT.js com runWithTfjsTensors também pode exigir as seguintes mudanças nas entradas do modelo:

  1. Reordene as entradas: dependendo de como o conversor ordenou as entradas e saídas do modelo, talvez seja necessário mudar a ordem ao transmiti-las.
  2. Transpor entradas: também é possível que o conversor tenha mudado o layout das entradas e saídas do modelo em comparação com o que o TensorFlow.js usa. Talvez seja necessário transpor as entradas para corresponder ao modelo e as saídas para corresponder ao restante do pipeline.
  3. Renomear entradas: se você estiver usando entradas nomeadas, os nomes também podem ter mudado.

Você pode receber mais informações sobre as entradas e saídas do modelo com model.getInputDetails() e model.getOutputDetails().