description: An optimizer that applies loss scaling to prevent numeric underflow.

tf.keras.mixed_precision.LossScaleOptimizer

View source on GitHub

An optimizer that applies loss scaling to prevent numeric underflow.

Inherits From: Optimizer

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

Loss scaling is a technique to prevent numeric underflow in intermediate gradients when float16 is used. To prevent underflow, the loss is multiplied (or "scaled") by a certain factor called the "loss scale", which causes intermediate gradients to be scaled by the loss scale as well. The final gradients are divided (or "unscaled") by the loss scale to bring them back to their original value.

LossScaleOptimizer wraps another optimizer and applies loss scaling to it. By default, the loss scale is dynamically updated over time so you do not have to choose the loss scale. The minimize method automatically scales the loss, unscales the gradients, and updates the loss scale so all you have to do is wrap your optimizer with a LossScaleOptimizer if you use minimize. For example:

>>> 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

If a tf.GradientTape is used to compute gradients instead of minimize, you must scale the loss and gradients manually. This can be done with the LossScaleOptimizer.get_scaled_loss and LossScaleOptimizer.get_unscaled_gradients methods. For example:

>>> 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

Warning: If you forget to call get_scaled_loss or get_unscaled_gradients (or both) when using a tf.GradientTape, the model will likely converge to a worse quality. Please make sure you call each function exactly once.

When mixed precision with float16 is used, there is typically no risk of underflow affecting model quality if loss scaling is properly used. See the mixed precision guide for more information on how to use mixed precision.

LossScaleOptimizer will occasionally skip applying gradients to the variables, in which case the trainable variables will not change that step. This is done because the dynamic loss scale will sometimes be raised too high, causing overflow in the gradients. Typically, the first 2 to 15 steps of the model are skipped as the initial loss scale is very high, but afterwards steps will only be skipped on average 0.05% of the time (the fraction of steps skipped is 1 / dynamic_growth_steps).

