Générer des bundles de navigateur de taille optimisée avec TensorFlow.js

Aperçu

TensorFlow.js 3.0 prend en charge la création de bundles de navigateurs optimisés en termes de taille et orientés vers la production . En d’autres termes, nous souhaitons vous permettre d’envoyer plus facilement moins de JavaScript au navigateur.

Cette fonctionnalité est destinée aux utilisateurs ayant des cas d'utilisation en production qui bénéficieraient particulièrement de la suppression d'octets de leur charge utile (et sont donc prêts à faire des efforts pour y parvenir). Pour utiliser cette fonctionnalité, vous devez être familier avec les modules ES , les outils de regroupement JavaScript tels que webpack ou rollup et les concepts tels que l'élimination du tree-shaking/dead-code .

Ce didacticiel montre comment créer un module tensorflow.js personnalisé qui peut être utilisé avec un bundler pour générer une version optimisée en taille pour un programme utilisant tensorflow.js.

Terminologie

Dans le contexte de ce document, nous utiliserons quelques termes clés :

Modules ES - Le système de modules JavaScript standard . Introduit dans ES6/ES2015. Identifiable grâce à l'utilisation de déclarations d'importation et d'exportation .

Regroupement : prendre un ensemble d'actifs JavaScript et les regrouper/regrouper en un ou plusieurs actifs JavaScript utilisables dans un navigateur. Il s'agit de l'étape qui produit généralement les ressources finales qui sont envoyées au navigateur. Les applications effectueront généralement leur propre regroupement directement à partir des sources de bibliothèque transpilées . Les bundles courants incluent le rollup et le webpack . Le résultat final du regroupement est ce qu'on appelle un bundle (ou parfois un morceau s'il est divisé en plusieurs parties).

Tree-Shaking / Dead Code Elimination - Suppression du code qui n'est pas utilisé par l'application écrite finale. Cela se fait lors du regroupement, généralement lors de l’étape de minification.

Opérations (Ops) - Une opération mathématique sur un ou plusieurs tenseurs qui produit un ou plusieurs tenseurs en sortie. Les opérations sont du code de « haut niveau » et peuvent utiliser d’autres opérations pour définir leur logique.

Noyau - Une implémentation spécifique d'une opération liée à des capacités matérielles spécifiques. Les noyaux sont de « bas niveau » et spécifiques au backend. Certaines opérations ont un mappage un-à-un de l'opération au noyau tandis que d'autres utilisent plusieurs noyaux.

Portée et cas d'utilisation

Modèles graphiques d'inférence uniquement

Le principal cas d'utilisation dont nous avons entendu parler par les utilisateurs à ce sujet et que nous prenons en charge dans cette version est celui de l' inférence avec les modèles graphiques TensorFlow.js . Si vous utilisez un modèle de couches TensorFlow.js , vous pouvez le convertir au format de modèle graphique à l'aide du tfjs-converter . Le format de modèle graphique est plus efficace pour le cas d’utilisation de l’inférence.

Manipulation du Tensor de bas niveau avec tfjs-core

L'autre cas d'utilisation que nous prenons en charge concerne les programmes qui utilisent directement le package @tensorflow/tjfs-core pour la manipulation du tenseur de niveau inférieur.

Notre approche des constructions personnalisées

Nos principes fondamentaux lors de la conception de cette fonctionnalité sont les suivants :

• Utilisez au maximum le système de modules JavaScript (ESM) et permettez aux utilisateurs de TensorFlow.js de faire de même.
• Rendre TensorFlow.js aussi arborescent que possible par les bundles existants (par exemple, webpack, rollup, etc.). Cela permet aux utilisateurs de profiter de toutes les capacités de ces bundlers, y compris des fonctionnalités telles que le fractionnement de code.
• Dans la mesure du possible, maintenez la facilité d'utilisation pour les utilisateurs qui ne sont pas aussi sensibles à la taille du bundle . Cela signifie que les versions de production nécessiteront plus d'efforts, car la plupart des valeurs par défaut de nos bibliothèques prennent en charge la facilité d'utilisation par rapport aux versions optimisées en taille.

L'objectif principal de notre flux de travail est de produire un module JavaScript personnalisé pour TensorFlow.js qui contient uniquement les fonctionnalités requises pour le programme que nous essayons d'optimiser. Nous nous appuyons sur les bundles existants pour effectuer l'optimisation proprement dite.

Bien que nous nous appuyions principalement sur le système de modules JavaScript, nous fournissons également un outil CLI personnalisé pour gérer les parties qui ne sont pas faciles à spécifier via le système de modules dans le code destiné à l'utilisateur. Voici deux exemples :

• Spécifications du modèle stockées dans les fichiers model.json
• L'option du système de répartition du noyau spécifique au backend que nous utilisons.

Cela rend la génération d'une version tfjs personnalisée un peu plus complexe que le simple fait de pointer un bundler vers le package @tensorflow/tfjs standard.

Comment créer des offres groupées personnalisées de taille optimisée

Étape 1 : Déterminez les noyaux utilisés par votre programme

Cette étape nous permet de déterminer tous les noyaux utilisés par les modèles que vous exécutez ou le code de pré/post-traitement en fonction du backend que vous avez sélectionné.

Utilisez tf.profile pour exécuter les parties de votre application qui utilisent tensorflow.js et obtenir les noyaux. Cela ressemblera à quelque chose comme ça

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

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

Copiez cette liste de noyaux dans votre presse-papiers pour l'étape suivante.

Vous devez profiler le code en utilisant le(s) même(s) backend(s) que vous souhaitez utiliser dans votre bundle personnalisé.

Vous devrez répéter cette étape si votre modèle change ou si votre code de pré/post-traitement change.

Étape 2. Écrivez un fichier de configuration pour le module tfjs personnalisé

Voici un exemple de fichier de configuration.

Cela ressemble à ceci :

{
  "kernels": ["Reshape", "_FusedMatMul", "Identity"],
  "backends": [
      "cpu"
  ],
  "models": [
      "./model/model.json"
  ],
  "outputPath": "./custom_tfjs",
  "forwardModeOnly": true
}

• kernels : La liste des noyaux à inclure dans le bundle. Copiez ceci à partir du résultat de l'étape 1.
• backends : la liste des backends que vous souhaitez inclure. Les options valides incluent "cpu", "webgl" et "wasm".
• models : une liste de fichiers model.json pour les modèles que vous chargez dans votre application. Peut être vide si votre programme n'utilise pas tfjs_converter pour charger un modèle graphique.
• OutputPath : un chemin vers un dossier dans lequel placer les modules générés.
• forwardModeOnly : définissez ceci sur false si vous souhaitez inclure des dégradés pour les noyaux répertoriés précédemment.

Étape 3. Générez le module tfjs personnalisé

Exécutez l'outil de construction personnalisé avec le fichier de configuration comme argument. Vous devez avoir installé le package @tensorflow/tfjs pour avoir accès à cet outil.

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

Cela créera un dossier sur outputPath avec quelques nouveaux fichiers.

Étape 4. Configurez votre bundler pour qu'il alias tfjs vers le nouveau module personnalisé.

Dans les bundles comme webpack et rollup, nous pouvons alias les références existantes aux modules tfjs pour pointer vers nos modules tfjs personnalisés nouvellement générés. Il existe trois modules qui doivent être aliasés pour maximiser les économies de taille du bundle.

Voici un extrait de ce à quoi cela ressemble dans webpack ( exemple complet ici ) :

...

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

...

Et voici l'extrait de code équivalent pour le rollup ( exemple complet ici ) :

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

...

Si votre bundler ne prend pas en charge l'alias de module, vous devrez modifier vos instructions import pour importer tensorflow.js à partir de custom_tfjs.js généré qui a été créé à l'étape 3. Les définitions d'opération ne seront pas ébranlées, mais les noyaux seront toujours arborescents. -secoué. En général, ce sont les noyaux qui secouent les arbres qui permettent de réaliser les plus grandes économies sur la taille finale du paquet.

Si vous utilisez uniquement le package @tensoflow/tfjs-core, il vous suffit de créer un alias pour ce package.

Étape 5. Créez votre forfait

Exécutez votre bundler (par exemple webpack ou rollup ) pour produire votre bundle. La taille du bundle doit être plus petite que si vous exécutez le bundler sans alias de module. Vous pouvez également utiliser des visualiseurs comme celui-ci pour voir ce qui a été intégré à votre bundle final.

Étape 6. Testez votre application

Assurez-vous de tester que votre application fonctionne comme prévu !