Précision mixte

Voir sur TensorFlow.org Exécuter dans Google Colab Voir 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 la formation pour le faire fonctionner plus rapidement 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 précision inférieure, 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 types de précision inférieure doivent être utilisés chaque fois que possible sur ces appareils. Cependant, les variables et quelques calculs doivent toujours être en float32 pour des raisons numériques afin que le modèle s'entraîne avec la même qualité. L'API de précision mixte Keras vous permet d'utiliser un mélange de float16 ou 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

Bien 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 et Cloud TPU récents. 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 dotés d'une capacité de calcul 7.0 ou supérieure verront les meilleures performances bénéficier d'une 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 de la 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 exemples de GPU qui bénéficieront le plus d'une précision mixte incluent les GPU RTX, le V100 et l'A100.

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 générera une erreur sinon.

nvidia-smi -L

GPU 0: Tesla V100-SXM2-16GB (UUID: GPU-99e10c4d-de77-42ee-4524-6c41c4e5e47d)

Tous les Cloud TPU prennent en charge bfloat16.

Même sur les CPU 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 fonctionnera 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é politique 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 ultérieurement à 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: 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 en float16 pour les performances, mais les variables doivent être conservées en 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 politique mixed_float16 améliorera considérablement les performances des GPU NVIDIA avec une capacité de calcul d'au moins 7.0. La règle s'exécutera sur d'autres GPU et CPU, mais n'améliorera peut-être pas les performances. Pour les TPU, la mixed_bfloat16 doit être utilisée à la place.

Construire le modèle

Ensuite, commençons à construire 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'environnement d'exécution TensorFlow domine généralement le temps d'exécution, ce qui rend toute amélioration des performances sur le GPU négligeable. Par conséquent, construisons deux grandes couches Dense avec 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 a donc la politique mixed_float16 car vous avez précédemment défini la politique globale sur mixed_float16 . Cela amènera les couches denses à effectuer des calculs float16 et à avoir des variables float32. Ils lancent leurs entrées sur float16 afin d'effectuer des calculs float16, ce qui fait que leurs sorties sont float16 en conséquence. Leurs variables sont float32 et seront converties en float16 lorsque les couches sont appelées pour éviter les erreurs de non-concordance de 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 devrait être float32. Étant donné que la politique de dtype est mixed_float16 , l'activation softmax aurait normalement un dtype de calcul float16 et des tenseurs float16 de sortie.

Cela peut être corrigé 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 être la politique float32 , qui effectue des 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 bien, mais un softmax à la fin du modèle devrait être dans float32. La raison en est que si le tenseur intermédiaire allant du softmax à la perte est float16 ou bfloat16, des problèmes numériques peuvent survenir.

Vous pouvez remplacer le dtype de n'importe quel calque par 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'inutiles pour ce modèle spécifique, les sorties du modèle peuvent être converties en float32 avec ce qui suit :

# 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, puis 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

Cet exemple convertit les données d'entrée de int8 en float32. Vous n'effectuez pas de conversion en float16 car la division par 255 se fait sur le CPU, 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 les calculs de traitement des entrées 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 type de calcul.

Les poids initiaux du modèle sont récupérés. Cela permettra de s'entraîner à nouveau à 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 [==============================] - 2s 78ms/step - loss: 4.9609 - accuracy: 0.4132 - val_loss: 0.6643 - val_accuracy: 0.8437
Epoch 2/5
6/6 [==============================] - 0s 34ms/step - loss: 0.7752 - accuracy: 0.7789 - val_loss: 0.3098 - val_accuracy: 0.9175
Epoch 3/5
6/6 [==============================] - 0s 34ms/step - loss: 0.3620 - accuracy: 0.8848 - val_loss: 0.3149 - val_accuracy: 0.8969
Epoch 4/5
6/6 [==============================] - 0s 34ms/step - loss: 0.2998 - accuracy: 0.9066 - val_loss: 0.2988 - val_accuracy: 0.9068
Epoch 5/5
6/6 [==============================] - 0s 33ms/step - loss: 0.2298 - accuracy: 0.9285 - val_loss: 0.5062 - val_accuracy: 0.8414
313/313 - 0s - loss: 0.5163 - accuracy: 0.8392
Test loss: 0.5163048505783081
Test accuracy: 0.8392000198364258

Notez que le modèle imprime le temps par étape dans les journaux : par exemple, "25 ms/étape". La première époque peut être plus lente car TensorFlow passe un certain 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, changez la politique de mixed_float16 en float32 dans la section "Définition de la politique dtype", puis réexécutez toutes les cellules jusqu'à ce point. Sur les GPU avec capacité de calcul 7.X, vous devriez voir le temps par étape augmenter de manière significative, indiquant qu'une précision mixte a accéléré le modèle. Assurez-vous de redéfinir la politique 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 dans le modèle jouet de ce guide lors de l'utilisation d'une précision mixte par rapport à float32. Cela est dû à l'utilisation de TensorFloat-32 , qui utilise automatiquement des calculs de précision inférieure 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 toujours des améliorations significatives des performances à partir d'une 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 par rapport à l'exécution d'une précision mixte sur des GPU, en particulier des GPU pré-Ampère. En effet, les TPU effectuent 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 voient généralement moins de gains de performances avec une précision mitigée sur les modèles du monde réel.

