Como gerar pacotes de navegador com tamanho otimizado com o TensorFlow.js

Visão geral

O TensorFlow.js 3.0 oferece suporte para a criação de pacotes de navegador orientados à produção e otimizados para tamanho . Em outras palavras, queremos facilitar o envio de menos JavaScript para o navegador.

Esse recurso é voltado para usuários com casos de uso de produção que se beneficiariam particularmente de reduzir 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 os Módulos ES , ferramentas de agrupamento JavaScript, como webpack ou rollup , e conceitos como eliminação de código morto/destruição de árvore .

Este tutorial demonstra como criar um módulo tensorflow.js personalizado que pode ser usado com um empacotador 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/agrupá-los em um ou mais ativos JavaScript que podem ser usados ​​em um navegador. Essa é a etapa que geralmente produz os ativos finais que são exibidos ao navegador. Os aplicativos geralmente fazem seu próprio pacote diretamente de fontes de biblioteca transpiladas . Os bundlers comuns incluem rollup e webpack . O resultado final do agrupamento é conhecido como um pacote (ou às vezes como um pedaço se for dividido em várias partes)

Tree-Shaking / Dead Code Eliminação - 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 um op vinculado a recursos de hardware específicos. Os kernels são de 'baixo nível' e específicos de back-end. Algumas operações têm um mapeamento um-para-um de op para kernel, enquanto outras operações usam vários kernels.

Escopo e Casos de Uso

Modelos de gráfico apenas de inferência

O principal caso de uso que ouvimos de usuários relacionados a isso e que estamos dando suporte nesta versão é o de fazer inferência com modelos de gráfico do TensorFlow.js . Se você estiver usando um modelo de camadas do TensorFlow.js , poderá convertê-lo no formato de modelo de gráfico usando o tfjs-converter . O formato do modelo de 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 suportamos são programas que usam diretamente o pacote @tensorflow/tjfs-core para manipulação de tensor de nível inferior.

Nossa abordagem para construções personalizadas

Nossos princípios fundamentais ao projetar essa 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 instável possível pelos bundlers 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, pois muitos dos padrões em nossas bibliotecas oferecem suporte à facilidade de uso em compilações otimizadas de tamanho.

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 os 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
• O sistema de despacho de kernel específico de back-end que usamos.

Isso torna a geração de uma compilaçã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 qualquer modelo que você executa ou código de pré/pós-processamento de acordo com o backend selecionado.

Use tf.profile para executar as partes de seu aplicativo que usam tensorflow.js e obter os kernels. Vai parecer algo assim

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 for alterado ou seu código de pré/pós-processamento for alterado.

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

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

Se parece com isso:

{
  "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-ends que você deseja incluir. As opções válidas incluem "cpu", "webgl" e "wasm".
• models: uma lista de arquivos model.json para modelos que você carrega em seu aplicativo. Pode estar vazio se seu programa não usar tfjs_converter para carregar um modelo de gráfico.
• outputPath: Um caminho para uma pasta para 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 para alias tfjs ao novo módulo customizado.

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

Aqui está um trecho do que parece 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 de 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 sacudidas em árvore, mas os kernels ainda serão em árvore - abalado. Geralmente, os kernels que agitam a árvore são o que fornece a maior economia no tamanho final do pacote.

Se você estiver usando apenas o pacote @tensoflow/tfjs-core, precisará apenas alias desse pacote.

Etapa 5. Crie seu pacote

Execute seu bundler (por exemplo webpack ou rollup ) para produzir seu bundle. O tamanho do bundle deve ser menor do que se você executar o bundler 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!