tfp.optimizer.convergence_criteria.ConvergenceCriterion

Base class for stopping rules.

tfp.optimizer.convergence_criteria.ConvergenceCriterion(
    min_num_steps=None, name=None
)

A convergence criterion determines when an optimization has converged given its history of losses, gradients, and parameter values. Each criterion is responsible for propagating from step to step whatever state it needs to represent the relevant aspects of that history (for example, a moving average of previous loss values or gradients). In particular, subclasses must implement:

  • _bootstrap(loss, grads, parameters): takes the initial loss, gradients, and values of parameters, and returns a (structure of) Tensor(s) representing the initial values of any auxiliary quantities tracked by the convergence criterion.
  • _one_step(step, loss, grads, window_size, auxiliary_state): At integer step >= 1, takes the current loss, gradients, and values of parameters, along with any auxiliary state carried over from the previous step, and returns (has_converged, updated_auxiliary_state), where has_converged is a boolean Tensor, and updated_auxiliary_state is a (structure of) Tensor(s) matching auxiliary_state, containing whatever information must be propagated to the next timestep.

min_num_steps optional int Tensor minimum number of steps before stopping. If set, subclass return values of has_converged=True will be ignored until step >= min_num_steps. Default value: None.
name optional Python str name prefixed to ops created by this class.

min_num_steps

name

Methods

bootstrap

View source

bootstrap(
    loss, grads, parameters
)

Returns a structure of Tensors for the rule's state at step 0.

The shape of the Tensors specifying loss, grads, and parameters may optionally be prefixed by one or more batch dimension(s).

Args
loss float Tensor initial value of loss being optimized.
grads list of float Tensor gradients of loss wrt parameters.
parameters list of float Tensor initial values of parameters being optimized.

Returns
initial_auxiliary_state (Structure of) Tensor(s) representing the initial auxiliary state carried forward by this criterion.

one_step

View source

one_step(
    step, loss, grads, parameters, auxiliary_state
)

Updates tracked quantities for a new step, and determines if converged.

The shape of the Tensors specifying loss, grads, and parameters may optionally be prefixed by one or more batch dimension(s). In this case, the returned value has_converged will have shape equal to the broadcast batch shape of whichever of those quantities is used by this convergence criterion, and the quantities defining the convergence criterion ( min_num_steps, etc.).

Args
step integer Tensor index of the current step, where step >= 1 (on step 0, initial_state should be called instead).
loss float Tensor value of loss at the current step.
grads list of float Tensor gradients of loss wrt parameters.
parameters list of float Tensor current values of parameters being optimized.
auxiliary_state the (structure of) Tensor(s) containing state carried forward from the previous step.

Returns
has_converged boolean Tensor indicating whether the optimization has converged.
updated_auxiliary_state (Structure of) Tensor(s) representing updated quantities tracked by the convergence criterion. This should match the structure of the value returned by bootstrap.