混合精度训练导致 NaN 损失答案

【问题标题】：Mixed precision training leads to NaN-loss混合精度训练导致 NaN 损失
【发布时间】：2021-07-13 11:35:09
【问题描述】：

我一直在关注Mixed Precision Guide。因此，我正在设置：

keras.mixed_precision.set_global_policy(mixed_precision)

并像这样包装优化器：

if mixed_precision.startswith('mixed'):
    logger.info(f'Using LossScaleOptimizer for mixed-precision policy "{mixed_precision}"')
    optimizer = keras.mixed_precision.LossScaleOptimizer(optimizer)

我的模型有简单的Dense 层作为输出，我将其设置为“float32”

# Set dtype explicitly in last layer for mixed-precision training (float32 for numeric stability).
self.output_dense = layers.Dense(vocab_size, dtype=tf.float32)

还有一个自定义的train_step() 实现，我对此进行了修改：

with tf.GradientTape() as tape:
    model_loss = self.loss_fn(
        inputs,
        y_true=y_true,
        mask=mask
    )

    is_mixed_precision = isinstance(self.optimizer, mixed_precision.LossScaleOptimizer)

    # We always want to return the unmodified model_loss for Tensorboard
    if is_mixed_precision:
        loss = self.optimizer.get_scaled_loss(model_loss)
    else:
        loss = model_loss

    gradients = tape.gradient(loss, self.trainable_variables)

    if is_mixed_precision:
        gradients = self.optimizer.get_unscaled_gradients(gradients)

    return model_loss, gradients

但是，过了一段时间我的损失还是变成了NaN：

在外我正在确认该政策是否已被模型识别：

logger.info(f'Mixed-precision policy: {mixed_precision}')
logger.info(f'Compute dtype:          {model.compute_dtype}')
logger.info(f'Variable dtype:         {model.variable_dtype}')

keras.py:216] Mixed-precision policy: mixed_float16
keras.py:217] Compute dtype:          float16
keras.py:218] Variable dtype:         float32

但我可以说这是由于NaN 丢失..

有什么明显的我做错或错过了吗？知道如何在这里找到问题吗？

【问题讨论】：

是否可以提供一些可重现的代码，也许是mnist？
@M.Innat 嗨！感谢您的回复。幸运的是，我认为我已经找到了问题（请参阅我的答案）。 :)

标签： tensorflow keras

【解决方案1】：

经过重新思考，我想我能够找到问题所在。它位于我的自定义多头注意力层中。更具体地说，问题似乎是我使用 value.dtype.min 的掩蔽，以便对 logits 应用掩码，例如：

logits += value.dtype.min * (1.0 - mask)

有趣的是，这甚至在一开始就奏效了。想想看，你应该很有可能从一开始就遇到下溢，但我能够训练模型一段时间，直到出现NaN。

无论如何，我的解决方案是给它一些空间，所以我只需将dtype 的最小值除以二：

logits += (value.dtype.min / 2.0) * (1.0 - mask)

【讨论】：