Optimisez les performances du GPU TensorFlow avec le profileur TensorFlow

Aperçu

Ce guide vous montrera comment utiliser TensorFlow Profiler avec TensorBoard pour obtenir un aperçu et obtenir les performances maximales de vos GPU, et déboguer lorsqu'un ou plusieurs de vos GPU sont sous-utilisés.

Si vous êtes nouveau sur le Profiler :

Gardez à l’esprit que transférer les calculs vers le GPU n’est pas toujours bénéfique, en particulier pour les petits modèles. Il peut y avoir des frais généraux dus à :

  • Transfert de données entre l'hôte (CPU) et l'appareil (GPU) ; et
  • En raison de la latence impliquée lorsque l'hôte lance les noyaux GPU.

Flux de travail d'optimisation des performances

Ce guide explique comment déboguer les problèmes de performances en commençant par un seul GPU, puis en passant à un seul hôte avec plusieurs GPU.

Il est recommandé de déboguer les problèmes de performances dans l'ordre suivant :

  1. Optimisez et déboguez les performances sur un GPU :
    1. Vérifiez si le pipeline d’entrée constitue un goulot d’étranglement.
    2. Déboguer les performances d'un GPU.
    3. Activez la précision mixte (avec fp16 (float16)) et activez éventuellement XLA .
  2. Optimisez et déboguez les performances sur l'hôte unique multi-GPU.

Par exemple, si vous utilisez une stratégie de distribution TensorFlow pour entraîner un modèle sur un hôte unique avec plusieurs GPU et remarquez une utilisation sous-optimale du GPU, vous devez d'abord optimiser et déboguer les performances d'un GPU avant de déboguer le système multi-GPU.

Comme base pour obtenir du code performant sur les GPU, ce guide suppose que vous utilisez déjà tf.function . Les API Keras Model.compile et Model.fit utiliseront automatiquement tf.function sous le capot. Lors de l'écriture d'une boucle d'entraînement personnalisée avec tf.GradientTape , reportez-vous à Meilleures performances avec tf.function pour savoir comment activer tf.function s.

Les sections suivantes traitent des approches suggérées pour chacun des scénarios ci-dessus afin d'aider à identifier et à résoudre les goulots d'étranglement en matière de performances.

1. Optimisez les performances sur un GPU

