fl4health.model_bases.fedrep_base module

class FedRepModel(base_module, head_module, flatten_features=False)

Bases: SequentiallySplitExchangeBaseModel

Implementation of the FedRep model structure: https://arxiv.org/pdf/2102.07078.pdf.

The architecture is fairly straightforward. The global module represents the first set of layers. These are learned with FedAvg. The local_prediction_head are the last layers, these are not exchanged with the server. The approach resembles FENDA, but vertical rather than parallel models. It also resembles MOON, but with partial weight exchange for weight aggregation.

freeze_base_module()

Any parameters in the base_module are fixed by setting requires_grad to False.

Return type:

None

freeze_head_module()

Any parameters in the head_module are fixed by setting requires_grad to False.

Return type:

None

unfreeze_base_module()

Any parameters in the base_module are unfrozen by setting requires_grad to True.

Return type:

None

unfreeze_head_module()

Any parameters in the head_module are unfrozen by setting requires_grad to True.

Return type:

None