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。