mixed_precision.ipynb

Kernel: Python 3

Copyright 2019 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Precisión mixta

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver el código fuente en GitHub

Descargar el bloc de notas

Descripción general

La precisión mixta es el uso de tipos de punto variable de 16 y 32 bits en un modelo durante el entrenamiento para que funcione más rápido y utilice menos memoria. Al mantener ciertas partes del modelo en los tipos de 32 bits para conseguir estabilidad numérica, el modelo tendrá un tiempo de paso menor y se entrenará igual de bien en términos de métricas de evaluación como la precisión. Esta guía describe cómo utilizar la API de precisión mixta de Keras para acelerar sus modelos. El uso de esta API puede mejorar el rendimiento en más de 3 veces en GPUs modernas, 60% en las TPUs y más de 2 veces en las últimas CPUs de Intel.

Actualmente, la mayoría de los modelos utilizan el dtype float32, que consume 32 bits de memoria. Sin embargo, existen dos dtypes de menor precisión, float16 y bfloat16, cada uno de los cuales consume 16 bits de memoria. Los aceleradores modernos pueden ejecutar operaciones más rápidamente en los dtypes de 16 bits, ya que disponen de hardware especializado para ejecutar cálculos de 16 bits y los dtypes de 16 bits pueden leerse desde la memoria más rápidamente.

Las GPUs de NVIDIA pueden ejecutar operaciones en float16 más rápido que en float32, y las TPUs y CPUs de Intel compatibles pueden ejecutar operaciones en bfloat16 más rápido que en float32. Por lo tanto, estos dtypes de menor precisión deben utilizarse siempre que sea posible en esos dispositivos. Sin embargo, las variables y algunos cálculos deberían seguir estando en float32 por razones numéricas para que el modelo se entrene con la misma calidad. La API de precisión mixta de Keras le permite utilizar una mezcla de float16 o bfloat16 con float32, para obtener los beneficios de rendimiento de float16/bfloat16 y los beneficios de estabilidad numérica de float32.

Nota: En esta guía, el término "estabilidad numérica" se refiere a cómo se ve afectada la calidad de un modelo por el uso de un dtype de menor precisión en vez de un dtype de mayor precisión. Una operación es "numéricamente inestable" en float16 o bfloat16 si su ejecución en uno de esos dtypes hace que el modelo tenga peor precisión de evaluación u otras métricas comparado con la ejecución de la operación en float32.

Preparación

import tensorflow as tf

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

Hardware compatible

Aunque la precisión mixta funcionará en la mayoría del hardware, sólo acelerará los modelos en GPUs deNVIDIA recientes, TPUs Cloud y CPUs Intel recientes. Las GPU de NVIDIA admiten una combinación de float16 y float32, mientras que las TPU y las CPU de Intel admiten una combinación de bfloat16 y float32.

Entre las GPUs de NVIDIA, las que tienen capacidad de cálculo 7.0 o superior son las que más se benefician de la precisión mixta porque disponen de unidades de hardware especiales, denominadas Tensor Cores, para acelerar las multiplicaciones y convoluciones de matrices float16. Las GPU más antiguas no ofrecen ninguna ventaja de rendimiento matemático por el uso de la precisión mixta, aunque el ahorro de la memoria y el ancho de banda puede permitir algunos incrementos en la velocidad. Puede consultar la capacidad de cálculo de su GPU en la página web de NVIDIA CUDA GPU. Entre las GPU que más se beneficiarán de la precisión mixta se incluyen las GPU RTX, la V100 y la A100.

Entre las CPU de Intel, a partir de los procesadores Intel Xeon de 4ª generación (nombre en clave Sapphire Rapids), se obtendrá el mayor beneficio de rendimiento de la precisión mixta, ya que pueden acelerar los cálculos bfloat16 utilizando instrucciones AMX (requiere Tensorflow 2.12 o posterior).

Nota: Si ejecuta esta guía en Google Colab, el tiempo de ejecución de la GPU generalmente tiene una P100 conectada. La P100 tiene capacidad de cálculo de 6.0 y no se espera que muestre un aumento significativo de la velocidad. Si se utiliza el tiempo de ejecución de la CPU, es posible que se produzca una ralentización, ya que es probable que el tiempo de ejecución tenga una CPU sin AMX.

Puede verificar su tipo de GPU haciendo lo siguiente. El comando sólo existe si los controladores de NVIDIA están instalados, por lo que lo siguiente producirá un error en caso contrario.

