Generazione di bundle browser con dimensioni ottimizzate con TensorFlow.js

Panoramica

TensorFlow.js 3.0 offre il supporto per la creazione di bundle di browser ottimizzati per le dimensioni e orientati alla produzione . Per dirla in un altro modo, vogliamo rendere più semplice l'invio di meno JavaScript al browser.

Questa funzionalità è rivolta agli utenti con casi d'uso in produzione che trarrebbero particolare vantaggio dalla riduzione dei byte dal loro carico utile (e sono quindi disposti a impegnarsi per raggiungere questo obiettivo). Per utilizzare questa funzionalità è necessario avere familiarità con i moduli ES , gli strumenti di raggruppamento JavaScript come webpack o rollup e concetti come l'eliminazione del tree-shaking/dead-code .

Questo tutorial dimostra come creare un modulo tensorflow.js personalizzato che può essere utilizzato con un bundler per generare una build con dimensioni ottimizzate per un programma utilizzando tensorflow.js.

Terminologia

Nel contesto di questo documento utilizzeremo alcuni termini chiave:

Moduli ES : il sistema di moduli JavaScript standard . Introdotto in ES6/ES2015. Identificabile mediante l'uso di dichiarazioni di importazione ed esportazione .

Raggruppamento : prendere una serie di risorse JavaScript e raggrupparle/raggrupparle in una o più risorse JavaScript utilizzabili in un browser. Questo è il passaggio che solitamente produce le risorse finali che vengono fornite al browser. Le applicazioni generalmente verranno raggruppate direttamente dalle fonti della libreria transpilate . I bundler più comuni includono rollup e webpack . Il risultato finale del raggruppamento è noto come bundle (o talvolta come blocco se è suddiviso in più parti)

Tree-Shaking/Eliminazione del codice morto : rimozione del codice che non viene utilizzato dall'applicazione scritta finale. Questa operazione viene eseguita durante il raggruppamento, in genere nella fase di minimizzazione.

Operazioni (Ops) - Un'operazione matematica su uno o più tensori che produce uno o più tensori come output. Le operazioni sono codice di "alto livello" e possono utilizzare altre operazioni per definire la propria logica.

Kernel : un'implementazione specifica di un'operazione legata a funzionalità hardware specifiche. I kernel sono di "basso livello" e specifici del backend. Alcune operazioni hanno una mappatura uno a uno dall'operazione al kernel mentre altre operazioni utilizzano più kernel.

Ambito e casi d'uso

Modelli grafici solo di inferenza

Il caso d'uso principale di cui abbiamo sentito parlare dagli utenti in relazione a questo e che supportiamo in questa versione è quello di fare inferenza con i modelli grafici TensorFlow.js . Se utilizzi un modello a livelli TensorFlow.js , puoi convertirlo nel formato del modello grafico utilizzando tfjs-converter . Il formato del modello grafico è più efficiente per il caso d'uso dell'inferenza.

Manipolazione del tensore di basso livello con tfjs-core

L'altro caso d'uso che supportiamo riguarda i programmi che utilizzano direttamente il pacchetto @tensorflow/tjfs-core per la manipolazione del tensore di livello inferiore.

Il nostro approccio alle costruzioni personalizzate

I nostri principi fondamentali nella progettazione di questa funzionalità includono quanto segue:

  • Sfrutta al massimo il sistema di moduli JavaScript (ESM) e consenti agli utenti di TensorFlow.js di fare lo stesso.
  • Rendi TensorFlow.js quanto più facilmente modificabile dagli alberi dai bundler esistenti (ad esempio webpack, rollup, ecc.). Ciò consente agli utenti di sfruttare tutte le funzionalità di questi bundler, comprese funzionalità come la suddivisione del codice.
  • Per quanto possibile, mantenere la facilità d'uso per gli utenti che non sono così sensibili alle dimensioni del bundle . Ciò significa che le build di produzione richiederanno uno sforzo maggiore poiché molte delle impostazioni predefinite nelle nostre librerie supportano la facilità d'uso delle build ottimizzate per le dimensioni.

L'obiettivo principale del nostro flusso di lavoro è produrre un modulo JavaScript personalizzato per TensorFlow.js che contenga solo le funzionalità richieste per il programma che stiamo cercando di ottimizzare. Ci affidiamo ai bundler esistenti per eseguire l'ottimizzazione vera e propria.

