tff.learning.templates.FinalizerProcess

A stateful process for finalization of a round of training.

Inherits From: MeasuredProcess, IterativeProcess

tff.learning.templates.FinalizerProcess(
    initialize_fn, next_fn
)

A FinalizerProcess is a tff.templates.MeasuredProcess that formalizes the type signature of initialize_fn and next_fn for the work performed by server in a learning process after aggregating model updates from clients.

The initialize_fn and next_fn must have the following type signatures:

  - initialize_fn: ( -> S@SERVER)
  - next_fn: (<S@SERVER,
               A@SERVER,
               B@SERVER>
              ->
              <state=S@SERVER,
               result=A@SERVER,
               measurements=M@SERVER>)

FinalizerProcess requires next_fn with a second and a third input argument, which are both placed at SERVER. The second type A represents the current server parameters to be updated, while the third type B represents an update to the parameter A, and often matches the type of B.

The result field of the returned tff.templates.MeasuredProcessOutput must be placed at SERVER, be of type matching that of second input argument (B) and represents the updated ("finalized") model parameters.

initialize_fn A no-arg tff.Computation that returns the initial state of the measured process. Let the type of this state be called S.
next_fn A tff.Computation that represents the iterated function. The first or only argument must match the state type S. The return value must be a MeasuredProcessOutput whose state member matches the state type S.
next_is_multi_arg An optional boolean indicating that next_fn will receive more than just the state argument (if True) or only the state argument (if False). This parameter is primarily used to provide better error messages.

TypeError If initialize_fn and next_fn are not instances of tff.Computation.
TemplateInitFnParamNotEmptyError If initialize_fn has any input arguments.
TemplateStateNotAssignableError If the state returned by either initialize_fn or next_fn is not assignable to the first input argument of next_fn.
TemplateNotMeasuredProcessOutputError If next_fn does not return a MeasuredProcessOutput.

initialize A no-arg tff.Computation that returns the initial state.
next A tff.Computation that runs one iteration of the process.

Its first argument should always be the current state (originally produced by tff.templates.MeasuredProcess.initialize), and the return type must be a tff.templates.MeasuredProcessOutput.

state_type The tff.Type of the state of the process.