Gerando pacotes de navegador com tamanho otimizado com TensorFlow.js

Visão geral

O TensorFlow.js 3.0 oferece suporte para a criação de pacotes de navegador orientados à produção e com tamanho otimizado . Em outras palavras, queremos tornar mais fácil para você enviar menos JavaScript para o navegador.

Esse recurso é voltado para usuários com casos de uso de produção que se beneficiariam particularmente com a redução de bytes de sua carga útil (e, portanto, estão dispostos a se esforçar para conseguir isso). Para usar esse recurso, você deve estar familiarizado com módulos ES , ferramentas de empacotamento de JavaScript, como webpack ou rollup , e conceitos como agitação de árvore/eliminação de código morto .

Este tutorial demonstra como criar um módulo tensorflow.js customizado que pode ser usado com um bundler para gerar uma compilação de tamanho otimizado para um programa usando tensorflow.js.

Terminologia

No contexto deste documento, existem alguns termos-chave que usaremos:

Módulos ES - O sistema de módulos JavaScript padrão . Introduzido em ES6/ES2015. Identificável pelo uso de declarações de importação e exportação .

Agrupamento - pegar um conjunto de ativos JavaScript e agrupá-los/agregá-los em um ou mais ativos JavaScript que podem ser usados ​​em um navegador. Esta é a etapa que geralmente produz os ativos finais que são veiculados ao navegador. Os aplicativos geralmente farão seu próprio empacotamento diretamente das fontes da biblioteca transpilada . Os empacotadores comuns incluem rollup e webpack . O resultado final do agrupamento é conhecido como pacote (ou às vezes como um pedaço, se for dividido em várias partes)

Tree-Shaking / Dead Code Elimination - Remoção de código que não é usado pelo aplicativo escrito final. Isso é feito durante o agrupamento, normalmente na etapa de minificação.

Operações (Ops) - Uma operação matemática em um ou mais tensores que produz um ou mais tensores como saída. As operações são códigos de “alto nível” e podem usar outras operações para definir sua lógica.

Kernel – Uma implementação específica de uma operação vinculada a recursos de hardware específicos. Os kernels são de 'baixo nível' e específicos de backend. Algumas operações têm um mapeamento individual de operação para kernel, enquanto outras operações usam vários kernels.

Escopo e casos de uso

Modelos gráficos apenas de inferência

O principal caso de uso que ouvimos dos usuários relacionado a isso e que oferecemos suporte nesta versão é fazer inferência com modelos de gráfico do TensorFlow.js . Se você estiver usando um modelo de camadas TensorFlow.js , poderá convertê-lo para o formato de modelo gráfico usando o tfjs-converter . O formato do modelo gráfico é mais eficiente para o caso de uso de inferência.

Manipulação de tensor de baixo nível com tfjs-core

O outro caso de uso que oferecemos suporte são programas que usam diretamente o pacote @tensorflow/tjfs-core para manipulação de tensores de nível inferior.

Nossa abordagem para construções personalizadas

Nossos princípios básicos ao projetar esta funcionalidade incluem o seguinte:

  • Aproveite ao máximo o sistema de módulos JavaScript (ESM) e permita que os usuários do TensorFlow.js façam o mesmo.
  • Torne o TensorFlow.js o mais abalável possível pelos empacotadores existentes (por exemplo, webpack, rollup, etc.). Isso permite que os usuários aproveitem todos os recursos desses empacotadores, incluindo recursos como divisão de código.
  • Tanto quanto possível, mantenha a facilidade de uso para usuários que não são tão sensíveis ao tamanho do pacote . Isso significa que as compilações de produção exigirão mais esforço, já que muitos dos padrões em nossas bibliotecas oferecem suporte à facilidade de uso em vez de compilações com tamanho otimizado.

O objetivo principal do nosso fluxo de trabalho é produzir um módulo JavaScript personalizado para TensorFlow.js que contenha apenas a funcionalidade necessária para o programa que estamos tentando otimizar. Contamos com empacotadores existentes para fazer a otimização real.

Embora dependamos principalmente do sistema de módulos JavaScript, também fornecemos uma ferramenta CLI personalizada para lidar com partes que não são fáceis de especificar por meio do sistema de módulos no código voltado para o usuário. Dois exemplos disso são:

  • Especificações do modelo armazenadas em arquivos model.json
  • A operação do sistema de despacho de kernel específico de back-end que usamos.

Isso torna a geração de uma construção tfjs personalizada um pouco mais complicada do que apenas apontar um empacotador para o pacote @tensorflow/tfjs normal.

Como criar pacotes personalizados com tamanho otimizado

Etapa 1: Determine quais kernels seu programa está usando

