p2pfl.learning.frameworks.flax.flax_learner module

Flax Learner for P2PFL.

class p2pfl.learning.frameworks.flax.flax_learner.FlaxLearner(model, data, self_addr='unknown-node', aggregator=None)[source]

Bases: Learner

Learner for Flax models in P2PFL.

Parameters:

  • model (FlaxModel) – The FlaxModel instance.

  • data (P2PFLDataset) – The P2PFLDataset instance.

  • self_addr (str) – The address of this node.

evaluate()[source]

Evaluate the Flax model.

Return type:

Dict[str, float]

fit()[source]

Fit the model.

Return type:

P2PFLModel

property flax_model: FlaxModel

Retrieve the Flax model.

get_framework()[source]

Retrieve the learner name.

Return type:

str

Returns:

The name of the learner class.

interrupt_fit()[source]

Interrupt the fit process.

Return type:

None

train_step(state, x, y)[source]

Perform a single training step.

Return type:

Tuple[TrainState, float, float]