LossScaleOptimizer

[源代码]

BaseLossScaleOptimizer 类

tf_keras.mixed_precision.LossScaleOptimizer()

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

损失缩放是一种在使用 float16 时防止中间梯度发生数值下溢的技术。为了防止下溢，损失乘以（或“缩放”）一个称为“损失比例”的特定因子，这也会导致中间梯度按损失比例进行缩放。最终梯度除以（或“取消缩放”）损失比例，将其恢复到原始值。

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

>>> 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 未在内部优化器上设置。