p2pfl.learning.frameworks.pytorch.callbacks.fedprox_callback module

FedProx Callback for PyTorch Lightning.

class p2pfl.learning.frameworks.pytorch.callbacks.fedprox_callback.FedProxCallback[source]

Bases: Callback, P2PFLCallback

PyTorch Lightning Callback to implement the FedProx algorithm.

This callback modifies the gradients before the optimizer step by adding the gradient of the proximal term: mu * (w - w_t).

static get_name()[source]

Return the name of the callback.

Return type:

str

on_before_optimizer_step(trainer, pl_module, optimizer)[source]

Add the proximal gradient term: mu * (w - w_global).

Parameters:

  • trainer (Trainer) – The trainer

  • pl_module (LightningModule) – The model

  • optimizer (Optimizer) – The optimizer

Return type:

None

on_train_start(trainer, pl_module)[source]

Store the initial model parameters and proximal coefficient.

Parameters:

  • trainer (Trainer) – The trainer

  • pl_module (LightningModule) – The model.

Return type:

None