Когда я использую смешанную точность для моей модели tf.keras, потеря моей модели совсем не уменьшается. Я заметил, что мои градиенты часто либо заканчиваются на значениях "nan" или "-inf" или "inf" после использования смешанной точности.
Чтобы использовать смешанную точность, я использовал следующие фрагменты кода:
from tensorflow.keras.mixed_precision import experimental as mixed_precision
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)
self.optimizer = tf.keras.optimizers.Adam(learning_rate=self.learning_rate)
self.optimizer = mixed_precision.LossScaleOptimizer(self.optimizer, loss_scale='dynamic')
for batch_step, batch in enumerate(self.batch_iterator_train):
with tf.GradientTape() as tape:
logits, _ = self.model(batch['images'], is_training=True)
loss_value = self.loss_fn(labels=batch['labels'], logits=logits)
scaled_loss = self.optimizer.get_scaled_loss(loss_value)
scaled_grads = tape.gradient(scaled_loss, self.model.trainable_variables)
gradients = self.optimizer.get_unscaled_gradients(scaled_grads)
self.optimizer.apply_gradients(zip(gradients, self.model.trainable_variables))
Моя модель довольно большая, и она использует слои CNN, если это имеет значение. Я думал, что добавление lossScaleOptimizer решит эту проблему недостаточного заполнения градиентов, как описано https://www.tensorflow.org/guide/keras/mixed_precision, но, похоже, это не работает.
Во время отладки я заметил, что loss_scale начинается с максимального значения по умолчанию, а затем снижается до 1,0; Я предполагаю, потому что это обнаружило градиенты нана. Но тогда даже при масштабе потерь 1,0 я все равно не получу никакого сокращения потерь.
Буду очень признателен за любые идеи или мысли. спасибо!
Моя модель работает нормально, если я, кстати, не использую смешанную точность.