tf.keras смешанная шкала потери точности, по-прежнему вызывающая недостаточный уровень - PullRequest
0 голосов
/ 30 апреля 2020

Когда я использую смешанную точность для моей модели 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 я все равно не получу никакого сокращения потерь.

Буду очень признателен за любые идеи или мысли. спасибо!

Моя модель работает нормально, если я, кстати, не использую смешанную точность.

...