ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more


An optimizer that applies loss scaling.

Inherits From: Optimizer

Loss scaling is a process that multiplies the loss by a multiplier called the loss scale, and divides each gradient by the same multiplier. The pseudocode for this process is:

loss = ...
loss *= loss_scale
grads = gradients(loss, vars)
grads /= loss_scale

Mathematically, loss scaling has no effect, but can help avoid numerical underflow in intermediate gradients when float16 tensors are used for mixed precision training. By multiplying the loss, each intermediate gradient will have the same multiplier applied.

The loss scale can either be a fixed constant, chosen by the user, or be dynamically determined. Dynamically determining the loss scale is convenient as a loss scale does not have to be explicitly chosen. However it reduces performance.

This optimizer wraps another optimizer and applies loss scaling to it via a LossScale. Loss scaling is applied whenever gradients are computed, such as through minimize().

use_locking Bool. If True apply use locks to prevent concurrent updates to variables.
name A non-empty string. The name to use for accumulators created for the optimizer.

ValueError If name is malformed.



View source

Apply gradients to variables.

This is the second part of minimize(). It returns an Operation that conditionally applies gradients if all gradient values are finite. Otherwise no update is performed (nor is global_step incremented).

grads_and_vars List of (gradient, variable) pairs as returned by compute_gradients().
global_step Optional Variable to increment by one after the variables have been updated.
name Optional name for the returned operation. Default to the name passed to the Optimizer constructor.

An Operation that conditionally applies the specified gradients. If global_step was not None, that operation also increments global_step.

RuntimeError If you should use _distributed_apply() instead.


View source

Compute gradients of loss for the variables in var_list.

This adjusts the dynamic range of the gradient evaluation by scaling up the loss value. The gradient values are then scaled back down by the reciprocal of the loss scale. This is useful in reduced precision training where small gradient values would otherwise underflow the representable range.

loss A Tensor containing the value to minimize or a callable taking no arguments which returns the value to minimize. When eager execution is enabled it must be a callable.
var_list Optional list or tuple of tf.Variable to update to minimize loss. Defaults to the list of variables collected in the graph under the key GraphKeys.TRAINABLE_VARIABLES.
gate_gradients How to gate the computation of gradients. Can be GATE_NONE, GATE_OP, or GATE_GRAPH.