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 em CPU, acelerados via XNNPack em WebAssembly (Wasm) e 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. Criar pipelines atuais: integre-se aos pipelines do TensorFlow.js existentes com 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;

// Host LiteRT's Wasm files on your server.
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 =
    await new Tensor(image, /* shape */ [1, 3, 224, 224]).moveTo('webgpu');

// Run the model
const outputs = model(inputTensor);
// You can also use model([inputTensor])
// or model({'input_tensor_name': inputTensor})

// Clean up and get outputs
inputTensor.delete();
const outputTensorCpu = await outputs[0].moveTo('wasm');
const outputData = outputTensorCpu.toTypedArray();
outputTensorCpu.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 NPM @litertjs/core e @litertjs/tfjs-interop.
  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().