!nvidia-smi -L

Todos los TPUs de la nube admiten bfloat16.

Incluso en CPUs Intel antiguas, otras CPUs x86 sin AMX y GPUs antiguas, donde no se espera un aumento de velocidad, las APIs de precisión mixta pueden seguir utilizándose para realizar pruebas unitarias, depuración o simplemente para probar la API. Sin embargo, mixed_bfloat16 en CPUs que no utilicen instrucciones AMX y mixed_float16 en todas las CPUs x86 funcionarán significativamente más lento.

Preparar la política dtype

Para usar la precisión mixta en Keras, necesita crear una tf.keras.mixed_precision.Policy, generalmente denominada dtype policy. Las políticas dtype especifican los dtypes en los que se ejecutarán las capas. En esta guía, construirá una política a partir de la cadena 'mixed_float16' y la establecerá como política global. Esto hará que las capas creadas posteriormente utilicen una precisión mixta con una mezcla de float16 y float32.

policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)

Para abreviar, puede pasar directamente una cadena a set_global_policy, lo cual se hace normalmente en la práctica.

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

La política especifica dos aspectos importantes de una capa: el dtype en el que se realizan los cálculos de la capa y el dtype de las variables de una capa. Arriba, creó una política mixed_float16 (es decir, una mixed_precision.Policy que se creó pasando la cadena 'mixed_float16' a su constructor). Con esta política, las capas utilizan cálculos en float16 y variables en float32. Los cálculos se realizan en float16 por cuestiones de rendimiento, pero las variables deben mantenerse en float32 por estabilidad numérica. Puede consultar directamente estas propiedades de la política.

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

Como se mencionó anteriormente, la política mixed_float16 mejorará el rendimiento de forma más significativa en las GPU de NVIDIA con una capacidad de cálculo de al menos 7.0. La política funcionará en otras GPU y CPU, pero es posible que no mejore el rendimiento. Para TPUs y CPUs, la política mixed_bfloat16 debe utilizarse como alternativa.

Construyendo el modelo

A continuación, vamos a construir un modelo sencillo. Los modelos simples muy pequeños no suelen beneficiarse de la precisión mixta, ya que la sobrecarga del tiempo de ejecución de TensorFlow suele dominar el tiempo de ejecución, haciendo que cualquier mejora de rendimiento en la GPU sea insignificante. Por lo tanto, vamos a construir dos grandes capas Dense con 4096 unidades cada una si se utiliza una GPU.

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)

Cada capa tiene una política y utiliza la política global de forma predeterminada. Por lo tanto, cada una de las capas Dense tiene la política mixed_float16 porque usted estableció previamente la política global a mixed_float16. Esto hará que las capas densas hagan cálculos float16 y tengan variables float32. Ellas convierten sus entradas a float16 para hacer cálculos float16, lo que causa que sus salidas sean float16 como consecuencia. Sus variables son float32 y serán convertidas a float16 cuando las capas sean llamadas para evitar errores por desajustes de tipo.

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)

A continuación, cree las predicciones de salida. Normalmente, puede crear las predicciones de salida de la siguiente manera, pero esto no siempre es numéricamente estable con 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)

Una activación softmax al final del modelo debería ser float32. Dado que la política de dtype es mixed_float16, la activación softmax tendría normalmente un dtype de cálculo float16 y tensores de salida float16.

Esto puede solucionarse separando las capas Dense y softmax, y pasando dtype='float32' a la capa 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)

Si pasa dtype='float32' al constructor de la capa softmax, la política dtype de la capa se convierte en la política float32, que realiza los cálculos y mantiene las variables en float32. De forma equivalente, podría haber pasado dtype=mixed_precision.Policy('float32'); las capas siempre convierten el argumento dtype en una política. Como la capa Activation no tiene variables, el dtype variable de la política se ignora, pero el dtype de cálculo de la política de float32 hace que softmax y la salida del modelo sean float32.

Si se añade un softmax en float16 en medio de un modelo está bien, pero un softmax al final del modelo debe estar en float32. La razón es que si el tensor intermedio que fluye desde el softmax a la pérdida es float16 o bfloat16, pueden producirse problemas numéricos.

Puede anular el dtype de cualquier capa para que sea float32 pasando dtype='float32' si cree que no será numéricamente estable con cálculos en float16. Pero generalmente, esto sólo es necesario en la última capa del modelo, ya que la mayoría de las capas tienen suficiente precisión con mixed_float16 y mixed_bfloat16.