Sebbene ci affidiamo principalmente al sistema di moduli JavaScript, forniamo anche uno strumento CLI personalizzato per gestire parti che non sono facili da specificare tramite il sistema di moduli nel codice rivolto all'utente. Due esempi di questo sono:

  • Specifiche del modello archiviate nei file model.json
  • L'op per il sistema di dispacciamento del kernel specifico del backend che utilizziamo.

Ciò rende la generazione di una build tfjs personalizzata un po' più complessa rispetto al semplice puntamento di un bundler al normale pacchetto @tensorflow/tfjs.

Come creare pacchetti personalizzati con dimensioni ottimizzate

Passaggio 1: determina quali kernel utilizza il tuo programma

Questo passaggio ci consente di determinare tutti i kernel utilizzati da qualsiasi modello eseguito o dal codice di pre/post-elaborazione dato il backend selezionato.

Utilizza tf.profile per eseguire le parti della tua applicazione che utilizzano tensorflow.js e ottenere i kernel. Sembrerà qualcosa del genere

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

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

Copia l'elenco dei kernel negli appunti per il passaggio successivo.

Devi profilare il codice utilizzando gli stessi backend che desideri utilizzare nel tuo pacchetto personalizzato.

Dovrai ripetere questo passaggio se il tuo modello cambia o il tuo codice di pre/post-elaborazione cambia.

Passaggio 2. Scrivere un file di configurazione per il modulo tfjs personalizzato

Ecco un file di configurazione di esempio.

Sembra questo:

{
  "kernels": ["Reshape", "_FusedMatMul", "Identity"],
  "backends": [
      "cpu"
  ],
  "models": [
      "./model/model.json"
  ],
  "outputPath": "./custom_tfjs",
  "forwardModeOnly": true
}
  • kernel: l'elenco dei kernel da includere nel bundle. Copialo dall'output del passaggio 1.
  • backend: l'elenco dei backend che desideri includere. Le opzioni valide includono "cpu", "webgl" e "wasm".
  • models: un elenco di file model.json per i modelli caricati nell'applicazione. Può essere vuoto se il tuo programma non utilizza tfjs_converter per caricare un modello grafico.
  • outputPath: percorso di una cartella in cui inserire i moduli generati.
  • forwardModeOnly: impostalo su false se desideri includere gradienti per i kernel elencati in precedenza.

Passaggio 3. Genera il modulo tfjs personalizzato

Esegui lo strumento di creazione personalizzata con il file di configurazione come argomento. È necessario avere installato il pacchetto @tensorflow/tfjs per avere accesso a questo strumento.

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

Questo creerà una cartella in outputPath con alcuni nuovi file.

Passaggio 4. Configura il tuo bundler per eseguire l'alias tfjs sul nuovo modulo personalizzato.

In bundler come webpack e rollup possiamo creare alias i riferimenti esistenti ai moduli tfjs per puntare ai nostri moduli tfjs personalizzati appena generati. Sono tre i moduli a cui è necessario assegnare un alias per ottenere il massimo risparmio in termini di dimensioni del pacchetto.

Ecco uno snippet di come appare nel webpack ( esempio completo qui ):

...

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'),
  }
}

...

Ed ecco lo snippet di codice equivalente per il rollup ( esempio completo qui ):

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 il tuo bundler non supporta l'aliasing dei moduli, dovrai modificare le istruzioni import per importare tensorflow.js dal custom_tfjs.js generato che è stato creato nel passaggio 3. Le definizioni op non verranno modificate ad albero, ma i kernel saranno comunque ad albero -scosso. Generalmente i chicchi che scuotono gli alberi sono quelli che offrono il maggiore risparmio nella dimensione finale del pacchetto.

Se stai utilizzando solo il pacchetto @tensoflow/tfjs-core, devi solo creare l'alias di quel pacchetto.

Passaggio 5. Crea il tuo pacchetto

Esegui il tuo bundler (ad esempio webpack o rollup ) per produrre il tuo bundle. La dimensione del bundle dovrebbe essere inferiore rispetto a quando si esegue il bundler senza alias del modulo. Puoi anche utilizzare visualizzatori come questo per vedere cosa è stato inserito nel tuo pacchetto finale.

Passaggio 6. Testa la tua app

Assicurati di verificare che la tua app funzioni come previsto!