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 (dict[str, Any] | None) – The logs.

Return type:

None

on_train_begin(logs=None)

Initialize control variates and replace the optimizer with custom one.

Parameters:

logs (dict[str, Any] | None)

Return type:

None

on_train_end(logs=None)

Update local control variate (c_i) and compute deltas.

Parameters:

logs (dict[str, Any] | None) – The logs.

Return type:

None

set_additional_info(info)

Set additional information required for SCAFFOLD.

Parameters:

info (dict[str, Any])

Return type:

None

class p2pfl.learning.frameworks.tensorflow.callbacks.scaffold_callback.ScaffoldOptimizerWrapper(*args, **kwargs)

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.