Incluso si el modelo no termina en un softmax, las salidas deben ser float32. Aunque no es necesario para este modelo específico, las salidas del modelo se pueden convertir a float32 con lo siguiente:

# 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)

Posteriormente, finalice y compile el modelo, y genere los datos de entrada:

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

En este ejemplo los datos de entrada se convierten de int8 a float32. No se convierte a float16 porque la división por 255 se realiza en la CPU, que ejecuta las operaciones en float16 más lentamente que las operaciones en float32. En este caso, la diferencia de rendimiento es insignificante, pero en general se debe ejecutar el procesamiento matemático de entrada en float32 si se ejecuta en la CPU. La primera capa del modelo convertirá las entradas a float16, ya que cada capa convierte las entradas de punto flotante a su dtype de cálculo.

Se recuperan los pesos iniciales del modelo. Esto permitirá volver a entrenar desde cero al cargar los pesos.

initial_weights = model.get_weights()

Entrenar el modelo con Model.fit

A continuación, entrene el modelo:

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])

Observe que el modelo imprime el tiempo por paso en los registros: por ejemplo, "25ms/paso". La primera época puede ser más lenta ya que TensorFlow pasa un poco de tiempo optimizando el modelo, pero después el tiempo por paso debería estabilizarse.

Si está ejecutando esta guía en Colab, puede comparar el rendimiento de la precisión mixta con float32. Para ello, cambie la política de mixed_float16 a float32 en la sección "Configuración de la política dtype" y, a continuación, vuelva a ejecutar todas las celdas hasta este punto. En las GPU con capacidad de cálculo 7.X, debería ver que el tiempo por paso aumenta significativamente, lo que indica que la precisión mixta aceleró el modelo. Asegúrese de volver a cambiar la política a mixed_float16 y vuelva a ejecutar las celdas antes de continuar con la guía.

En GPUs con capacidad de cálculo de al menos 8.0 (GPUs Ampere y superiores), es probable que no observe ninguna mejora de rendimiento en el modelo que se presenta en esta guía cuando utilice precisión mixta en comparación con float32. Esto se debe al uso de TensorFloat-32, que utiliza automáticamente matemáticas de menor precisión en ciertas operaciones de float32 como tf.linalg.matmul. TensorFloat-32 ofrece algunas de las ventajas de rendimiento de la precisión mixta cuando se utiliza float32. Sin embargo, en los modelos del mundo real, todavía se experimentan mejoras significativas en el rendimiento de la precisión mixta debido al ahorro de ancho de banda de memoria y las operaciones que TensorFloat-32 no admite.

Si ejecuta precisión mixta en una TPU, no verá tanta ganancia de rendimiento en comparación con la ejecución de precisión mixta en GPUs, especialmente GPUs pre-Ampere. Esto se debe a que las TPU realizan ciertas operaciones en bfloat16 incluso con la política dtype predeterminada de float32. Esto es similar a cómo las GPUs Ampere utilizan TensorFloat-32 de forma predeterminada. En comparación con las GPUs Ampere, las TPUs normalmente obtienen menos ganancias de rendimiento con precisión mixta en modelos reales.

Para muchos modelos reales, la precisión mixta también permite duplicar el tamaño del lote sin quedarse sin memoria, ya que los tensores float16 ocupan la mitad de la memoria. Sin embargo, esto no se aplica a este modelo experimental, ya que es posible ejecutar el modelo en cualquier tipo de datos en el que cada lote consista en el conjunto de datos MNIST de 60,000 imágenes.

Pérdida de escala

La pérdida de escala es una técnica que tf.keras.Model.fit realiza automáticamente con la política mixed_float16 para evitar el desbordamiento numérico. En esta sección se describe qué es la pérdida de escala y en la siguiente sección se describe cómo utilizarla con un bucle de entrenamiento personalizado.

Nota: Cuando se utiliza la política mixed_bfloat16, no es necesario realizar la pérdida de escala.

Subdesbordamiento y desbordamiento

El tipo de datos float16 tiene un rango dinámico estrecho en comparación con float32. Esto significa que los valores por encima de $65504$ se desbordarán hasta el infinito y los valores por debajo de $6.0 \times 10^{-8}$ se desbordarán por debajo de cero. float32 y bfloat16 tienen un rango dinámico mucho mayor, por lo que el desbordamiento y el subdesbordamiento no son un problema.

