Conozca lo último en aprendizaje automático, IA generativa y más en el Simposio WiML 2023.

Se usó la API de Cloud Translation para traducir esta página.

Precisión mixta

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver fuente en GitHub

Descargar libreta

Descripción general

La precisión mixta es el uso de tipos de coma flotante de 16 y 32 bits en un modelo durante el entrenamiento para que se ejecute más rápido y use menos memoria. Al mantener ciertas partes del modelo en los tipos de 32 bits para la estabilidad numérica, el modelo tendrá un tiempo de paso más bajo y se entrenará igualmente en términos de métricas de evaluación, como la precisión. Esta guía describe cómo usar 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 las GPU modernas y en un 60 % en las TPU.

Hoy en día, la mayoría de los modelos utilizan el dtype float32, que ocupa 32 bits de memoria. Sin embargo, hay dos tipos de d de menor precisión, float16 y bfloat16, cada uno de los cuales ocupa 16 bits de memoria. Los aceleradores modernos pueden ejecutar operaciones más rápido en los tipos de d de 16 bits, ya que tienen hardware especializado para ejecutar cálculos de 16 bits y los tipos de d de 16 bits se pueden leer de la memoria más rápido.

Las GPU de NVIDIA pueden ejecutar operaciones en float16 más rápido que en float32, y las TPU pueden ejecutar operaciones en bfloat16 más rápido que en float32. Por lo tanto, estos tipos de d de menor precisión deben usarse siempre que sea posible en esos dispositivos. Sin embargo, las variables y algunos cálculos aún deben estar en float32 por razones numéricas para que el modelo entrene con la misma calidad. La API de precisión mixta de Keras le permite usar una combinación de float16 o bfloat16 con float32 para obtener los beneficios de rendimiento de float16/bfloat16 y los beneficios de estabilidad numérica de float32.

Configuración

import tensorflow as tf

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

Hardware compatible

Si bien la precisión mixta se ejecutará en la mayoría del hardware, solo acelerará los modelos en GPU NVIDIA y TPU en la nube recientes. Las GPU NVIDIA admiten el uso de una combinación de float16 y float32, mientras que las TPU admiten una combinación de bfloat16 y float32.

Entre las GPU NVIDIA, aquellas con capacidad de cómputo 7.0 o superior verán el mayor beneficio de rendimiento de la precisión mixta porque tienen unidades de hardware especiales, llamadas Tensor Cores, para acelerar las multiplicaciones y convoluciones de la matriz float16. Las GPU más antiguas no ofrecen ningún beneficio de rendimiento matemático por usar precisión mixta, sin embargo, los ahorros de memoria y ancho de banda pueden permitir algunas aceleraciones. Puede buscar la capacidad informática de su GPU en la página web de GPU CUDA de NVIDIA. Los ejemplos de GPU que se beneficiarán más de la precisión mixta incluyen GPU RTX, V100 y A100.

Puede verificar su tipo de GPU con lo siguiente. El comando solo existe si los controladores NVIDIA están instalados, de lo contrario, lo siguiente generará un error.

nvidia-smi -L

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

Todas las Cloud TPU admiten bfloat16.

Incluso en CPU y GPU más antiguas, donde no se espera una aceleración, las API de precisión mixta aún se pueden usar para pruebas unitarias, depuración o simplemente para probar la API. Sin embargo, en las CPU, la precisión mixta se ejecutará significativamente más lentamente.

Configuración de la política de tipo de d

Para utilizar la precisión mixta en Keras, debe crear una tf.keras.mixed_precision.Policy , normalmente conocida como política de tipo d. Las políticas de Dtype especifican las capas de dtypes en las que se ejecutarán. En esta guía, construirá una política a partir de la cadena 'mixed_float16' y la establecerá como la política global. Esto hará que las capas creadas posteriormente usen precisión mixta con una combinación de float16 y 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

Para abreviar, puede pasar directamente una cadena a set_global_policy , lo que normalmente se hace 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 creada pasando la cadena 'mixed_float16' a su constructor). Con esta política, las capas usan cálculos float16 y variables float32. Los cálculos se realizan en float16 para el rendimiento, pero las variables deben mantenerse en float32 para la 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)