Dans un cas idéal, votre programme devrait avoir une utilisation élevée du GPU, une communication minimale du CPU (l'hôte) vers le GPU (le périphérique) et aucune surcharge du pipeline d'entrée.

La première étape de l'analyse des performances consiste à obtenir un profil pour un modèle fonctionnant avec un seul GPU.

La page de présentation du profileur de TensorBoard, qui présente une vue générale des performances de votre modèle lors de l'exécution d'un profil, peut donner une idée de la distance entre votre programme et le scénario idéal.

TensorFlow Profiler Overview Page

Les chiffres clés auxquels il faut prêter attention sur la page de présentation sont :

  1. Quelle part du temps d'étape correspond à l'exécution réelle de l'appareil ?
  2. Le pourcentage d'opérations placées sur l'appareil par rapport à l'hôte
  3. Combien de noyaux utilisent fp16

Atteindre des performances optimales signifie maximiser ces chiffres dans les trois cas. Pour avoir une compréhension approfondie de votre programme, vous devrez vous familiariser avec la visionneuse de traces Profiler de TensorBoard. Les sections ci-dessous présentent certains modèles courants de visionneuse de trace que vous devez rechercher lors du diagnostic des goulots d'étranglement des performances.

Vous trouverez ci-dessous une image d'une vue de trace de modèle exécutée sur un GPU. Dans les sections TensorFlow Name Scope et TensorFlow Ops , vous pouvez identifier différentes parties du modèle, telles que la passe avant, la fonction de perte, le calcul de passe arrière/gradient et la mise à jour du poids de l'optimiseur. Vous pouvez également exécuter les opérations sur le GPU à côté de chaque Stream , qui font référence aux flux CUDA. Chaque flux est utilisé pour des tâches spécifiques. Dans cette trace, Stream#118 est utilisé pour lancer des noyaux de calcul et des copies d'appareil à appareil. Stream#119 est utilisé pour la copie hôte vers périphérique et Stream#120 pour la copie périphérique vers hôte.

La trace ci-dessous montre les caractéristiques communes d'un modèle performant.

image

Par exemple, la chronologie de calcul du GPU ( Stream#118 ) semble « occupée » avec très peu d’intervalles. Il existe un minimum de copies d'un hôte à l'autre ( Stream #119 ) et d'un appareil à l'autre ( Stream #120 ), ainsi que des écarts minimes entre les étapes. Lorsque vous exécutez le profileur pour votre programme, vous ne pourrez peut-être pas identifier ces caractéristiques idéales dans votre vue de trace. Le reste de ce guide couvre les scénarios courants et comment les résoudre.

1. Déboguer le pipeline d'entrée

La première étape du débogage des performances du GPU consiste à déterminer si votre programme est lié aux entrées. Le moyen le plus simple de comprendre cela est d'utiliser l' analyseur de pipeline d'entrée du Profiler, sur TensorBoard, qui fournit un aperçu du temps passé dans le pipeline d'entrée.

image

Vous pouvez entreprendre les actions potentielles suivantes si votre pipeline d'entrée contribue de manière significative au temps d'étape :

  • Vous pouvez utiliser le guide spécifique à tf.data pour apprendre à déboguer votre pipeline d'entrée.
  • Un autre moyen rapide de vérifier si le pipeline d'entrée constitue le goulot d'étranglement consiste à utiliser des données d'entrée générées aléatoirement qui ne nécessitent aucun prétraitement. Voici un exemple d'utilisation de cette technique pour un modèle ResNet. Si le pipeline d'entrée est optimal, vous devriez bénéficier de performances similaires avec des données réelles et avec des données aléatoires/synthétiques générées. La seule surcharge dans le cas des données synthétiques sera due à la copie des données d'entrée qui peuvent à nouveau être préextraites et optimisées.

Reportez-vous également aux meilleures pratiques pour optimiser le pipeline de données d'entrée .

2. Déboguer les performances d'un GPU

Plusieurs facteurs peuvent contribuer à une faible utilisation du GPU. Vous trouverez ci-dessous quelques scénarios couramment observés lors de l’examen du visualiseur de traces et des solutions potentielles.

1. Analyser les écarts entre les étapes

Une observation courante lorsque votre programme ne fonctionne pas de manière optimale est l'écart entre les étapes d'entraînement. Dans l'image de la vue de trace ci-dessous, il y a un grand écart entre les étapes 8 et 9, ce qui signifie que le GPU est inactif pendant cette période.

image

Si votre visualiseur de traces affiche de grands écarts entre les étapes, cela peut indiquer que votre programme est limité aux entrées. Dans ce cas, vous devez vous référer à la section précédente sur le débogage de votre pipeline d'entrée si vous ne l'avez pas déjà fait.

Cependant, même avec un pipeline d'entrée optimisé, vous pouvez toujours avoir des écarts entre la fin d'une étape et le début d'une autre en raison de conflits de threads CPU. tf.data utilise des threads d'arrière-plan pour paralléliser le traitement du pipeline. Ces threads peuvent interférer avec l'activité côté hôte du GPU qui se produit au début de chaque étape, comme la copie de données ou la planification d'opérations GPU.

Si vous remarquez de grandes lacunes du côté de l'hôte, qui planifie ces opérations sur le GPU, vous pouvez définir la variable d'environnement TF_GPU_THREAD_MODE=gpu_private . Cela garantit que les noyaux GPU sont lancés à partir de leurs propres threads dédiés et ne sont pas mis en file d'attente derrière le travail tf.data .

Les écarts entre les étapes peuvent également être provoqués par des calculs de métriques, des rappels Keras ou des opérations en dehors de tf.function qui s'exécutent sur l'hôte. Ces opérations n'ont pas d'aussi bonnes performances que les opérations à l'intérieur d'un graphique TensorFlow. De plus, certaines de ces opérations s'exécutent sur le processeur et copient les tenseurs depuis le GPU.

Si, après avoir optimisé votre pipeline d'entrée, vous remarquez toujours des écarts entre les étapes dans la visionneuse de trace, vous devez examiner le code du modèle entre les étapes et vérifier si la désactivation des rappels/métriques améliore les performances. Certains détails de ces opérations figurent également dans la visionneuse de trace (côté appareil et hôte). La recommandation dans ce scénario est d'amortir la surcharge de ces opérations en les exécutant après un nombre fixe d'étapes au lieu de chaque étape. Lors de l'utilisation de la méthode Model.compile dans l'API tf.keras , la définition de l'indicateur steps_per_execution le fait automatiquement. Pour les boucles d'entraînement personnalisées, utilisez tf.while_loop .

2. Améliorer l'utilisation des appareils

1. Petits noyaux GPU et retards de lancement du noyau hôte

L'hôte met les noyaux en file d'attente pour qu'ils soient exécutés sur le GPU, mais il existe une latence (environ 20 à 40 μs) avant que les noyaux ne soient réellement exécutés sur le GPU. Dans un cas idéal, l'hôte met suffisamment de noyaux en file d'attente sur le GPU pour que le GPU passe la majeure partie de son temps à s'exécuter, plutôt que d'attendre que l'hôte mette en file d'attente davantage de noyaux.

La page de présentation du profileur sur TensorBoard indique la durée pendant laquelle le GPU est resté inactif en raison de l'attente du lancement des noyaux par l'hôte. Dans l'image ci-dessous, le GPU est inactif pendant environ 10 % du temps d'étape en attendant le lancement des noyaux.

image

La visionneuse de traces de ce même programme affiche de petits écarts entre les noyaux lorsque l'hôte est en train de lancer des noyaux sur le GPU.

image

En lançant de nombreuses petites opérations sur le GPU (comme une addition scalaire, par exemple), l'hôte risque de ne pas suivre le GPU. L'outil TensorFlow Stats dans TensorBoard pour le même profil affiche 126 224 opérations Mul prenant 2,77 secondes. Ainsi, chaque noyau dure environ 21,9 μs, ce qui est très petit (à peu près le même temps que la latence de lancement) et peut potentiellement entraîner des retards de lancement du noyau hôte.

image

Si votre visionneuse de traces affiche de nombreux petits écarts entre les opérations sur le GPU, comme dans l'image ci-dessus, vous pouvez :

  • Concaténez de petits tenseurs et utilisez des opérations vectorisées ou utilisez une taille de lot plus grande pour que chaque noyau lancé fasse plus de travail, ce qui occupera le GPU plus longtemps.
  • Assurez-vous que vous utilisez tf.function pour créer des graphiques TensorFlow, afin de ne pas exécuter d'opérations en mode purement impatient. Si vous utilisez Model.fit (par opposition à une boucle d'entraînement personnalisée avec tf.GradientTape ), alors tf.keras.Model.compile le fera automatiquement pour vous.
  • Fusionnez les noyaux en utilisant XLA avec tf.function(jit_compile=True) ou le clustering automatique. Pour plus de détails, consultez la section Activer la précision mixte et XLA ci-dessous pour savoir comment activer XLA pour obtenir des performances plus élevées. Cette fonctionnalité peut conduire à une utilisation élevée de l’appareil.
2. Placement des opérations TensorFlow

La page de présentation du profileur vous montre le pourcentage d'opérations placées sur l'hôte par rapport à l'appareil (vous pouvez également vérifier le placement d'opérations spécifiques en regardant la visionneuse de trace . Comme dans l'image ci-dessous, vous voulez le pourcentage d'opérations sur l'hôte être très petit par rapport à l'appareil.

image

Idéalement, la plupart des opérations de calcul intensives devraient être placées sur le GPU.

Pour savoir à quels appareils les opérations et les tenseurs de votre modèle sont affectés, définissez tf.debugging.set_log_device_placement(True) comme première instruction de votre programme.

Notez que dans certains cas, même si vous spécifiez une opération à placer sur un périphérique particulier, son implémentation peut remplacer cette condition (exemple : tf.unique ). Même pour la formation sur un seul GPU, la spécification d'une stratégie de distribution, telle que tf.distribute.OneDeviceStrategy , peut entraîner un placement plus déterministe des opérations sur votre appareil.

L'une des raisons pour lesquelles la majorité des opérations sont placées sur le GPU est d'éviter des copies de mémoire excessives entre l'hôte et le périphérique (des copies de mémoire pour les données d'entrée/sortie du modèle entre l'hôte et le périphérique sont attendues). Un exemple de copie excessive est illustré dans la vue de trace ci-dessous sur les flux GPU #167 , #168 et #169 .

image

Ces copies peuvent parfois nuire aux performances si elles bloquent l’exécution des noyaux GPU. Les opérations de copie de mémoire dans la visionneuse de trace contiennent plus d'informations sur les opérations qui sont à l'origine de ces tenseurs copiés, mais il n'est pas toujours facile d'associer une memCopy à une opération. Dans ces cas, il est utile d'examiner les opérations à proximité pour vérifier si la copie de la mémoire a lieu au même endroit à chaque étape.

3. Des noyaux plus efficaces sur les GPU

Une fois que l'utilisation du GPU de votre programme est acceptable, l'étape suivante consiste à envisager d'augmenter l'efficacité des noyaux GPU en utilisant des cœurs Tensor ou des opérations de fusion.

1. Utiliser les cœurs tenseurs

Les GPU NVIDIA® modernes disposent de cœurs Tensor spécialisés qui peuvent améliorer considérablement les performances des noyaux éligibles.

Vous pouvez utiliser les statistiques du noyau GPU de TensorBoard pour visualiser quels noyaux GPU sont éligibles à Tensor Core et quels noyaux utilisent des cœurs Tensor. L'activation fp16 (voir la section Activation de la précision mixte ci-dessous) est un moyen de faire en sorte que les noyaux GEMM (General Matrix Multiply) de votre programme (matmul ops) utilisent le Tensor Core. Les noyaux GPU utilisent efficacement les Tensor Cores lorsque la précision est de fp16 et que les dimensions du tenseur d'entrée/sortie sont divisibles par 8 ou 16 (pour int8 ).

Pour d'autres recommandations détaillées sur la manière de rendre les noyaux efficaces pour les GPU, reportez-vous au guide des performances du Deep Learning NVIDIA® .

2. Opérations de fusion

Utilisez tf.function(jit_compile=True) pour fusionner des opérations plus petites afin de former des noyaux plus gros, conduisant à des gains de performances significatifs. Pour en savoir plus, reportez-vous au guide XLA .

3. Activer la précision mixte et XLA

Après avoir suivi les étapes ci-dessus, l'activation de la précision mixte et de XLA sont deux étapes facultatives que vous pouvez suivre pour améliorer davantage les performances. L’approche suggérée consiste à les activer un par un et à vérifier que les avantages en termes de performances sont ceux attendus.

1. Activer une précision mixte

Le guide de précision TensorFlow Mixed montre comment activer la précision fp16 sur les GPU. Activez AMP sur les GPU NVIDIA® pour utiliser les cœurs Tensor et obtenez des accélérations globales jusqu'à 3 fois par rapport à l'utilisation d'une seule précision fp32 (float32) sur Volta et les architectures GPU plus récentes.

Assurez-vous que les dimensions de la matrice/du tenseur satisfont aux exigences d'appel des noyaux qui utilisent des cœurs Tensor. Les noyaux GPU utilisent efficacement les Tensor Cores lorsque la précision est de fp16 et que les dimensions d'entrée/sortie sont divisibles par 8 ou 16 (pour int8).

Notez qu'avec cuDNN v7.6.3 et versions ultérieures, les dimensions de convolution seront automatiquement complétées si nécessaire pour exploiter les cœurs Tensor.

Suivez les bonnes pratiques ci-dessous pour maximiser les avantages en termes de performances de la précision fp16 .

1. Utilisez les noyaux fp16 optimaux

Avec fp16 activé, les noyaux de multiplications matricielles (GEMM) de votre programme doivent utiliser la version fp16 correspondante qui utilise les cœurs Tensor. Cependant, dans certains cas, cela ne se produit pas et vous ne ressentez pas l'accélération attendue suite à l'activation fp16 , car votre programme revient à une implémentation inefficace.

image

La page de statistiques du noyau GPU indique quelles opérations sont éligibles à Tensor Core et quels noyaux utilisent réellement le Tensor Core efficace. Le guide NVIDIA® sur les performances du deep learning contient des suggestions supplémentaires sur la façon d'exploiter les cœurs Tensor. De plus, les avantages de l'utilisation fp16 apparaîtront également dans les noyaux qui étaient auparavant limités en mémoire, car désormais les opérations prendront la moitié du temps.

2. Mise à l'échelle des pertes dynamiques et statiques

Une mise à l'échelle des pertes est nécessaire lors de l'utilisation fp16 pour éviter un sous-débordement dû à une faible précision. Il existe deux types de mise à l'échelle des pertes, dynamique et statique, tous deux expliqués plus en détail dans le guide Mixed Precision . Vous pouvez utiliser la stratégie mixed_float16 pour activer automatiquement la mise à l'échelle des pertes dans l'optimiseur Keras.

Lorsque vous essayez d'optimiser les performances, il est important de garder à l'esprit que la mise à l'échelle dynamique des pertes peut introduire des opérations conditionnelles supplémentaires exécutées sur l'hôte et entraîner des écarts qui seront visibles entre les étapes dans la visionneuse de trace. D'un autre côté, la mise à l'échelle des pertes statiques n'entraîne pas de tels frais généraux et peut constituer une meilleure option en termes de performances, avec la nécessité de spécifier la valeur correcte de l'échelle de perte statique.

2. Activez XLA avec tf.function(jit_compile=True) ou le clustering automatique

Comme dernière étape pour obtenir les meilleures performances avec un seul GPU, vous pouvez expérimenter l'activation de XLA, qui fusionnera les opérations et entraînera une meilleure utilisation des appareils et une empreinte mémoire réduite. Pour plus de détails sur la façon d'activer XLA dans votre programme avec tf.function(jit_compile=True) ou le clustering automatique, reportez-vous au guide XLA .

Vous pouvez définir le niveau JIT global sur -1 (désactivé), 1 ou 2 . Un niveau plus élevé est plus agressif et peut réduire le parallélisme et utiliser plus de mémoire. Définissez la valeur sur 1 si vous avez des restrictions de mémoire. Notez que XLA ne fonctionne pas bien pour les modèles avec des formes de tenseur d'entrée variables, car le compilateur XLA devrait continuer à compiler les noyaux chaque fois qu'il rencontre de nouvelles formes.

2. Optimiser les performances sur l'hôte unique multi-GPU

L'API tf.distribute.MirroredStrategy peut être utilisée pour faire évoluer la formation de modèles d'un GPU vers plusieurs GPU sur un seul hôte. (Pour en savoir plus sur la façon d'effectuer une formation distribuée avec TensorFlow, reportez-vous aux guides Formation distribuée avec TensorFlow , Utiliser un GPU et Utiliser des TPU ainsi qu'au didacticiel Formation distribuée avec Keras .)

Même si la transition d’un GPU à plusieurs GPU devrait idéalement être évolutive dès le départ, vous pouvez parfois rencontrer des problèmes de performances.

Lorsque vous passez d'une formation avec un seul GPU à plusieurs GPU sur le même hôte, idéalement, vous devriez bénéficier d'une mise à l'échelle des performances avec uniquement la surcharge supplémentaire de communication par gradient et une utilisation accrue des threads de l'hôte. En raison de cette surcharge, vous n'aurez pas une accélération exacte de 2x si vous passez de 1 à 2 GPU, par exemple.

La vue de trace ci-dessous montre un exemple de surcharge de communication supplémentaire lors de l'entraînement sur plusieurs GPU. Il y a une certaine surcharge pour concaténer les dégradés, les communiquer entre les répliques et les diviser avant de procéder à la mise à jour du poids.

image

La liste de contrôle suivante vous aidera à obtenir de meilleures performances lors de l'optimisation des performances dans le scénario multi-GPU :

  1. Essayez de maximiser la taille du lot, ce qui entraînera une utilisation plus élevée des appareils et amortira les coûts de communication sur plusieurs GPU. L'utilisation du profileur de mémoire permet d'avoir une idée de la proximité de votre programme par rapport à l'utilisation maximale de la mémoire. Notez que même si une taille de lot plus élevée peut affecter la convergence, cela est généralement contrebalancé par les avantages en termes de performances.
  2. Lors du passage d'un seul GPU à plusieurs GPU, le même hôte doit désormais traiter beaucoup plus de données d'entrée. Ainsi, après (1), il est recommandé de revérifier les performances du pipeline d’entrée et de s’assurer qu’il ne s’agit pas d’un goulot d’étranglement.
  3. Vérifiez la chronologie du GPU dans la vue de trace de votre programme pour détecter tout appel AllReduce inutile, car cela entraîne une synchronisation sur tous les appareils. Dans la vue de trace présentée ci-dessus, AllReduce est effectué via le noyau NCCL et il n'y a qu'un seul appel NCCL sur chaque GPU pour les dégradés de chaque étape.
  4. Recherchez les opérations de copie D2H, H2D et D2D inutiles qui peuvent être minimisées.
  5. Vérifiez la durée de l'étape pour vous assurer que chaque réplique effectue le même travail. Par exemple, il peut arriver qu'un GPU (généralement GPU0 ) soit sursouscrit parce que l'hôte finit par y consacrer plus de travail par erreur.
  6. Enfin, vérifiez l'étape de formation sur tous les GPU dans votre vue de trace pour toutes les opérations qui s'exécutent de manière séquentielle. Cela se produit généralement lorsque votre programme inclut des dépendances de contrôle d'un GPU à un autre. Dans le passé, le débogage des performances dans cette situation était résolu au cas par cas. Si vous observez ce comportement dans votre programme, signalez un problème GitHub avec des images de votre vue de trace.

1. Optimiser le dégradé AllReduce

Lors d'un entraînement avec une stratégie synchrone, chaque appareil reçoit une partie des données d'entrée.

Après avoir calculé les passages aller et retour dans le modèle, les gradients calculés sur chaque appareil doivent être agrégés et réduits. Ce gradient AllReduce se produit après le calcul du gradient sur chaque appareil et avant que l'optimiseur ne mette à jour les poids du modèle.

Chaque GPU concatène d'abord les dégradés entre les couches du modèle, les communique entre les GPU à l'aide de tf.distribute.CrossDeviceOps ( tf.distribute.NcclAllReduce est la valeur par défaut), puis renvoie les dégradés après réduction par couche.

L'optimiseur utilisera ces dégradés réduits pour mettre à jour les poids de votre modèle. Idéalement, ce processus devrait se produire en même temps sur tous les GPU pour éviter toute surcharge.

Le temps nécessaire à AllReduce doit être approximativement le même que :

(number of parameters * 4bytes)/ (communication bandwidth)

Ce calcul est utile pour vérifier rapidement si les performances que vous obtenez lors de l'exécution d'une tâche de formation distribuée sont conformes aux attentes ou si vous devez effectuer un débogage supplémentaire des performances. Vous pouvez obtenir le nombre de paramètres de votre modèle à partir de Model.summary .

Notez que chaque paramètre de modèle a une taille de 4 octets puisque TensorFlow utilise fp32 (float32) pour communiquer les dégradés. Même lorsque fp16 est activé, NCCL AllReduce utilise les paramètres fp32 .

Pour bénéficier des avantages de la mise à l’échelle, le temps d’étape doit être beaucoup plus élevé que ces frais généraux. Une façon d'y parvenir consiste à utiliser une taille de lot plus élevée, car la taille du lot affecte le temps d'étape, mais n'a pas d'impact sur la surcharge de communication.

2. Conflit de thread hôte GPU

Lors de l'exécution de plusieurs GPU, le travail du CPU consiste à occuper tous les appareils en lançant efficacement les noyaux GPU sur les appareils.

Cependant, lorsqu'il existe de nombreuses opérations indépendantes que le CPU peut planifier sur un GPU, le CPU peut décider d'utiliser un grand nombre de ses threads hôtes pour occuper un GPU, puis lancer des noyaux sur un autre GPU dans un ordre non déterministe. . Cela peut provoquer une distorsion ou une mise à l'échelle négative, ce qui peut affecter négativement les performances.

La visionneuse de traces ci-dessous montre la surcharge lorsque le processeur échelonne le lancement inefficace du noyau GPU, car GPU1 est inactif, puis commence à exécuter des opérations après le démarrage GPU2 .

image

La vue de trace de l'hôte montre que l'hôte lance les noyaux sur GPU2 avant de les lancer sur GPU1 (notez que les opérations tf_Compute* ci-dessous ne sont pas indicatives des threads CPU).

image

Si vous rencontrez ce type d'étalement des noyaux GPU dans la vue de trace de votre programme, l'action recommandée est la suivante :

  • Définissez la variable d'environnement TensorFlow TF_GPU_THREAD_MODE sur gpu_private . Cette variable d'environnement indiquera à l'hôte de garder les threads d'un GPU privés.
  • Par défaut, TF_GPU_THREAD_MODE=gpu_private définit le nombre de threads à 2, ce qui est suffisant dans la plupart des cas. Cependant, ce nombre peut être modifié en définissant la variable d'environnement TensorFlow TF_GPU_THREAD_COUNT sur le nombre de threads souhaité.