Por ejemplo:

x = tf.constant(256, dtype='float16')
(x ** 2).numpy()  # Overflow

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

En la práctica, el desbordamiento con float16 raramente ocurre. Además, el desbordamiento también se produce muy pocas veces durante el siguiente paso. Sin embargo, durante el paso anterior, los gradientes pueden desbordarse por debajo de cero. La pérdida de escala es una técnica para prevenir este desbordamiento.

Resumen de la pérdida de escala

El concepto básico de la pérdida de escala es simple: basta con multiplicar la pérdida por un número grande, por ejemplo $1024$, y se obtiene el valor de loss scale. Esto hará que los gradientes también se escalen por $1024$, reduciendo considerablemente la posibilidad de que haya un subdesbordamiento. Cuando los gradientes finales estén calculados, divídalos por $1024$ para que devuelvan sus valores correctos.

El pseudocódigo de este proceso es:

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

Puede ser complicado elegir una pérdida de escala. Si la pérdida de escala es demasiado baja, los gradientes pueden subdesbordarse por debajo de cero. Si es demasiado alta, ocurre lo contrario: los gradientes pueden desbordarse hasta el infinito.

Para resolver esto, TensorFlow determina dinámicamente la pérdida de escala por lo que no tiene que elegir uno manualmente. Si utiliza tf.keras.Model.fit, la pérdida de escala la realiza usted para que no tenga que hacer ningún trabajo extra. Si utiliza un bucle de entrenamiento personalizado, debe utilizar explícitamente el optimizador especial tf.keras.mixed_precision.LossScaleOptimizer para utilizar la pérdida de escala. Esto se describe en la siguiente sección.

Cómo entrenar el modelo con un bucle de entrenamiento personalizado

Hasta ahora, entrenó un modelo Keras con precisión mixta usando tf.keras.Model.fit. A continuación, utilizará la precisión mixta con un bucle de entrenamiento personalizado. Si aún no conoce lo que es un bucle de entrenamiento personalizado, lea primero la Guía de entrenamiento personalizado.

Ejecutar un bucle de entrenamiento personalizado con precisión mixta requiere dos cambios respecto a ejecutarlo en float32:

1. Construya el modelo con precisión mixta (ya lo ha hecho)

2. Utilice explícitamente la pérdida de escala mediante mixed_float16.

Para el paso (2), utilizará la clase tf.keras.mixed_precision.LossScaleOptimizer, que envuelve un optimizador y aplica la pérdida de escala. De forma predeterminada, determina dinámicamente la pérdida de escala por lo que no tiene que elegir uno. Construya un LossScaleOptimizer como se indica a continuación.

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

Si lo desea, es posible elegir una pérdida de escala explícita o personalizar el comportamiento de la pérdida de escala, pero es muy recomendable mantener el comportamiento de pérdida de escala de forma predeterminada, ya que se ha demostrado que funciona bien en todos los modelos conocidos. Consulte la documentación tf.keras.mixed_precision.LossScaleOptimizer si desea personalizar el comportamiento de la pérdida de escala.

Después, defina el objeto de pérdida y el 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)

Después defina la función de paso de entrenamiento. Utilizará dos nuevos métodos del optimizador de pérdida de escala para escalar la pérdida y desescalar los gradientes:

• get_scaled_loss(loss): Multiplica la pérdida por la pérdida de escala

• get_unscaled_gradients(gradients): Toma una lista de gradientes escalados como entrada, y divide cada uno por la pérdida de escala para desescalarlos

Estas funciones deben utilizarse para prevenir el subdesbordamiento en los gradientes. LossScaleOptimizer.apply_gradients aplicará los gradientes si ninguno de ellos tiene Infs o NaN. También actualizará la pérdida de escala, reduciéndola a la mitad si los gradientes tenían Inf o NaN y potencialmente la incrementará si no es así.

@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

El LossScaleOptimizer probablemente se saltará los primeros pasos al inicio del entrenamiento. La pérdida de escala se inicia en un nivel alto para que la pérdida de escala óptima pueda determinarse rápidamente. Después de unos pocos pasos, la pérdida de escala se estabilizará y se saltará muy pocos pasos. Este proceso se produce automáticamente y no afectará a la calidad del entrenamiento.

A continuación, defina el paso de prueba:

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

Cargue los pesos iniciales del modelo, para poder volver a entrenarlo desde cero:

model.set_weights(initial_weights)