Pour de nombreux modèles réels, la précision mixte vous permet également de doubler la taille du lot sans manquer de mémoire, car les tenseurs float16 occupent 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 dtype où chaque lot se compose de l'ensemble de données MNIST complet 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-dépassement numérique. Cette section décrit ce qu'est la mise à l'échelle des pertes et la section suivante explique comment l'utiliser avec une boucle d'apprentissage personnalisée.

Débordement 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 jusqu'à l'infini et que les valeurs inférieures \(6.0 \times 10^{-8}\) jusqu'à zéro. float32 et bfloat16 ont une plage dynamique beaucoup plus élevée, de sorte que le débordement et le sous-dépassement 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épassement se produit également rarement pendant le passage vers l'avant. Cependant, lors de la passe arrière, les gradients peuvent sous-dépasser jusqu'à zéro. La mise à l'échelle des pertes est une technique pour empêcher ce sous-dépassement.

Aperçu 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, par exemple \(1024\), et vous obtenez la valeur de l' échelle des pertes . Cela entraînera également une mise à l'échelle des gradients par \(1024\) , ce qui réduira considérablement le risque de débordement. Une fois les gradients finaux calculés, divisez-les par \(1024\) pour les ramener à leurs valeurs correctes.

Le pseudo-code 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

Choisir une échelle de perte peut être délicat. Si l'échelle de perte est trop faible, les gradients peuvent encore passer à zéro. S'il est trop élevé, c'est l'inverse qui pose problème : les gradients 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 de formation 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îner le modèle avec une boucle d'entraînement personnalisée

Jusqu'à présent, vous avez formé un modèle Keras avec une précision mixte à l'aide de tf.keras.Model.fit . Ensuite, vous utiliserez la précision mixte avec une boucle d'entraînement personnalisée. Si vous ne savez pas déjà 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 modifications par rapport à son exécution dans float32 :

  1. Construisez le modèle avec une précision mitigée (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 encapsule un optimiseur et applique une 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 autrement 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é efficace 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 les tf.data.Dataset :

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 d'étape d'apprentissage. Vous utiliserez deux nouvelles méthodes de l'optimiseur d'échelle de perte pour mettre à l'échelle la perte et annuler l'échelle des 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 en tant qu'entrées et divise chacun par l'échelle de perte pour les réduire à l'échelle

Ces fonctions doivent être utilisées afin d'éviter le sous-dépassement dans les gradients. LossScaleOptimizer.apply_gradients appliquera alors des gradients si aucun d'entre eux n'a Inf s ou NaN s. Il mettra également à jour l'échelle de perte, en la réduisant de moitié si les gradients avaient Inf s ou NaN s et en l'augmentant potentiellement dans le cas contraire.

@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

Le LossScaleOptimizer probablement les premières étapes au début de la formation. L'échelle de perte commence à un niveau élevé afin que l'échelle de perte optimale puisse ê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 réentraîner à 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=4.869325160980225, test accuracy=0.7221999764442444
Epoch 1: loss=0.4893573224544525, test accuracy=0.878000020980835
Epoch 2: loss=0.36011582612991333, test accuracy=0.9440000057220459
Epoch 3: loss=0.27391332387924194, test accuracy=0.9318000078201294
Epoch 4: loss=0.247697651386261, test accuracy=0.933899998664856

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 une 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 du lot augmente généralement le débit de formation, c'est-à-dire les éléments de formation par seconde sur lesquels votre modèle peut s'exécuter.

S'assurer que les cœurs GPU Tensor 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 de 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 Tensor Cores soit utilisé.

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

Vous devriez essayer d'utiliser Tensor Cores lorsque cela est possible. Si vous souhaitez en savoir plus, le guide des performances d'apprentissage en profondeur NVIDIA décrit les exigences exactes pour l'utilisation de Tensor Cores 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 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 des Cloud TPU, car les Tensors bfloat16 utilisent la moitié de la mémoire. Doubler la taille du lot peut augmenter le débit de formation.

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 également au type float32 comme pour 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 une capacité de calcul 7.0, car cela améliorera les performances jusqu'à 3 fois.

  • 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 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 l' optimizer.get_scaled_loss pour mettre à l'échelle la perte et l' optimizer.get_unscaled_gradients pour réduire l'échelle des gradients.

  • Doublez la taille du lot d'apprentissage 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 optimiser les performances

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