Três obstáculos ao migrar do TensorFlow.js para o LiteRT.js
Executar modelos de IA no navegador é empolgante, mas logo nos deparamos com limitações. Quando processamos imagens de alta resolução ou iniciamos operações pesadas, a interface começa a travar. O TensorFlow.js (TF.js) utiliza o backend WebGL e bindings de kernel JavaScript, o que gera um overhead significativo em operações de matrizes em larga escala.
Por outro lado, o LiteRT.js compila o runtime nativo em C++ para WebAssembly (Wasm) e o porta para o navegador. É uma estrutura que resolve o gargalo de forma fundamental. Além disso, o processo de trazer modelos PyTorch para a web torna-se mais simples. Antigamente, era preciso converter de PyTorch para ONNX e depois passar pelo TensorFlow para obter o formato TF.js, o que causava quebra na compatibilidade de operadores e perda de precisão. Agora, basta extrair um arquivo .tflite padrão com uma única biblioteca (ai-edge-torch).
Se você não quiser descartar seus recursos existentes, pode usar o pacote @litertjs/tfjs-interop fornecido pelo Google. A abordagem consiste em manter o pipeline de pré e pós-processamento de dados do TF.js e substituir apenas a execução principal de previsão do modelo pelo LiteRT.js.
Resolvendo a inconsistência entre NCHW e NHWC
Os modelos TFLite convertidos do PyTorch geralmente exigem a estrutura de canal prioritário NCHW (canal, altura, largura). No entanto, o array ImageData do Canvas no navegador está no formato NHWC (altura, largura, canal), onde os pixels são listados sequencialmente. É necessário um código de pré-processamento que processe essa lacuna de forma síncrona sem sobrecarregar a thread principal.
Primeiro, crie um Float32Array multiplicando o número total de pixels do ImageData por 3. Em seguida, especifique os offsets de achatamento de cada canal (vermelho, verde, azul) como 0, o total de pixels e o dobro do total de pixels, respectivamente. Por fim, normalize os valores de pixel entre 0 e 255 dividindo-os por 255.0 e atribua-os aos offsets de cada canal. Esse é o processo de reorganizar os dados de entrada em um buffer de tensor achatado NCHW.
`javascript
/**
- Utilitário de pré-processamento para converter rapidamente buffers de ImageData em formato NHWC para NCHW de alta performance
- @param {ImageData} imageData - Dados de pixel originais obtidos do HTML5 Canvas
- @param {number} width - Resolução horizontal da imagem de entrada exigida pelo modelo alvo
- @param {number} height - Resolução vertical da imagem de entrada exigida pelo modelo alvo
- @returns {Float32Array} Buffer de tensor achatado reorganizado no formato NCHW
*/
export function preprocessNHWCToNCHW(imageData, width, height) {
const { data } = imageData;
const totalPixels = width * height;
const nchwBuffer = new Float32Array(totalPixels * 3);
const rChannelOffset = 0;
const gChannelOffset = totalPixels;
const bChannelOffset = totalPixels * 2;
for (let i = 0; i < totalPixels; i++) {
const srcIndex = i * 4;
nchwBuffer[rChannelOffset + i] = data[srcIndex] / 255.0;
nchwBuffer[gChannelOffset + i] = data[srcIndex + 1] / 255.0;
nchwBuffer[bChannelOffset + i] = data[srcIndex + 2] / 255.0;
}
return nchwBuffer;
}
`
Evitando travamentos na tela com Web Workers e Zero-Copy
A situação mais terrível ao executar deep learning no frontend é o travamento da UI. Para que o navegador exiba animações fluidas a 60 quadros por segundo, o loop de eventos deve processar tarefas síncronas em menos de 16,6ms. No entanto, operações de tensores bloqueiam facilmente o loop de renderização principal (que é single-threaded).
O LiteRT.js tem uma velocidade de execução básica cerca de 3 vezes mais rápida do que as ferramentas baseadas em JavaScript. Ao adicionar hardware de aceleração como WebGPU ou WebNN, a velocidade aumenta de 5 a até 60 vezes em relação ao modo CPU. O backend WebNN, que utiliza NPU dedicada, exige que o recurso de JavaScript Promise Integration (JSPI) esteja ativado para conectar o escalonador de kernel WebAssembly síncrono ao loop de controle de hardware assíncrono do navegador. Para utilizar esses recursos de aceleração e manter a thread principal responsiva, a inicialização da biblioteca e todo o pipeline de inferência devem ser isolados dentro de um Web Worker.
Nesse momento, se você simplesmente passar o buffer na comunicação de dados entre threads, ocorrerá uma cópia interna de memória, acumulando overhead na CPU e na memória heap. É necessário usar Transferable Objects para transferir a própria propriedade da área de endereço de memória física, eliminando assim a latência. O buffer cuja propriedade foi transferida é imediatamente invalidado no contexto de origem, garantindo também a segurança entre threads.
`javascript
// litert-worker.js - Módulo de Web Worker dedicado a operações de inferência em segundo plano
import { loadLiteRt, loadAndCompile, Tensor } from '@litertjs/core';
let compiledModel = null;
let isLoaded = false;
self.onmessage = async (event) => {
const { type, payload } = event.data;
switch (type) {
case 'LOAD_MODEL':
try {
await loadLiteRt(payload.wasmDirectory, { jspi: payload.enableJspi || false });
compiledModel = await loadAndCompile(payload.modelUrl, {
accelerator: payload.accelerator || 'webgpu'
});
isLoaded = true;
self.postMessage({ type: 'MODEL_READY' });
} catch (err) {
self.postMessage({ type: 'ERROR', error: Initialization failed: ${err.message} });
}
break;
case 'RUN_INFERENCE':
if (!isLoaded || !compiledModel) {
self.postMessage({ type: 'ERROR', error: 'Model has not been loaded' });
return;
}
try {
const rawInputData = payload.bufferData;
const inputShape = payload.shape;
const inputTensor = new Tensor(rawInputData, inputShape);
const results = await compiledModel.run(inputTensor);
const cpuOutputTensor = await results[0].moveTo('wasm');
const outputBuffer = cpuOutputTensor.toTypedArray();
inputTensor.delete();
cpuOutputTensor.delete();
results[0].delete();
self.postMessage(
{
type: 'INFERENCE_COMPLETE',
payload: {
data: outputBuffer,
shape: results[0].shape
}
},
[outputBuffer.buffer]
);
} catch (err) {
self.postMessage({ type: 'ERROR', error: `Inference failed: ${err.message}` });
}
break;
default:
self.postMessage({ type: 'UNKNOWN_OP' });
}
};
`
`javascript
// litert-bridge.js - Classe orquestradora de IA para a thread principal
export class LiteRtBridge {
constructor(workerPath) {
this.worker = new Worker(workerPath);
this.promiseMap = new Map();
this.tokenCounter = 0;
this.worker.onmessage = (event) => {
const { type, payload, error } = event.data;
if (type === 'MODEL_READY') {
if (this.initResolve) this.initResolve();
} else if (type === 'INFERENCE_COMPLETE') {
const currentToken = this.tokenCounter;
const promiseHandler = this.promiseMap.get(currentToken);
if (promiseHandler) {
promiseHandler.resolve(payload);
this.promiseMap.delete(currentToken);
}
} else if (type === 'ERROR') {
const currentToken = this.tokenCounter;
const promiseHandler = this.promiseMap.get(currentToken);
if (promiseHandler) {
promiseHandler.reject(new Error(error));
this.promiseMap.delete(currentToken);
} else if (this.initReject) {
this.initReject(new Error(error));
}
}
};
}
bootstrap(wasmDirectory, modelUrl, accelerator = 'webgpu') {
return new Promise((resolve, reject) => {
this.initResolve = resolve;
this.initReject = reject;
this.worker.postMessage({
type: 'LOAD_MODEL',
payload: { wasmDirectory, modelUrl, accelerator, enableJspi: true }
});
});
}
execute(inputFloat32Array, inputShape) {
return new Promise((resolve, reject) => {
this.tokenCounter++;
this.promiseMap.set(this.tokenCounter, { resolve, reject });
this.worker.postMessage(
{
type: 'RUN_INFERENCE',
payload: {
bufferData: inputFloat32Array,
shape: inputShape
}
},
[inputFloat32Array.buffer]
);
});
}
}
`
Não confie no Garbage Collection automático
Ao usar o TF.js, o padrão era eliminar tensores dentro de um escopo de chamada síncrona usando tf.tidy(). No entanto, quando códigos assíncronos ou Promises estavam envolvidos, isso criava bugs onde o escopo era deixado antes do término do trabalho assíncrono, resultando em tensores corrompidos ou coleta não realizada.
O LiteRT.js é ainda mais rígido. Ele não é alvo do Garbage Collection (GC) do motor do navegador. O espaço de memória virtual linear do WebAssembly e os buffers WebGPU não podem ser rastreados pelo motor JavaScript, como o V8, em relação ao estado da heap. Se você não chamar explicitamente .delete() na instância do tensor após o uso, a memória do navegador se acumulará infinitamente. Em um serviço que transmite quadros de vídeo de alta definição dezenas de vezes por segundo, a aba travará em poucos minutos.
É aconselhável gerenciar isso manualmente, criando uma classe rastreadora de escopo (scope tracker) que registre o tempo de vida dos tensores gerados em todo o pipeline assíncrono e garanta sua destruição coletiva.
`javascript
/**
- Gerenciador de escopo de memória assíncrona que medeia o rastreamento manual e a destruição garantida de tensores da heap WebAssembly
*/
export class LiteRtScopeTracker {
constructor() {
this.trackList = new Set();
}
/**
- Registra tensores criados ou movidos na lista de gerenciamento de ciclo de vida
- @param {Tensor} tensor - Tensor do LiteRT.js a ser rastreado e destruído
- @returns {Tensor} Retorna o objeto tensor para suportar a escrita de código inline
*/
register(tensor) {
if (tensor && typeof tensor.delete === 'function') {
this.trackList.add(tensor);
}
return tensor;
}
/**
- Força a estrutura de gerenciamento seguro de pipeline de tensores dentro de um bloco de execução assíncrono
- @param {Function} asyncCallable - Função de lógica de negócio de inferência assíncrona
- @returns {Promise<*>} Resultado final dos dados brutos retornados pelo bloco de execução
*/
async enforceScope(asyncCallable) {
try {
const outputResult = await asyncCallable(this);
if (Array.isArray(outputResult)) {
outputResult.forEach((item) => this.trackList.delete(item));
} else {
this.trackList.delete(outputResult);
}
return outputResult;
} finally {
this.disposeAll();
}
}
/**
- Libera permanentemente da área Wasm todos os tensores TFLite nativos vinculados ao escopo de gerenciamento
*/
disposeAll() {
for (const tensor of this.trackList) {
try {
tensor.delete();
} catch (err) {
console.error('An error occurred while cleaning the native Wasm tensor memory:', err);
}
}
this.trackList.clear();
}
}
`
`javascript
// Exemplo de implementação de processamento de inferência de IA multithreaded assíncrono, robusto e seguro, utilizando o gerenciador de escopo de memória
export async function runRobustVisionInference(rawPixelArray, compiledModel) {
const scopeTracker = new LiteRtScopeTracker();
try {
return await scopeTracker.enforceScope(async (scope) => {
const inputTensor = scope.register(new Tensor(rawPixelArray, [1, 3, 224, 224]));
const predictionResults = await compiledModel.run(inputTensor);
predictionResults.forEach((tensor) => scope.register(tensor));
const firstOutputTensor = predictionResults[0];
const wasmTransferTensor = scope.register(await firstOutputTensor.moveTo('wasm'));
const targetJsArray = wasmTransferTensor.toTypedArray();
return targetJsArray;
});
} catch (err) {
console.error('Fatal crash occurred during the model pipeline execution:', err);
throw err;
}
}
`
Carregamento condicional de arquivos Wasm pesados
O problema do aumento do tamanho do download dos recursos também não pode ser ignorado. Para permitir que o tree-shaking ocorra no build bundler, é necessário eliminar as referências estáticas no estilo CommonJS e escrever o código-fonte com base na sintaxe de módulos ES6 (import/export). Também é importante configurar sideEffects: false nas ferramentas de build para deixar o bundle mais leve.
O runtime principal do LiteRT.js, @litertjs/core, carrega de forma seletiva três builds de kernel WebAssembly, dependendo do desempenho do dispositivo. Em navegadores modernos como Chrome ou Edge, ele seleciona o módulo que suporta multithreading e SIMD (litert_wasm_simd.wasm), enquanto em ambientes legados como o Safari, ele carrega o módulo de fallback padrão (litert_wasm.wasm). Se a compilação via GPU falhar, o runtime XNNPACK, que transfere todos os operadores de hardware para o sandbox de CPU Wasm, entra em ação como dispositivo auxiliar.
Para evitar atrasos no carregamento inicial, deve-se verificar os aceleradores por especificação e carregar os módulos dinamicamente.
`javascript
// litert-loader.js - Engine de detecção de dispositivo em tempo real e acoplamento dinâmico de aceleradores
export async function bootstrapHighPerformanceInferenceEngine() {
const supportsWebGpu = 'gpu' in navigator;
let chosenAccelerator = 'wasm';
if (supportsWebGpu) {
try {
const gpuAdapter = await navigator.gpu.requestAdapter();
if (gpuAdapter) {
const info = await gpuAdapter.requestDevice();
if (info) {
chosenAccelerator = 'webgpu';
}
}
} catch (e) {
console.warn("GPU profile probe failed, resolving execution chain to fallback WASM.");
}
}
const { loadLiteRt, loadAndCompile } = await import('@litertjs/core');
const cdnWasmHostPath = 'https://cdn.jsdelivr.net/npm/@litertjs/core/wasm/';
await loadLiteRt(cdnWasmHostPath, {
jspi: chosenAccelerator === 'webnn'
});
return {
loadAndCompile,
chosenAccelerator
};
}
`
Coloque a transferência de propriedade de memória via Web Workers e a destruição explícita de objetos na área Wasm como a base do seu projeto. Assim que você começar a controlar o fluxo de dados diretamente, poderá implantar serviços de IA on-device em produção sem medo de que o navegador trave.