Précision mixte

Voir sur TensorFlow.org Exécuter dans Google Colab Afficher la source sur GitHub Télécharger le cahier

Aperçu

La précision mixte est l'utilisation de types à virgule flottante 16 bits et 32 ​​bits dans un modèle pendant l'entraînement pour le rendre plus rapide et utiliser moins de mémoire. En conservant certaines parties du modèle dans les types 32 bits pour la stabilité numérique, le modèle aura un temps de pas inférieur et s'entraînera également en termes de métriques d'évaluation telles que la précision. Ce guide décrit comment utiliser l'API de précision mixte Keras pour accélérer vos modèles. L'utilisation de cette API peut améliorer les performances de plus de 3 fois sur les GPU modernes et de 60% sur les TPU.

Aujourd'hui, la plupart des modèles utilisent le dtype float32, qui prend 32 bits de mémoire. Cependant, il existe deux dtypes de moindre précision, float16 et bfloat16, chacun prenant 16 bits de mémoire à la place. Les accélérateurs modernes peuvent exécuter des opérations plus rapidement dans les dtypes 16 bits, car ils disposent d'un matériel spécialisé pour exécuter des calculs 16 bits et les dtypes 16 bits peuvent être lus plus rapidement à partir de la mémoire.

Les GPU NVIDIA peuvent exécuter des opérations dans float16 plus rapidement que dans float32, et les TPU peuvent exécuter des opérations dans bfloat16 plus rapidement que float32. Par conséquent, ces dtypes de moindre précision doivent être utilisés autant que possible sur ces appareils. Cependant, les variables et quelques calculs devraient toujours être dans float32 pour des raisons numériques afin que le modèle s'entraîne à la même qualité. L'API de précision mixte Keras vous permet d'utiliser un mélange de float16 ou de bfloat16 avec float32, pour obtenir les avantages de performance de float16 / bfloat16 et les avantages de stabilité numérique de float32.

Installer

import tensorflow as tf

from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import mixed_precision

Matériel pris en charge

Alors que la précision mixte fonctionnera sur la plupart des matériels, elle n'accélérera que les modèles sur les GPU NVIDIA récents et les Cloud TPU. Les GPU NVIDIA prennent en charge l'utilisation d'un mélange de float16 et float32, tandis que les TPU prennent en charge un mélange de bfloat16 et float32.

Parmi les GPU NVIDIA, ceux avec une capacité de calcul 7.0 ou supérieure verront le plus grand avantage en termes de performances de la précision mixte, car ils disposent d'unités matérielles spéciales, appelées Tensor Cores, pour accélérer les multiplications et les convolutions de la matrice float16. Les GPU plus anciens n'offrent aucun avantage en termes de performances mathématiques pour l'utilisation d'une précision mixte, mais les économies de mémoire et de bande passante peuvent permettre certaines accélérations. Vous pouvez rechercher la capacité de calcul de votre GPU sur la page Web du GPU CUDA de NVIDIA. Les GPU RTX, le V100 et l'A100 sont des exemples de GPU qui bénéficieront le plus de la précision mixte.

Vous pouvez vérifier votre type de GPU avec ce qui suit. La commande n'existe que si les pilotes NVIDIA sont installés, donc ce qui suit soulèvera une erreur dans le cas contraire.

nvidia-smi -L
GPU 0: NVIDIA Tesla V100-SXM2-16GB (UUID: GPU-5fbd083a-e390-c0a4-23e2-c42f67b39b4a)

Tous les Cloud TPU prennent en charge bfloat16.

Même sur les processeurs et les GPU plus anciens, où aucune accélération n'est attendue, les API de précision mixte peuvent toujours être utilisées pour les tests unitaires, le débogage ou simplement pour essayer l'API. Sur les processeurs, la précision mixte s'exécutera cependant beaucoup plus lentement.

Définition de la stratégie dtype