LossScaleOptimizer delegates all public Optimizer methods to the inner optimizer. Additionally, in methods minimize and get_gradients, it scales the loss and unscales the gradients. In methodsminimizeand apply_gradients`, it additionally updates the loss scale and skips applying gradients if any gradient has a nonfinite value.

Hyperparameters

Hyperparameters can be accessed and set on the LossScaleOptimizer, which will be delegated to the wrapped optimizer.

>>> 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

However, accessing or setting non-hyperparameters is not delegated to the LossScaleOptimizer. In an Adam optimizer, beta_1 is a hyperparameter but epsilon is not, as the Adam optimizer only calls Optimizer._set_hyper on 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

In the above example, despite epsilon being set on the LossScaleOptimizer, the old epsilon value will still be used when training as epsilon was not set on the inner optimizer.

Methods

add_slot

View source

add_slot(
    var, slot_name, initializer='zeros'
)

Add a new slot variable for var.

add_weight

View source

add_weight(
    name, shape, dtype=None, initializer='zeros', trainable=None,
    synchronization=tf.VariableSynchronization.AUTO,
    aggregation=tf.compat.v1.VariableAggregation.NONE
)

apply_gradients

View source

apply_gradients(
    grads_and_vars, name=None, experimental_aggregate_gradients=True
)

Apply gradients to variables.

This is the second part of minimize(). It returns an Operation that applies gradients.

The method sums gradients from all replicas in the presence of tf.distribute.Strategy by default. You can aggregate gradients yourself by passing experimental_aggregate_gradients=False.

Example:

grads = tape.gradient(loss, vars)
grads = tf.distribute.get_replica_context().all_reduce('sum', grads)
# Processing aggregated gradients.
optimizer.apply_gradients(zip(grads, vars),
    experimental_aggregate_gradients=False)

from_config

View source

@classmethod
from_config(
    config, custom_objects=None
)

Creates an optimizer from its config.

This method is the reverse of get_config, capable of instantiating the same optimizer from the config dictionary.

get_config

View source

get_config()

Returns the config of the optimizer.

An optimizer config is a Python dictionary (serializable) containing the configuration of an optimizer. The same optimizer can be reinstantiated later (without any saved state) from this configuration.

get_gradients

View source

get_gradients(
    loss, params
)

Returns gradients of loss with respect to params.

Should be used only in legacy v1 graph mode.

get_scaled_loss

View source

get_scaled_loss(
    loss
)

Scales the loss by the loss scale.

This method is only needed if you compute gradients manually, e.g. with tf.GradientTape. In that case, call this method to scale the loss before passing the loss to tf.GradientTape. If you use LossScaleOptimizer.minimize or LossScaleOptimizer.get_gradients, loss scaling is automatically applied and this method is unneeded.

If this method is called, get_unscaled_gradients should also be called. See the tf.keras.mixed_precision.LossScaleOptimizer doc for an example.

get_slot

View source

get_slot(
    var, slot_name
)

get_slot_names

View source

get_slot_names()

A list of names for this optimizer's slots.

get_unscaled_gradients

View source

get_unscaled_gradients(
    grads
)

Unscales the gradients by the loss scale.

This method is only needed if you compute gradients manually, e.g. with tf.GradientTape. In that case, call this method to unscale the gradients after computing them with tf.GradientTape. If you use LossScaleOptimizer.minimize or LossScaleOptimizer.get_gradients, loss scaling is automatically applied and this method is unneeded.

If this method is called, get_scaled_loss should also be called. See the tf.keras.mixed_precision.LossScaleOptimizer doc for an example.

get_updates

View source

get_updates(
    loss, params
)

get_weights

View source

get_weights()

Returns the current weights of the optimizer.

The weights of an optimizer are its state (ie, variables). This function returns the weight values associated with this optimizer as a list of Numpy arrays. The first value is always the iterations count of the optimizer, followed by the optimizer's state variables in the order they were created. The returned list can in turn be used to load state into similarly parameterized optimizers.

For example, the RMSprop optimizer for this simple model returns a list of three values-- the iteration count, followed by the root-mean-square value of the kernel and bias of the single Dense layer:

>>> opt = tf.keras.optimizers.RMSprop()
>>> m = tf.keras.models.Sequential([tf.keras.layers.Dense(10)])
>>> m.compile(opt, loss='mse')
>>> data = np.arange(100).reshape(5, 20)
>>> labels = np.zeros(5)
>>> print('Training'); results = m.fit(data, labels)
Training ...
>>> len(opt.get_weights())
3

minimize

View source

minimize(
    loss, var_list, grad_loss=None, name=None, tape=None
)

Minimize loss by updating var_list.

This method simply computes gradient using tf.GradientTape and calls apply_gradients(). If you want to process the gradient before applying then call tf.GradientTape and apply_gradients() explicitly instead of using this function.

set_weights

View source

set_weights(
    weights
)

Set the weights of the optimizer.

The weights of an optimizer are its state (ie, variables). This function takes the weight values associated with this optimizer as a list of Numpy arrays. The first value is always the iterations count of the optimizer, followed by the optimizer's state variables in the order they are created. The passed values are used to set the new state of the optimizer.

For example, the RMSprop optimizer for this simple model takes a list of three values-- the iteration count, followed by the root-mean-square value of the kernel and bias of the single Dense layer:

>>> opt = tf.keras.optimizers.RMSprop()
>>> m = tf.keras.models.Sequential([tf.keras.layers.Dense(10)])
>>> m.compile(opt, loss='mse')
>>> data = np.arange(100).reshape(5, 20)
>>> labels = np.zeros(5)
>>> print('Training'); results = m.fit(data, labels)
Training ...
>>> new_weights = [np.array(10), np.ones([20, 10]), np.zeros([10])]
>>> opt.set_weights(new_weights)
>>> opt.iterations
<tf.Variable 'RMSprop/iter:0' shape=() dtype=int64, numpy=10>

variables

View source

variables()

Returns variables of this Optimizer based on the order created.