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

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