Aperçu
Ce guide décrit les mécanismes permettant de définir des opérations personnalisées (ops), des noyaux et des dégradés dans TensorFlow.js. Il vise à fournir un aperçu des principaux concepts et des indicateurs de code qui démontrent les concepts en action.
A qui s'adresse ce guide ?
Il s'agit d'un guide assez avancé qui aborde certains éléments internes de TensorFlow.js, il peut être particulièrement utile pour les groupes de personnes suivants :
- Utilisateurs avancés de TensorFlow.js intéressés par la personnalisation du comportement de diverses opérations mathématiques (par exemple, les chercheurs remplaçant les implémentations de dégradé existantes ou les utilisateurs qui doivent corriger les fonctionnalités manquantes dans la bibliothèque)
- Les utilisateurs créent des bibliothèques qui étendent TensorFlow.js (par exemple, une bibliothèque générale d'algèbre linéaire construite sur les primitives TensorFlow.js ou un nouveau backend TensorFlow.js).
- Utilisateurs souhaitant contribuer à de nouvelles opérations sur tensorflow.js et souhaitant obtenir un aperçu général du fonctionnement de ces mécanismes.
Il ne s'agit pas d'un guide d'utilisation générale de TensorFlow.js, car il aborde les mécanismes de mise en œuvre internes. Vous n'avez pas besoin de comprendre ces mécanismes pour utiliser TensorFlow.js
Vous devez être à l'aise (ou prêt à essayer) la lecture du code source de TensorFlow.js pour tirer le meilleur parti de ce guide.
Terminologie
Pour ce guide, quelques termes clés sont utiles à décrire dès le départ.
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 option liée à des capacités matérielles/plateformes 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.
Gradient / GradFunc — La définition du « mode arrière » d'une opération/noyau qui calcule la dérivée de cette fonction par rapport à certaines entrées. Les dégradés sont du code de « haut niveau » (non spécifique au backend) et peuvent appeler d'autres opérations ou noyaux.
Registre du noyau - Une carte d'un tuple (nom du noyau, nom du backend) vers une implémentation du noyau.
Registre de dégradé — Une carte d'un nom de noyau à une implémentation de dégradé .
Organisation du code
Les opérations et les dégradés sont définis dans tfjs-core .
Les noyaux sont spécifiques au backend et sont définis dans leurs dossiers backend respectifs (par exemple tfjs-backend-cpu ).
Les opérations personnalisées, les noyaux et les dégradés n'ont pas besoin d'être définis dans ces packages. Mais ils utiliseront souvent des symboles similaires dans leur mise en œuvre.
Implémentation d'opérations personnalisées
Une façon de considérer une opération personnalisée est simplement comme une fonction JavaScript qui renvoie une sortie tensorielle, souvent avec des tenseurs en entrée.
- Certaines opérations peuvent être complètement définies en termes d'opérations existantes et doivent simplement importer et appeler ces fonctions directement. Voici un exemple .
- L'implémentation d'une opération peut également être distribuée vers des noyaux spécifiques au backend. Cela se fait via
Engine.runKernelet sera décrit plus en détail dans la section « implémentation de noyaux personnalisés ». Voici un exemple .
Implémentation de noyaux personnalisés
Les implémentations de noyau spécifiques au backend permettent une implémentation optimisée de la logique pour une opération donnée. Les noyaux sont invoqués par les opérations appelant tf.engine().runKernel() . Une implémentation du noyau est définie par quatre éléments
- Un nom de noyau.
- Le backend dans lequel le noyau est implémenté.
- Entrées : arguments tensoriels de la fonction noyau.
- Attributs : arguments non tenseurs de la fonction noyau.
Voici un exemple d' implémentation du noyau . Les conventions utilisées pour l'implémentation sont spécifiques au backend et sont mieux comprises en examinant l'implémentation et la documentation de chaque backend.
Généralement, les noyaux fonctionnent à un niveau inférieur aux tenseurs et lisent et écrivent directement dans la mémoire qui sera finalement enveloppée dans des tenseurs par tfjs-core.
Une fois qu'un noyau est implémenté, il peut être enregistré auprès de TensorFlow.js en utilisant la fonction registerKernel de tfjs-core. Vous pouvez enregistrer un noyau pour chaque backend dans lequel vous souhaitez que ce noyau fonctionne. Une fois enregistré, le noyau peut être invoqué avec tf.engine().runKernel(...) et TensorFlow.js veillera à être envoyé à l'implémentation dans le backend actif actuel.
Implémentation de dégradés personnalisés
Les dégradés sont généralement définis pour un noyau donné (identifié par le même nom de noyau utilisé dans un appel à tf.engine().runKernel(...) ). Cela permet à tfjs-core d'utiliser un registre pour rechercher les définitions de dégradé pour n'importe quel noyau au moment de l'exécution.
La mise en œuvre de dégradés personnalisés est utile pour :
- Ajout d'une définition de dégradé qui peut ne pas être présente dans la bibliothèque
- Remplacement d'une définition de dégradé existante pour personnaliser le calcul du dégradé pour un noyau donné.
Vous pouvez voir des exemples d' implémentations de dégradés ici .
Une fois que vous avez implémenté un dégradé pour un appel donné, il peut être enregistré auprès de TensorFlow.js en utilisant la fonction registerGradient de tfjs-core.
L'autre approche pour implémenter des dégradés personnalisés qui contourne le registre de dégradés (et permet ainsi de calculer des dégradés pour des fonctions arbitraires de manière arbitraire consiste à utiliser tf.customGrad .
Voici un exemple d'opération dans la bibliothèque utilisant customGrad