tf.keras.mixed_precision.LossScaleOptimizer

Оптимизатор, применяющий масштабирование потерь для предотвращения числового переполнения.

Масштабирование потерь — это техника предотвращения числового переполнения промежуточных градиентов при использовании float16. Для предотвращения переполнения потери умножаются (или «масштабируются») на определённый множитель, называемый «масштабом потерь», что приводит к масштабированию промежуточных градиентов также на этот масштаб. Конечные градиенты делятся (или «растягиваются») на масштаб потерь, чтобы вернуть их к исходному значению.

LossScaleOptimizer оборачивает другой оптимизатор и применяет к нему масштабирование потерь. По умолчанию масштаб потерь динамически обновляется со временем, поэтому вам не нужно выбирать масштаб потерь. Метод minimize автоматически масштабирует потери, не масштабирует градиенты и обновляет масштаб потерь, поэтому всё, что вам нужно сделать, — обернуть свой оптимизатор LossScaleOptimizer, если вы используете minimize. Например:

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

При использовании смешанной точности с float16, как правило, нет риска переполнения, влияющего на качество модели, если масштабирование потерь используется надлежащим образом. Более подробную информацию о том, как использовать смешанную точность, см. в руководстве по смешанной точности.

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.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 вызывает Optimizer._set_hyper только для beta_1.

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 не был задан во внутреннем оптимизаторе.

Методы

get_scaled_loss

Просмотреть исходный код

get_scaled_loss(
    loss
)

Масштабирует потерю на масштаб потерь.

Этот метод необходим только если вы вычисляете градиенты вручную, например, с помощью tf.GradientTape. В этом случае вызовите этот метод, чтобы масштабировать потерю перед передачей её в tf.GradientTape. Если вы используете LossScaleOptimizer.minimize или LossScaleOptimizer.get_gradients, масштабирование потерь применяется автоматически, и этот метод не нужен.

Если этот метод вызывается, то следует также вызвать get_unscaled_gradients. См. документацию к tf.keras.mixed_precision.LossScaleOptimizer для примера.

get_unscaled_gradients

Просмотреть исходный код

get_unscaled_gradients(
    grads
)

Рассмасштабирует градиенты на масштаб потерь.

Этот метод необходим только если вы вычисляете градиенты вручную, например, с помощью tf.GradientTape. В этом случае вызовите этот метод, чтобы расмасштабировать градиенты после их вычисления с помощью tf.GradientTape. Если вы используете LossScaleOptimizer.minimize или LossScaleOptimizer.get_gradients, масштабирование потерь применяется автоматически, и этот метод не нужен.

Если этот метод вызывается, то следует также вызвать get_scaled_loss. См. документацию к tf.keras.mixed_precision.LossScaleOptimizer для примера.