Compute dtype: float16
Variable dtype: float32

Como se mencionó anteriormente, la política mixed_float16 mejorará significativamente el rendimiento en las GPU NVIDIA con capacidad de cómputo de al menos 7.0. La política se ejecutará en otras GPU y CPU, pero es posible que no mejore el rendimiento. Para las TPU, se debe usar la política mixed_bfloat16 en su lugar.

Construyendo el modelo

A continuación, comencemos a construir un modelo simple. Los modelos de juguete muy pequeños normalmente no se benefician de la precisión mixta, porque la sobrecarga del tiempo de ejecución de TensorFlow suele dominar el tiempo de ejecución, lo que hace que cualquier mejora en el rendimiento de la GPU sea insignificante. Por lo tanto, construyamos dos capas Dense grandes con 4096 unidades cada una si se usa 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)

The model will run with 4096 units on a GPU

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 configuró la política global en mixed_float16 anteriormente. Esto hará que las capas densas realicen cálculos float16 y tengan variables float32. Convierten sus entradas en float16 para realizar cálculos de float16, lo que hace que sus salidas sean float16 como resultado. Sus variables son float32 y se convertirán en float16 cuando se llame a las capas para evitar errores por discrepancias de tipo de d.

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

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)

Outputs dtype: float16

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

Esto se puede arreglar 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)

Outputs dtype: float32

Pasar dtype='float32' al constructor de la capa softmax anula la política dtype de la capa para que sea la política float32 , que realiza cálculos y mantiene las variables en float32. De manera equivalente, podría haber pasado dtype=mixed_precision.Policy('float32') ; las capas siempre convierten el argumento dtype en una política. Debido a que la capa de Activation no tiene variables, el dtype variable de la política se ignora, pero el dtype de cómputo de la política de float32 hace que softmax y la salida del modelo sean float32.

Agregar un softmax float16 en el 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 softmax hasta loss es float16 o bfloat16, pueden ocurrir 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 los cálculos de float16. Pero normalmente, esto solo 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 aún deben ser float32. Si bien no es necesario para este modelo específico, las salidas del modelo se pueden convertir en 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)

A continuación, termine y compile el modelo y genere 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

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

Se recuperan los pesos iniciales del modelo. Esto permitirá volver a entrenar desde cero cargando las pesas.

initial_weights = model.get_weights()

Entrenando al modelo con Model.fit

A continuación, entrena 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])

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

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 dedica algún tiempo a optimizar el modelo, pero luego el tiempo por paso debería estabilizarse.

Si está ejecutando esta guía en Colab, puede comparar el rendimiento de precisión mixta con float32. Para hacerlo, cambie la política de mixed_float16 a float32 en la sección "Configuración de la política dtype", luego vuelva a ejecutar todas las celdas hasta este punto. En las GPU con capacidad de cómputo 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 las GPU con una capacidad de cómputo de al menos 8,0 (GPU de amperios y superiores), es probable que no vea ninguna mejora en el rendimiento del modelo de juguete de esta guía al usar precisión mixta en comparación con float32. Esto se debe al uso de TensorFloat-32 , que automáticamente usa matemáticas de menor precisión en ciertas operaciones de float32 como tf.linalg.matmul . TensorFloat-32 brinda algunas de las ventajas de rendimiento de la precisión mixta cuando se usa float32. Sin embargo, en los modelos del mundo real, normalmente verá mejoras significativas en el rendimiento de la precisión mixta debido al ahorro de ancho de banda de la memoria y las operaciones que TensorFloat-32 no admite.

Si ejecuta una precisión mixta en una TPU, no verá una ganancia de rendimiento tan grande en comparación con la ejecución de una precisión mixta en las GPU, especialmente en las GPU anteriores a Ampere. Esto se debe a que las TPU realizan ciertas operaciones en bfloat16 bajo el capó, incluso con la política de dtype predeterminada de float32. Esto es similar a cómo las GPU Ampere usan TensorFloat-32 de forma predeterminada. En comparación con las GPU Ampere, las TPU suelen ver menos ganancias de rendimiento con precisión mixta en modelos del mundo real.

