p2pfl.learning.frameworks.tensorflow.callbacks.scaffold_callback module

Callback for SCAFFOLD operations (Keras).

class p2pfl.learning.frameworks.tensorflow.callbacks.scaffold_callback.SCAFFOLDCallback

Bases: Callback, P2PFLCallback

Callback for SCAFFOLD operations to use with TensorFlow Keras.

At the beginning of the training, the callback initializes control variates and substitutes the optimizer with a custom one to apply control variate adjustments. After training, it updates the local control variate (c_i) and computes the deltas.

static get_name()

Return the name of the callback.

Return type:

str

on_train_batch_end(batch, logs=None)

Increment the local step counter after each batch.

Parameters:

• batch (Any) – The batch.

• logs (Optional[Dict[str, Any]]) – The logs.

Return type:

None

on_train_begin(logs=None)

Initialize control variates and replace the optimizer with custom one.

Return type:

None

on_train_end(logs=None)

Update local control variate (c_i) and compute deltas.

Parameters:

logs (Optional[Dict[str, Any]]) – The logs.

Return type:

None

set_additional_info(info)

Set additional information required for SCAFFOLD.

Return type:

None

class p2pfl.learning.frameworks.tensorflow.callbacks.scaffold_callback.ScaffoldOptimizerWrapper(optimizer, c_i, c, eta_l)

Bases: Optimizer

Wraps an optimizer to only redefine apply_gradients, delegating other calls.

apply_gradients(grads_and_vars, name=None, **kwargs)

Apply gradients with SCAFFOLD adjustments.