Por último, ejecute el bucle de entrenamiento personalizado:

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()))

Consejos sobre el rendimiento de la GPU

Estos son algunos consejos de rendimiento cuando se emplea la precisión mixta en las GPU.

Incremente el tamaño de su lote

Si no afecta a la calidad del modelo, pruebe a ejecutar con el doble del tamaño del lote cuando utilice la precisión mixta. Como los tensores float16 utilizan la mitad de la memoria, esto frecuentemente le permitirá duplicar el tamaño del lote sin quedarse sin memoria. Aumentar el tamaño del lote normalmente aumenta el rendimiento del entrenamiento, es decir, los elementos de entrenamiento por segundo con los que puede funcionar el modelo.

Garantice el uso de los Tensor Cores de la GPU

Como ya se mencionó, las GPUs de NVIDIA modernas utilizan una unidad de hardware especial denominada Tensor Cores que puede multiplicar matrices float16 con gran rapidez. Sin embargo, los Tensor Cores requieren que ciertas dimensiones de los tensores sean múltiplos de 8. En los siguientes ejemplos, un argumento aparece en negrita si y sólo si necesita ser múltiplo de 8 para que se puedan utilizar los Tensor Cores.

• tf.keras.layers.Dense(units=64)

• tf.keras.layers.Conv2d(filters=48, kernel_size=7, stride=3)

  • Y de forma similar para otras capas convolucionales, como tf.keras.layers.Conv3d

• tf.keras.layers.LSTM(units=64)

  • Y similar para otras RNNs, como tf.keras.layers.GRU

• tf.keras.Model.fit(epochs=2, batch_size=128)

Debe intentar utilizar Tensor Cores siempre que sea posible. Si desea obtener más información, consulte la Guía de rendimiento del deep learning de NVIDIA, donde se describen los requisitos exactos para utilizar los Tensor Cores, así como otra información de rendimiento relacionada con los Tensor Cores.

XLA

XLA es un compilador que puede aumentar aún más el rendimiento de precisión mixta, así como el rendimiento de float32 hasta un nivel menor. Consulte la Guía XLA para obtener más información.

Consejos sobre el rendimiento de la TPU en la nube

Al igual que con las GPUs, debería intentar duplicar el tamaño del lote cuando utilice TPUs en la nube, ya que los tensores bfloat16 utilizan la mitad de la memoria. Si duplicas el tamaño del lote, puede aumentar el rendimiento del entrenamiento.

Las TPU no requieren ningún otro ajuste específico de precisión mixta para obtener un rendimiento óptimo, ya que requieren el uso de XLA. Las TPU se benefician de que ciertas dimensiones sean múltiplos de $128$, pero esto se aplica tanto al tipo float32 como a la precisión mixta. Consulte la Guía de rendimiento de TPU en la nube para obtener consejos generales sobre el rendimiento de las TPU, que se aplican tanto a la precisión mixta como a los tensores float32.

Resumen

• Debe utilizar la precisión mixta si utiliza TPUs, GPUs de NVIDIA con al menos una capacidad de cálculo 7.0, o CPUs de Intel con soporte para instrucciones AMX, ya que mejorará su rendimiento hasta 3 veces.

• Puede utilizar la precisión mixta con las siguientes líneas:

  # On TPUs and CPUs, use 'mixed_bfloat16' instead
  mixed_precision.set_global_policy('mixed_float16')
  

• Si su modelo termina en softmax, asegúrese de que es float32. E independientemente de como termine su modelo, asegúrese de que la salida sea float32.

• Si utiliza un bucle de entrenamiento personalizado con mixed_float16, además de las líneas anteriores, necesita envolver su optimizador con un tf.keras.mixed_precision.LossScaleOptimizer. Luego llame a optimizer.get_scaled_loss para escalar la pérdida, y optimizer.get_unscaled_gradients para desescalar los gradientes.

• Si utiliza un bucle de entrenamiento personalizado con mixed_bfloat16, basta con establecer la global_policy mencionada anteriormente.

• Duplique el tamaño del lote de entrenamiento si no reduce la precisión de la evaluación

• En las GPU, asegúrese de que la mayoría de las dimensiones del tensor son múltiplos de $8$ para maximizar el rendimiento

Para ver un ejemplo de la precisión mixta que utiliza la API tf.keras.mixed_precision, consulte funciones y clases relacionadas con el rendimiento del entrenamiento. Revise los modelos oficiales, como Transformer, para obtener más detalles.