LossScaleOptimizer

[源代码]

BaseLossScaleOptimizer 类

tf_keras.mixed_precision.LossScaleOptimizer()

一个应用损失缩放以防止数值下溢的优化器。

损失缩放是一种在 Float16 被使用时，防止中间梯度数值下溢的技术。为了防止下溢，损失会乘以（或“缩放”）一个称为“损失尺度”的因子，该因子也会导致中间梯度被损失尺度缩放。最终的梯度会被除以（或“反缩放”）损失尺度，以恢复其原始值。

LossScaleOptimizer 包装另一个优化器并对其应用损失缩放。默认情况下，损失尺度会随时间动态更新，因此您无需选择损失尺度。minimize 方法会自动缩放损失，反缩放梯度，并更新损失尺度，因此如果您使用 minimize，只需用 LossScaleOptimizer 包装您的优化器即可。例如：

>>> opt = tf.keras.optimizers.experimental.SGD(0.25)
>>> opt = tf.keras.mixed_precision.LossScaleOptimizer(opt)
>>> var = tf.Variable(1.)
>>> loss_fn = lambda: var ** 2
>>> # 'minimize' applies loss scaling and updates the loss sale.
>>> opt.minimize(loss_fn, var_list=[var])
>>> var.numpy()
0.5

如果使用 tf.GradientTape 来计算梯度而不是 minimize，您必须手动缩放损失和梯度。这可以通过 LossScaleOptimizer.get_scaled_loss 和 LossScaleOptimizer.get_unscaled_gradients 方法来完成。例如：

>>> with tf.GradientTape() as tape:
...   loss = loss_fn()
...   scaled_loss = opt.get_scaled_loss(loss)
>>> scaled_grad = tape.gradient(scaled_loss, var)
>>> (grad,) = opt.get_unscaled_gradients([scaled_grad])
>>> opt.apply_gradients([(grad, var)])  # Loss scale is updated here
>>> var.numpy()
0.25

警告：如果您在使用 tf.GradientTape 时忘记调用 get_scaled_loss 或 get_unscaled_gradients（或两者都调用），模型很可能收敛到较差的质量。请确保每个函数都恰好调用一次。

当使用 Float16 进行混合精度训练时，如果正确使用了损失缩放，通常不会有影响模型质量的下溢风险。有关如何使用混合精度的更多信息，请参阅 混合精度指南。

参数

• inner_optimizer：要包装的 tf.keras.optimizers.Optimizer 或 tf.keras.optimizers.experimental.Optimizer 实例。
• dynamic：布尔值，指示是否使用动态损失缩放。如果为 True，损失尺度将使用一种算法随时间动态更新，该算法可将损失尺度保持在其近似最优值。如果为 False，则使用单个固定的损失尺度，并且必须指定 initial_scale，它将用作损失尺度。建议保持为 True，因为选择固定的损失尺度可能很棘手。目前，与固定损失缩放相比，动态损失缩放会带来一点性能开销。默认为 True。
• initial_scale：初始损失尺度。如果 dynamic 为 True，则默认为 2 ** 15。如果 dynamic 为 False，则必须指定此值，并且它将作为唯一的损失尺度，因为损失尺度不会随时间变化。使用动态损失缩放时，最好将其设置为一个非常大的数字，因为过高的损失尺度比过低的损失尺度降低得更快。
• dynamic_growth_steps：在使用动态损失缩放时，每 dynamic_growth_steps 步出现非无穷大梯度时，损失尺度会翻倍。如果遇到非无穷大梯度，计数将重置为零，该步的梯度将被跳过，并且损失尺度会减半。可以通过 LossScaleOptimizer.dynamic_counter 查询计数。仅当 dynamic 为 True 时才能指定此参数。默认为 2000。

LossScaleOptimizer 会偶尔跳过将梯度应用于变量，在这种情况下，该步的训练变量不会发生变化。这是因为动态损失尺度有时会过高，导致梯度溢出。通常，模型的初始 2 到 15 步会被跳过，因为初始损失尺度非常高，但之后，步长平均只有 0.05% 的几率被跳过（跳过的步长比例为 1 / dynamic_growth_steps）。

LossScaleOptimizer 将所有公共 Optimizer 方法委托给内部优化器。此外，在 minimize 和 get_gradients 方法中，它会缩放损失并反缩放梯度。在 minimize 和 apply_gradients 方法中，它还会更新损失尺度，并在任何梯度具有非无穷大值时跳过应用梯度。

超参数

如果包装的是 tf.keras.optimizers.Optimizer，则可以在 LossScaleOptimizer 上访问和设置超参数，这些超参数将被委托给被包装的优化器。

>>> opt = tf.keras.optimizers.legacy.Adam(beta_1=0.8, epsilon=1e-5)
>>> opt = tf.keras.mixed_precision.LossScaleOptimizer(opt)
>>> opt.beta_1  # Equivalent to `opt.inner_optimizer.beta_1`
0.8
>>> opt.beta_1 = 0.7  # Equivalent to `opt.inner_optimizer.beta_1 = 0.7`
>>> opt.beta_1
0.7
>>> opt.inner_optimizer.beta_1
0.7

然而，访问或设置非超参数的属性不会被委托给 LossScaleOptimizer。在 Adam 优化器中，beta_1 是一个超参数，但 epsilon 不是，因为 Adam 优化器只对 beta_1 调用 Optimizer._set_hyper。

>>> opt.inner_optimizer.epsilon
1e-5
>>> opt.epsilon
Traceback (most recent call last):
...
AttributeError: 'LossScaleOptimizer' object has no attribute 'epsilon'
>>> opt.epsilon = 1e-4  # This does NOT set epsilon on `opt.inner_optimizer`
>>> opt.inner_optimizer.epsilon
>>> 1e-5

在上面的示例中，尽管 epsilon 在 LossScaleOptimizer 上被设置了，但当训练时仍会使用旧的 epsilon 值，因为 epsilon 没有在内部优化器上设置。