Pour utiliser la précision mixte dans Keras, vous devez créer un tf.keras.mixed_precision.Policy , généralement appelé stratégie dtype . Les stratégies Dtype spécifient les couches dtypes dans lesquelles s'exécuteront. Dans ce guide, vous allez construire une stratégie à partir de la chaîne 'mixed_float16' et la définir comme stratégie globale. Cela amènera les couches créées par la suite à utiliser une précision mixte avec un mélange de float16 et float32.

policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)
INFO:tensorflow:Mixed precision compatibility check (mixed_float16): OK
Your GPU will likely run quickly with dtype policy mixed_float16 as it has compute capability of at least 7.0. Your GPU: NVIDIA Tesla V100-SXM2-16GB, compute capability 7.0

Pour faire court, vous pouvez directement passer une chaîne à set_global_policy , ce qui est généralement fait dans la pratique.

# Equivalent to the two lines above
mixed_precision.set_global_policy('mixed_float16')

La politique spécifie deux aspects importants d'une couche: le dtype dans lequel les calculs de la couche sont effectués et le dtype des variables d'une couche. Ci-dessus, vous avez créé une politique mixed_float16 (c'est-à-dire, une mixed_precision.Policy créée en passant la chaîne 'mixed_float16' à son constructeur). Avec cette politique, les couches utilisent des calculs float16 et des variables float32. Les calculs sont effectués dans float16 pour les performances, mais les variables doivent être conservées dans float32 pour la stabilité numérique. Vous pouvez interroger directement ces propriétés de la stratégie.

print('Compute dtype: %s' % policy.compute_dtype)
print('Variable dtype: %s' % policy.variable_dtype)
Compute dtype: float16
Variable dtype: float32

Comme mentionné précédemment, la stratégie mixed_float16 améliorera de manière plus significative les performances des GPU NVIDIA avec une capacité de calcul d'au moins 7,0. La stratégie s'exécutera sur d'autres GPU et CPU, mais n'améliorera peut-être pas les performances. Pour les TPU, la stratégie mixed_bfloat16 doit être utilisée à la place.

Construire le modèle

Ensuite, commençons à créer un modèle simple. Les très petits modèles de jouets ne bénéficient généralement pas d'une précision mixte, car la surcharge de l'exécution de TensorFlow domine généralement le temps d'exécution, rendant toute amélioration des performances du GPU négligeable. Par conséquent, construisons deux grandes couches Dense de 4096 unités chacune si un GPU est utilisé.

inputs = keras.Input(shape=(784,), name='digits')
if tf.config.list_physical_devices('GPU'):
  print('The model will run with 4096 units on a GPU')
  num_units = 4096
else:
  # Use fewer units on CPUs so the model finishes in a reasonable amount of time
  print('The model will run with 64 units on a CPU')
  num_units = 64
dense1 = layers.Dense(num_units, activation='relu', name='dense_1')
x = dense1(inputs)
dense2 = layers.Dense(num_units, activation='relu', name='dense_2')
x = dense2(x)
The model will run with 4096 units on a GPU

Chaque couche a une politique et utilise la politique globale par défaut. Chacune des couches Dense donc la politique mixed_float16 car vous mixed_float16 précédemment défini la politique globale sur mixed_float16 . Cela amènera les couches denses à faire des calculs float16 et à avoir des variables float32. Ils convertissent leurs entrées en float16 afin de faire des calculs float16, ce qui fait que leurs sorties sont float16 en conséquence. Leurs variables sont float32 et seront transtypées en float16 lorsque les couches sont appelées pour éviter les erreurs dues aux incohérences dtype.

print(dense1.dtype_policy)
print('x.dtype: %s' % x.dtype.name)
# 'kernel' is dense1's variable
print('dense1.kernel.dtype: %s' % dense1.kernel.dtype.name)
<Policy "mixed_float16">
x.dtype: float16
dense1.kernel.dtype: float32

Ensuite, créez les prédictions de sortie. Normalement, vous pouvez créer les prédictions de sortie comme suit, mais ce n'est pas toujours numériquement stable avec float16.

# INCORRECT: softmax and model output will be float16, when it should be float32
outputs = layers.Dense(10, activation='softmax', name='predictions')(x)
print('Outputs dtype: %s' % outputs.dtype.name)
Outputs dtype: float16

Une activation softmax à la fin du modèle doit être float32. Comme la politique dtype est mixed_float16 , l'activation softmax aurait normalement un dtype de calcul float16 et des tenseurs de sortie float16.

Cela peut être résolu en séparant les couches Dense et softmax, et en passant dtype='float32' à la couche softmax:

# CORRECT: softmax and model output are float32
x = layers.Dense(10, name='dense_logits')(x)
outputs = layers.Activation('softmax', dtype='float32', name='predictions')(x)
print('Outputs dtype: %s' % outputs.dtype.name)
Outputs dtype: float32

Passer dtype='float32' au constructeur de couche softmax remplace la politique dtype de la couche pour qu'elle devienne la politique float32 , qui effectue les calculs et conserve les variables dans float32. De manière équivalente, vous auriez pu passer à la place dtype=mixed_precision.Policy('float32') ; les couches convertissent toujours l'argument dtype en stratégie. Étant donné que la couche d' Activation n'a pas de variables, la variable dtype de la politique est ignorée, mais le dtype de calcul de la politique de float32 fait que softmax et la sortie du modèle sont float32.

Ajouter un softmax float16 au milieu d'un modèle est très bien, mais un softmax à la fin du modèle doit être dans float32. La raison en est que si le tenseur intermédiaire s'écoulant du softmax vers la perte est float16 ou bfloat16, des problèmes numériques peuvent survenir.

Vous pouvez remplacer le dtype de n'importe quelle couche pour qu'il soit float32 en passant dtype='float32' si vous pensez qu'il ne sera pas numériquement stable avec les calculs float16. Mais généralement, cela n'est nécessaire que sur la dernière couche du modèle, car la plupart des couches ont une précision suffisante avec mixed_float16 et mixed_bfloat16 .

Même si le modèle ne se termine pas par un softmax, les sorties doivent toujours être float32. Bien qu'elles ne soient pas nécessaires pour ce modèle spécifique, les sorties du modèle peuvent être converties en float32 avec les éléments suivants:

# The linear activation is an identity function. So this simply casts 'outputs'
# to float32. In this particular case, 'outputs' is already float32 so this is a
# no-op.
outputs = layers.Activation('linear', dtype='float32')(outputs)

Ensuite, terminez et compilez le modèle et générez les données d'entrée:

model = keras.Model(inputs=inputs, outputs=outputs)
model.compile(loss='sparse_categorical_crossentropy',
              optimizer=keras.optimizers.RMSprop(),
              metrics=['accuracy'])

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train.reshape(60000, 784).astype('float32') / 255
x_test = x_test.reshape(10000, 784).astype('float32') / 255
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11493376/11490434 [==============================] - 0s 0us/step

Cet exemple convertit les données d'entrée de int8 en float32. Vous ne convertissez pas en float16 car la division par 255 est sur le processeur, qui exécute les opérations float16 plus lentement que les opérations float32. Dans ce cas, la différence de performances est négligeable, mais en général, vous devez exécuter le traitement des entrées mathématiques dans float32 s'il s'exécute sur le processeur. La première couche du modèle convertira les entrées en float16, car chaque couche convertit les entrées en virgule flottante en son dtype de calcul.

Les poids initiaux du modèle sont récupérés. Cela permettra à nouveau de s'entraîner à partir de zéro en chargeant les poids.

initial_weights = model.get_weights()

Entraîner le modèle avec Model.fit

Ensuite, entraînez le modèle:

history = model.fit(x_train, y_train,
                    batch_size=8192,
                    epochs=5,
                    validation_split=0.2)
test_scores = model.evaluate(x_test, y_test, verbose=2)
print('Test loss:', test_scores[0])
print('Test accuracy:', test_scores[1])
Epoch 1/5
6/6 [==============================] - 1s 61ms/step - loss: 8.6004 - accuracy: 0.3090 - val_loss: 1.2854 - val_accuracy: 0.6582
Epoch 2/5
6/6 [==============================] - 0s 29ms/step - loss: 0.8481 - accuracy: 0.7540 - val_loss: 0.4211 - val_accuracy: 0.8837
Epoch 3/5
6/6 [==============================] - 0s 29ms/step - loss: 0.3853 - accuracy: 0.8770 - val_loss: 0.4489 - val_accuracy: 0.8570
Epoch 4/5
6/6 [==============================] - 0s 29ms/step - loss: 0.2945 - accuracy: 0.9076 - val_loss: 0.3176 - val_accuracy: 0.9022
Epoch 5/5
6/6 [==============================] - 0s 28ms/step - loss: 0.3327 - accuracy: 0.8990 - val_loss: 0.1799 - val_accuracy: 0.9427
313/313 - 0s - loss: 0.1818 - accuracy: 0.9425
Test loss: 0.1818101704120636
Test accuracy: 0.9424999952316284

Notez que le modèle imprime le temps par étape dans les journaux: par exemple, "25ms / étape". La première période peut être plus lente car TensorFlow passe du temps à optimiser le modèle, mais ensuite le temps par étape devrait se stabiliser.

Si vous exécutez ce guide dans Colab, vous pouvez comparer les performances de la précision mixte avec float32. Pour ce faire, modifiez la stratégie de mixed_float16 en float32 dans la section «Définition de la stratégie dtype», puis réexécutez toutes les cellules jusqu'à ce point. Sur les GPU dotés de la capacité de calcul 7.X, vous devriez voir le temps par étape augmenter considérablement, ce qui indique que la précision mixte accélère le modèle. Assurez-vous de mixed_float16 la stratégie sur mixed_float16 et réexécutez les cellules avant de continuer avec le guide.

Sur les GPU avec une capacité de calcul d'au moins 8,0 (GPU Ampère et plus), vous ne verrez probablement aucune amélioration des performances du modèle de jouet dans ce guide lorsque vous utilisez une précision mixte par rapport à float32. Cela est dû à l'utilisation de TensorFloat-32 , qui utilise automatiquement des mathématiques de moindre précision dans certaines opérations float32 telles que tf.linalg.matmul . TensorFloat-32 offre certains des avantages de performance de la précision mixte lors de l'utilisation de float32. Cependant, dans les modèles du monde réel, vous constaterez généralement des améliorations de performances significatives par rapport à la précision mixte en raison des économies de bande passante mémoire et des opérations que TensorFloat-32 ne prend pas en charge.

Si vous exécutez une précision mixte sur un TPU, vous ne verrez pas autant de gain de performances que l'exécution d'une précision mixte sur des GPU, en particulier des GPU pré-ampères. C'est parce que les TPU font certaines opérations dans bfloat16 sous le capot même avec la politique dtype par défaut de float32. Ceci est similaire à la façon dont les GPU Ampere utilisent TensorFloat-32 par défaut. Par rapport aux GPU Ampere, les TPU enregistrent généralement moins de gains de performances avec une précision mixte sur les modèles du monde réel.

Pour de nombreux modèles du monde réel, la précision mixte vous permet également de doubler la taille du lot sans manquer de mémoire, car les tenseurs float16 prennent la moitié de la mémoire. Cela ne s'applique cependant pas à ce modèle de jouet, car vous pouvez probablement exécuter le modèle dans n'importe quel type de dtype où chaque lot comprend l'ensemble de données MNIST de 60 000 images.

Mise à l'échelle des pertes

La mise à l'échelle des pertes est une technique que tf.keras.Model.fit exécute automatiquement avec la politique mixed_float16 pour éviter le sous- mixed_float16 numérique. Cette section décrit ce qu'est la mise à l'échelle des pertes et la section suivante décrit comment l'utiliser avec une boucle d'apprentissage personnalisée.

Dépassement inférieur et débordement

Le type de données float16 a une plage dynamique étroite par rapport à float32. Cela signifie que les valeurs supérieures à 65504 $ déborderont à l'infini et les valeurs inférieures à 6,0 $ \ fois 10 ^ {- 8} $ seront inférieures à zéro. float32 et bfloat16 ont une plage dynamique beaucoup plus élevée, de sorte que les débordements et sous-débordements ne sont pas un problème.

Par example:

x = tf.constant(256, dtype='float16')
(x ** 2).numpy()  # Overflow
inf
x = tf.constant(1e-5, dtype='float16')
(x ** 2).numpy()  # Underflow
0.0

En pratique, le débordement avec float16 se produit rarement. De plus, le sous-débordement se produit rarement pendant la passe avant. Cependant, lors du passage en arrière, les dégradés peuvent être inférieurs à zéro. La mise à l'échelle des pertes est une technique pour éviter ce sous-débit.

Vue d'ensemble de la mise à l'échelle des pertes

Le concept de base de la mise à l'échelle des pertes est simple: multipliez simplement la perte par un grand nombre, disons 1024 $, et vous obtenez la valeur de l' échelle de perte . Cela entraînera également une mise à l'échelle des dégradés de 1024 $, ce qui réduira considérablement le risque de sous-débordement. Une fois les dégradés finaux calculés, divisez-les par 1024 $ pour les ramener à leurs valeurs correctes.

Le pseudocode de ce processus est:

loss_scale = 1024
loss = model(inputs)
loss *= loss_scale
# Assume `grads` are float32. You do not want to divide float16 gradients.
grads = compute_gradient(loss, model.trainable_variables)
grads /= loss_scale

Le choix d'une échelle de perte peut être délicat. Si l'échelle de perte est trop faible, les gradients peuvent encore déborder jusqu'à zéro. S'il est trop élevé, le problème inverse se produit: les dégradés peuvent déborder à l'infini.

Pour résoudre ce problème, TensorFlow détermine dynamiquement l'échelle de perte afin que vous n'ayez pas à en choisir une manuellement. Si vous utilisez tf.keras.Model.fit , la mise à l'échelle des pertes est effectuée pour vous afin que vous n'ayez pas à faire de travail supplémentaire. Si vous utilisez une boucle d'entraînement personnalisée, vous devez explicitement utiliser le wrapper d'optimisation spécial tf.keras.mixed_precision.LossScaleOptimizer afin d'utiliser la mise à l'échelle des pertes. C'est décrit dans la section suivante.

Entraînement du modèle avec une boucle d'entraînement personnalisée

Jusqu'à présent, vous avez entraîné un modèle Keras avec une précision mixte à l'aide de tf.keras.Model.fit . Ensuite, vous utiliserez une précision mixte avec une boucle d'entraînement personnalisée. Si vous ne savez pas encore ce qu'est une boucle d'entraînement personnalisée, veuillez d'abord lire le guide d'entraînement personnalisé .

L'exécution d'une boucle d'entraînement personnalisée avec une précision mixte nécessite deux changements par rapport à son exécution dans float32:

  1. Construisez le modèle avec une précision mixte (vous l'avez déjà fait)
  2. Utilisez explicitement la mise à l'échelle des pertes si mixed_float16 est utilisé.

Pour l'étape (2), vous utiliserez la classe tf.keras.mixed_precision.LossScaleOptimizer , qui tf.keras.mixed_precision.LossScaleOptimizer un optimiseur et applique la mise à l'échelle des pertes. Par défaut, il détermine dynamiquement l'échelle de perte afin que vous n'ayez pas à en choisir une. Construisez un LossScaleOptimizer comme suit.

optimizer = keras.optimizers.RMSprop()
optimizer = mixed_precision.LossScaleOptimizer(optimizer)

Si vous le souhaitez, il est possible de choisir une échelle de perte explicite ou de personnaliser le comportement de mise à l'échelle des pertes, mais il est fortement recommandé de conserver le comportement de mise à l'échelle des pertes par défaut, car il s'est avéré bien fonctionner sur tous les modèles connus. Consultez la documentation tf.keras.mixed_precision.LossScaleOptimizer si vous souhaitez personnaliser le comportement de mise à l'échelle des pertes.

Ensuite, définissez l'objet de perte et lestf.data.Dataset s:

loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
train_dataset = (tf.data.Dataset.from_tensor_slices((x_train, y_train))
                 .shuffle(10000).batch(8192))
test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(8192)

Ensuite, définissez la fonction de l'étape d'entraînement. Vous utiliserez deux nouvelles méthodes de l'optimiseur d'échelle de perte pour mettre à l'échelle la perte et redimensionner les gradients:

  • get_scaled_loss(loss) : multiplie la perte par l'échelle de perte
  • get_unscaled_gradients(gradients) : prend une liste de dégradés mis à l'échelle comme entrées, et divise chacun par l'échelle de perte pour les redimensionner

Ces fonctions doivent être utilisées afin d'éviter un sous-débordement dans les dégradés. LossScaleOptimizer.apply_gradients appliquera alors des dégradés si aucun d'entre eux n'a Inf s ou NaN s. Il mettra également à jour l'échelle de perte, la divisant par deux si les gradients avaient Inf s ou NaN s et l'augmentant potentiellement autrement.

@tf.function
def train_step(x, y):
  with tf.GradientTape() as tape:
    predictions = model(x)
    loss = loss_object(y, predictions)
    scaled_loss = optimizer.get_scaled_loss(loss)
  scaled_gradients = tape.gradient(scaled_loss, model.trainable_variables)
  gradients = optimizer.get_unscaled_gradients(scaled_gradients)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  return loss

LossScaleOptimizer probablement les premières étapes au début de l'entraînement. L'échelle de perte est élevée au départ, de sorte que l'échelle de perte optimale peut être rapidement déterminée. Après quelques étapes, l'échelle de perte se stabilisera et très peu d'étapes seront sautées. Ce processus se produit automatiquement et n'affecte pas la qualité de la formation.

Maintenant, définissez l'étape de test:

@tf.function
def test_step(x):
  return model(x, training=False)

Chargez les poids initiaux du modèle afin de pouvoir vous entraîner à nouveau à partir de zéro:

model.set_weights(initial_weights)

Enfin, exécutez la boucle d'entraînement personnalisée:

for epoch in range(5):
  epoch_loss_avg = tf.keras.metrics.Mean()
  test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
      name='test_accuracy')
  for x, y in train_dataset:
    loss = train_step(x, y)
    epoch_loss_avg(loss)
  for x, y in test_dataset:
    predictions = test_step(x)
    test_accuracy.update_state(y, predictions)
  print('Epoch {}: loss={}, test accuracy={}'.format(epoch, epoch_loss_avg.result(), test_accuracy.result()))
Epoch 0: loss=3.7858102321624756, test accuracy=0.7168999910354614
Epoch 1: loss=0.5736480355262756, test accuracy=0.8765000104904175
Epoch 2: loss=0.3463210463523865, test accuracy=0.9394000172615051
Epoch 3: loss=0.33968204259872437, test accuracy=0.9503999948501587
Epoch 4: loss=0.2318723201751709, test accuracy=0.939300000667572

Conseils sur les performances du GPU

Voici quelques conseils de performance lors de l'utilisation de la précision mixte sur les GPU.

Augmenter la taille de votre lot

Si cela n'affecte pas la qualité du modèle, essayez d'exécuter avec le double de la taille du lot lorsque vous utilisez la précision mixte. Comme les tenseurs float16 utilisent la moitié de la mémoire, cela vous permet souvent de doubler la taille de votre lot sans manquer de mémoire. L'augmentation de la taille des lots augmente généralement le débit d'entraînement, c'est-à-dire les éléments d'entraînement par seconde sur lesquels votre modèle peut s'exécuter.

S'assurer que les cœurs de mesure GPU sont utilisés

Comme mentionné précédemment, les GPU NVIDIA modernes utilisent une unité matérielle spéciale appelée Tensor Cores qui peut multiplier très rapidement les matrices float16. Cependant, Tensor Cores nécessite que certaines dimensions des tenseurs soient un multiple de 8. Dans les exemples ci-dessous, un argument est en gras si et seulement s'il doit être un multiple de 8 pour que les Tensor Cores soient utilisés.

  • tf.keras.layers.Dense ( unités = 64 )
  • tf.keras.layers.Conv2d ( filtres = 48 , kernel_size = 7, foulée = 3)
    • Et de même pour d'autres couches convolutives, telles que tf.keras.layers.Conv3d
  • tf.keras.layers.LSTM ( unités = 64 )
    • Et similaire pour d'autres RNN, tels que tf.keras.layers.GRU
  • tf.keras.Model.fit (epochs = 2, batch_size = 128 )

Vous devriez essayer d'utiliser les cœurs Tensor lorsque cela est possible. Si vous souhaitez en savoir plus, le guide des performances de l'apprentissage profond NVIDIA décrit les exigences exactes pour l'utilisation de Tensor Core ainsi que d'autres informations sur les performances liées à Tensor Core.

XLA

XLA est un compilateur qui peut encore augmenter les performances de précision mixte, ainsi que les performances de float32 dans une moindre mesure. Reportez-vous au guide XLA pour plus de détails.

Conseils sur les performances de Cloud TPU

Comme pour les GPU, vous devriez essayer de doubler la taille de votre lot lorsque vous utilisez Cloud TPU, car les tenseurs bfloat16 utilisent la moitié de la mémoire. Le fait de doubler la taille du lot peut augmenter le débit d'entraînement.

Les TPU ne nécessitent aucun autre réglage mixte spécifique à la précision pour obtenir des performances optimales. Ils nécessitent déjà l'utilisation de XLA. Les TPU bénéficient du fait que certaines dimensions sont des multiples de 128 $, mais cela s'applique aussi bien au type float32 qu'à la précision mixte. Consultez le guide des performances de Cloud TPU pour obtenir des conseils généraux sur les performances du TPU, qui s'appliquent à la précision mixte ainsi qu'aux tenseurs float32.

Résumé

  • Vous devez utiliser une précision mixte si vous utilisez des TPU ou des GPU NVIDIA avec au moins la capacité de calcul 7.0, car cela améliorera les performances jusqu'à 3x.
  • Vous pouvez utiliser la précision mixte avec les lignes suivantes:

    # On TPUs, use 'mixed_bfloat16' instead
    mixed_precision.set_global_policy('mixed_float16')
    
  • Si votre modèle se termine par softmax, assurez-vous qu'il s'agit bien de float32. Et quelle que soit la fin de votre modèle, assurez-vous que la sortie est float32.

  • Si vous utilisez une boucle d'entraînement personnalisée avec mixed_float16 , en plus des lignes ci-dessus, vous devez envelopper votre optimiseur avec un tf.keras.mixed_precision.LossScaleOptimizer . Appelez ensuite optimizer.get_scaled_loss pour mettre à l'échelle la perte et optimizer.get_unscaled_gradients pour annuler l'échelle des dégradés.

  • Doublez la taille du lot de formation si cela ne réduit pas la précision de l'évaluation

  • Sur les GPU, assurez-vous que la plupart des dimensions du tenseur sont un multiple de 8 $ pour maximiser les performances

Pour plus d'exemples de précision mixte utilisant l'API tf.keras.mixed_precision , consultez le référentiel officiel des modèles . La plupart des modèles officiels, tels que ResNet et Transformer , fonctionneront avec une précision mixte en passant --dtype=fp16 .