Para muchos modelos del mundo real, la precisión mixta también le 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 de juguete, ya que es probable que pueda ejecutar el modelo en cualquier tipo de d donde cada lote consta del conjunto de datos MNIST completo de 60,000 imágenes.

Escalado de pérdidas

El escalado de pérdidas es una técnica que tf.keras.Model.fit realiza automáticamente con la política mixed_float16 para evitar el subdesbordamiento numérico. Esta sección describe qué es el escalado de pérdida y la siguiente sección describe cómo usarlo con un ciclo de entrenamiento personalizado.

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 \(65504\) se desbordarán hasta el infinito y los valores por debajo \(6.0 \times 10^{-8}\) se desbordarán hasta cero. float32 y bfloat16 tienen un rango dinámico mucho más alto, por lo que el desbordamiento y el subdesbordamiento no son un problema.

Por ejemplo:

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 la práctica, el desbordamiento con float16 rara vez ocurre. Además, el desbordamiento también rara vez ocurre durante el pase hacia adelante. Sin embargo, durante el paso hacia atrás, los gradientes pueden descender hasta cero. El escalado de pérdidas es una técnica para evitar este desbordamiento.

Descripción general del escalado de pérdidas

El concepto básico de la escala de pérdida es simple: simplemente multiplique la pérdida por un número grande, digamos \(1024\), y obtendrá el valor de la escala de pérdida . Esto hará que los gradientes también se escalen en \(1024\) , lo que reducirá en gran medida la posibilidad de subdesbordamiento. Una vez que se calculan los gradientes finales, divídalos por \(1024\) para devolverlos a sus valores correctos.

El pseudocódigo para 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

Elegir una escala de pérdida puede ser complicado. Si la escala de pérdida es demasiado baja, los gradientes aún pueden descender hasta cero. Si es demasiado alto, ocurre el problema contrario: los gradientes pueden desbordarse hasta el infinito.

Para resolver esto, TensorFlow determina dinámicamente la escala de pérdida para que no tengas que elegir una manualmente. Si usa tf.keras.Model.fit , la escala de pérdida se realiza por usted para que no tenga que hacer ningún trabajo adicional. Si usa un ciclo de entrenamiento personalizado, debe usar explícitamente el envoltorio del optimizador especial tf.keras.mixed_precision.LossScaleOptimizer para usar la escala de pérdida. Esto es descrito en la siguiente sección.

Entrenamiento del modelo con un bucle de entrenamiento personalizado

Hasta ahora, ha entrenado un modelo de Keras con precisión mixta utilizando tf.keras.Model.fit . A continuación, utilizará precisión mixta con un ciclo de entrenamiento personalizado. Si aún no sabe qué es un ciclo de entrenamiento personalizado, lea primero la guía de entrenamiento personalizado .

Ejecutar un ciclo de entrenamiento personalizado con precisión mixta requiere dos cambios en lugar de ejecutarlo en float32:

Construye el modelo con precisión mixta (ya lo hiciste)
Use explícitamente la escala de pérdida si se usa mixed_float16 .

Para el paso (2), utilizará la clase tf.keras.mixed_precision.LossScaleOptimizer , que envuelve un optimizador y aplica el escalado de pérdida. De forma predeterminada, determina dinámicamente la escala de pérdida para que no tenga que elegir una. Construya un LossScaleOptimizer de la siguiente manera.

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

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

A continuación, defina el objeto de pérdida y los tf.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)

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

get_scaled_loss(loss) : Multiplica la pérdida por la escala de pérdida
get_unscaled_gradients(gradients) : toma una lista de gradientes escalados como entradas y divide cada uno por la escala de pérdida para desescalarlos

Estas funciones deben utilizarse para evitar el subdesbordamiento de los gradientes. LossScaleOptimizer.apply_gradients aplicará gradientes si ninguno de ellos tiene Inf s o NaN s. También actualizará la escala de pérdida, reduciéndola a la mitad si los gradientes tenían Inf s o NaN s y potencialmente aumentándola de lo contrario.

@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

