p2pfl.learning.frameworks.flax.flax_learner module

Flax Learner for P2PFL.

class p2pfl.learning.frameworks.flax.flax_learner.FlaxLearner(model=None, data=None, aggregator=None)[source]

Bases: Learner

Learner for Flax models in P2PFL.

Parameters:

  • model (Optional[P2PFLModel]) – The FlaxModel instance.

  • data (Optional[P2PFLDataset]) – The P2PFLDataset instance.

  • self_addr – The address of this node.

evaluate(*args: Any, **kwargs: Any) → Any
Return type:

Any

fit(*args: Any, **kwargs: Any) → Any
Return type:

Any

property flax_model: FlaxModel

Retrieve the Flax model.

get_framework(*args: Any, **kwargs: Any) → Any
Return type:

Any

interrupt_fit(*args: Any, **kwargs: Any) → Any
Return type:

Any

train_step(*args: Any, **kwargs: Any) → Any
Return type:

Any