Esta etapa nos permite determinar todos os kernels usados ​​por quaisquer modelos que você executa ou código de pré/pós-processamento, dado o back-end que você selecionou.

Use tf.profile para executar as partes do seu aplicativo que usam tensorflow.js e obter os kernels. Será algo parecido com isto

const profileInfo = await tf.profile(() => {
  // You must profile all uses of tf symbols.
  runAllMyTfjsCode();
});

const kernelNames = profileInfo.kernelNames
console.log(kernelNames);

Copie essa lista de kernels para sua área de transferência para a próxima etapa.

Você precisa criar o perfil do código usando os mesmos back-ends que deseja usar em seu pacote personalizado.

Você precisará repetir esta etapa se seu modelo mudar ou seu código de pré/pós-processamento mudar.

Etapa 2. Escreva um arquivo de configuração para o módulo tfjs personalizado

Aqui está um exemplo de arquivo de configuração.

Parece assim:

{
  "kernels": ["Reshape", "_FusedMatMul", "Identity"],
  "backends": [
      "cpu"
  ],
  "models": [
      "./model/model.json"
  ],
  "outputPath": "./custom_tfjs",
  "forwardModeOnly": true
}
  • kernels: A lista de kernels a serem incluídos no pacote. Copie isso da saída da Etapa 1.
  • back-ends: a lista de back-end(s) que você deseja incluir. As opções válidas incluem “cpu”, “webgl” e “wasm”.
  • modelos: uma lista de arquivos model.json para modelos que você carrega em seu aplicativo. Pode estar vazio se o seu programa não usar tfjs_converter para carregar um modelo de gráfico.
  • outputPath: um caminho para uma pasta onde colocar os módulos gerados.
  • forwardModeOnly: defina como false se desejar incluir gradientes para os kernels listados anteriormente.

Etapa 3. Gere o módulo tfjs personalizado

Execute a ferramenta de compilação personalizada com o arquivo de configuração como argumento. Você precisa ter o pacote @tensorflow/tfjs instalado para ter acesso a esta ferramenta.

npx tfjs-custom-module  --config custom_tfjs_config.json

Isso criará uma pasta em outputPath com alguns novos arquivos.

Etapa 4. Configure seu bundler como alias tfjs para o novo módulo personalizado.

Em bundlers como webpack e rollup, podemos criar um alias às referências existentes aos módulos tfjs para apontar para nossos módulos tfjs personalizados recém-gerados. Existem três módulos que precisam ter alias para economia máxima no tamanho do pacote.

Aqui está um trecho de como fica no webpack ( exemplo completo aqui ):

...

config.resolve = {
  alias: {
    '@tensorflow/tfjs$':
        path.resolve(__dirname, './custom_tfjs/custom_tfjs.js'),
    '@tensorflow/tfjs-core$': path.resolve(
        __dirname, './custom_tfjs/custom_tfjs_core.js'),
    '@tensorflow/tfjs-core/dist/ops/ops_for_converter': path.resolve(
        __dirname, './custom_tfjs/custom_ops_for_converter.js'),
  }
}

...

E aqui está o trecho de código equivalente para rollup ( exemplo completo aqui ):

import alias from '@rollup/plugin-alias';

...

alias({
  entries: [
    {
      find: /@tensorflow\/tfjs$/,
      replacement: path.resolve(__dirname, './custom_tfjs/custom_tfjs.js'),
    },
    {
      find: /@tensorflow\/tfjs-core$/,
      replacement: path.resolve(__dirname, './custom_tfjs/custom_tfjs_core.js'),
    },
    {
      find: '@tensorflow/tfjs-core/dist/ops/ops_for_converter',
      replacement: path.resolve(__dirname, './custom_tfjs/custom_ops_for_converter.js'),
    },
  ],
}));

...

Se o seu bundler não suportar alias de módulo, você precisará alterar suas instruções import para importar tensorflow.js do custom_tfjs.js gerado que foi criado na Etapa 3. As definições de operação não serão alteradas em árvore, mas os kernels ainda serão em árvore -abalado. Geralmente, os kernels que agitam as árvores são o que proporciona a maior economia no tamanho final do pacote.

Se você estiver usando apenas o pacote @tensoflow/tfjs-core, precisará apenas criar um alias para esse pacote.

Etapa 5. Crie seu pacote

Execute seu bundler (por exemplo, webpack ou rollup ) para produzir seu pacote. O tamanho do pacote configurável deve ser menor do que se você executar o pacote configurável sem alias de módulo. Você também pode usar visualizadores como este para ver o que chegou ao seu pacote final.

Etapa 6. Teste seu aplicativo

Certifique-se de testar se seu aplicativo está funcionando conforme o esperado!