Es probable que LossScaleOptimizer omita los primeros pasos al comienzo del entrenamiento. La escala de pérdida comienza alta para que la escala de pérdida óptima pueda determinarse rápidamente. Después de unos pocos pasos, la escala de pérdida se estabilizará y se omitirán muy pocos pasos. Este proceso ocurre automáticamente y no afecta la calidad del entrenamiento.

Ahora, defina el paso de prueba:

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

Cargue los pesos iniciales del modelo, para que pueda volver a entrenar desde cero:

model.set_weights(initial_weights)

Finalmente, ejecute el ciclo 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()))

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

Consejos de rendimiento de la GPU

Estos son algunos consejos de rendimiento cuando se usa precisión mixta en GPU.

Aumentar el tamaño de su lote

Si no afecta la calidad del modelo, intente ejecutar con el doble del tamaño del lote cuando use precisión mixta. Como los tensores float16 usan la mitad de la memoria, esto a menudo le permite duplicar el tamaño de su lote sin quedarse sin memoria. El aumento del tamaño del lote generalmente aumenta el rendimiento del entrenamiento, es decir, los elementos de entrenamiento por segundo en los que se puede ejecutar su modelo.

Asegurarse de que se utilicen núcleos de tensor de GPU

Como se mencionó anteriormente, las GPU NVIDIA modernas usan una unidad de hardware especial llamada Tensor Cores que puede multiplicar matrices float16 muy rápidamente. Sin embargo, Tensor Cores requiere que ciertas dimensiones de los tensores sean un múltiplo de 8. En los ejemplos a continuación, un argumento está en negrita si y solo si necesita ser un múltiplo de 8 para que se use Tensor Cores.

tf.keras.layers.Dense( unidades=64 )
tf.keras.layers.Conv2d( filtros=48 , kernel_size=7, zancada=3)
- Y de manera similar para otras capas convolucionales, como tf.keras.layers.Conv3d
tf.keras.layers.LSTM( unidades=64 )
- Y similar para otras RNN, como tf.keras.layers.GRU
tf.keras.Model.fit(epochs=2, batch_size=128 )

Debe intentar usar Tensor Cores cuando sea posible. Si desea obtener más información, la guía de rendimiento de aprendizaje profundo de NVIDIA describe los requisitos exactos para usar Tensor Cores, así como otra información de rendimiento relacionada con Tensor Core.

XLA

XLA es un compilador que puede aumentar aún más el rendimiento de precisión mixta, así como el rendimiento de float32 en menor medida. Consulte la guía XLA para obtener más detalles.

Consejos de rendimiento de Cloud TPU

Al igual que con las GPU, debe intentar duplicar el tamaño de su lote cuando use Cloud TPU porque los tensores bfloat16 usan la mitad de la memoria. Duplicar el tamaño del lote puede aumentar el rendimiento del entrenamiento.

Las TPU no requieren ningún otro ajuste mixto específico de precisión para obtener un rendimiento óptimo. Ya requieren el uso de XLA. Las TPU se benefician de tener ciertas dimensiones que son 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 Cloud TPU para obtener sugerencias generales de rendimiento de TPU, que se aplican a la precisión mixta, así como a los tensores float32.

Resumen

Debe usar precisión mixta si usa TPU o NVIDIA GPU con al menos capacidad de cómputo 7.0, ya que mejorará el rendimiento hasta 3 veces.

Puede usar precisión mixta con las siguientes líneas:

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

Si su modelo termina en softmax, asegúrese de que sea float32. E independientemente de en qué termine su modelo, asegúrese de que la salida sea float32.
Si usa un ciclo de entrenamiento personalizado con mixed_float16 , además de las líneas anteriores, debe envolver su optimizador con un tf.keras.mixed_precision.LossScaleOptimizer . Luego llame optimizer.get_scaled_loss para escalar la pérdida y optimizer.get_unscaled_gradients para desescalar los gradientes.
Duplicar 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 sean un múltiplo de \(8\) para maximizar el rendimiento

Para obtener más ejemplos de precisión mixta utilizando la API tf.keras.mixed_precision , consulte el repositorio de modelos oficiales . La mayoría de los modelos oficiales, como ResNet y Transformer , se ejecutarán con precisión mixta pasando --dtype=fp16 .