Python tf.keras.mixed_precision.LossScaleOptimizer用法及代码示例

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

继承自：Optimizer

用法

tf.keras.mixed_precision.LossScaleOptimizer(
    inner_optimizer, dynamic=True, initial_scale=None, dynamic_growth_steps=None
)

损失缩放是一种在使用 float16 时防止中间梯度中的数值下溢的技术。为了防止下溢，损失乘以(或"scaled")一个称为"loss scale"的特定因子，这导致中间梯度也被损失比例缩放。最终梯度除以损失比例(或"unscaled")，使它们恢复到原始值。

LossScaleOptimizer 包装另一个优化器并对其应用损失缩放。默认情况下，损失比例随时间动态更新，因此您不必选择损失比例。 minimize 方法会自动缩放损失、取消缩放梯度并更新损失比例，因此如果您使用 minimize ，您所要做的就是用 LossScaleOptimizer 包装优化器。例如：

opt = tf.keras.optimizers.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 的混合精度时，如果正确使用损失缩放，通常不会有下溢影响模型质量的风险。有关如何使用混合精度的更多信息，请参阅混合精度指南。

LossScaleOptimizer 偶尔会跳过对变量应用梯度，在这种情况下，可训练变量不会改变该步骤。这样做是因为动态损失比例有时会提高得太高，导致梯度溢出。通常，模型的前 2 到 15 个步骤会被跳过，因为初始损失​​规模非常高，但之后的步骤平均只会被跳过 0.05% 的时间(跳过的步骤的比例是 1 / dynamic_growth_steps )。

LossScaleOptimizer 将所有公共 Optimizer 方法委托给内部优化器。此外，在 minimize 和 get_gradients 方法中，它会缩放损失并取消缩放梯度。在方法 minimize 和 apply_gradients 中，如果任何梯度具有非有限值，它还会更新损失比例并跳过应用梯度。

超参数

可以在 LossScaleOptimizer 上访问和设置超参数，这将委托给包装的优化器。

opt = tf.keras.optimizers.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。