Source code for fl4health.clients.fedper_client

from flwr.common.typing import Config

from fl4health.clients.basic_client import BasicClient
from fl4health.model_bases.sequential_split_models import SequentiallySplitExchangeBaseModel
from fl4health.parameter_exchange.layer_exchanger import FixedLayerExchanger
from fl4health.parameter_exchange.parameter_exchanger_base import ParameterExchanger




class FedPerClient(BasicClient):
    """
    Client to implement the FedPer method (https://arxiv.org/abs/1912.00818).

    Trains a global feature extractor shared by all clients through FedAvg and a private classifier that is unique to
    each client. The training is nearly identical to the ``BasicClient`` with the exception that our parameter
    exchanger needs to be a fixed layer exchanger that only exchanges the feature extraction base, which relies on the
    model being of type ``SequentiallySplitExchangeBaseModel``.
    """



    def get_parameter_exchanger(self, config: Config) -> ParameterExchanger:
        assert isinstance(self.model, SequentiallySplitExchangeBaseModel), (
            "Models for FedPer must be of type SequentiallySplitExchangeBaseModel to facilitate partial weight "
            f"exchange. The current model is of type {type(self.model)}."
        )
        return FixedLayerExchanger(self.model.layers_to_exchange())