fl4health.clients.partial_weight_exchange_client module

class PartialWeightExchangeClient(data_path, metrics, device, loss_meter_type=LossMeterType.AVERAGE, checkpoint_and_state_module=None, reporters=None, progress_bar=False, client_name=None, store_initial_model=False)

Bases: BasicClient

__init__(data_path, metrics, device, loss_meter_type=LossMeterType.AVERAGE, checkpoint_and_state_module=None, reporters=None, progress_bar=False, client_name=None, store_initial_model=False)

Client that only exchanges a subset of its parameters with the server in each communication round.

The strategy for selecting which parameters to exchange is determined by self.parameter_exchanger, which must be a subclass of PartialParameterExchanger.

Parameters:

• data_path (Path) – path to the data to be used to load the data for client-side training

• metrics (Sequence[Metric]) – Metrics to be computed based on the labels and predictions of the client model

• device (torch.device) – Device indicator for where to send the model, batches, labels etc. Often ‘cpu’ or ‘cuda’

• loss_meter_type (LossMeterType, optional) – Type of meter used to track and compute the losses over each batch. Defaults to LossMeterType.AVERAGE.

• checkpoint_and_state_module (ClientCheckpointAndStateModule | None, optional) – A module meant to handle both checkpointing and state saving. The module, and its underlying model and state checkpointing components will determine when and how to do checkpointing during client-side training. No checkpointing (state or model) is done if not provided. Defaults to None.

• reporters (Sequence[BaseReporter] | None, optional) – A sequence of FL4Health reporters which the client should send data to. Defaults to None.

• progress_bar (bool, optional) – Whether or not to display a progress bar during client training and validation. Uses tqdm. Defaults to False

• client_name (str | None, optional) – An optional client name that uniquely identifies a client. If not passed, a hash is randomly generated. Client state will use this as part of its state file name. Defaults to None.

• store_initial_model (bool) – Indicates whether the client should store a copy of the model weights at the beginning of each training round. The model copy might be required to select the subset of model parameters to be exchanged with the server, depending on the selection criterion used. Defaults to False.

get_parameter_exchanger(config)

This method configures and instantiates a PartialParameterExchanger and should be implemented by the user since there are various strategies to select parameters to be exchanged.

Parameters:

config (Config) – Configuration used to setup the weight exchanger properties for dynamic exchange

Returns:

This exchanger handles the exchange orchestration between clients and server during

federated training

Return type:

ParameterExchanger

get_parameters(config)

Determines which weights are sent back to the server for aggregation. This uses a parameter exchanger to determine parameters sent. Note that this overrides the basic client get_parameters function to send the initial model so that starting weights may be extracted and compared to current weights after local training

Parameters:

config (Config) – configuration used to setup the exchange

Returns:

The list of weights to be sent to the server from the client

Return type:

NDArrays

set_parameters(parameters, config, fitting_round)

Sets the local model parameters transferred from the server using a parameter exchanger to coordinate how parameters are set.

In the first fitting round, we assume the full model is being initialized and use the FullParameterExchanger() to set all model weights.

In other times, this approach uses a partial weight exchanger to set model weights.

Parameters:

• parameters (NDArrays) – parameters is the set of weights and their corresponding model component names, corresponding to the state dict names. These are woven together in the NDArrays object. These are unwound properly by the parameter exchanger

• config (Config) – configuration if required to control parameter exchange.

• fitting_round (bool) – Boolean that indicates whether the current federated learning round is a fitting round or an evaluation round. This is used to help determine which parameter exchange should be used for pulling parameters. A full parameter exchanger is only used if the current federated learning round is the very first fitting round. Otherwise, use a PartialParameterExchanger.

Return type:

None

setup_client(config)

Setup the components of the client necessary for client side training and parameter exchange. Mostly handled by a call to the basic client flow, but also sets up the initial model to facilitate storage of initial parameters during training

Parameters:

config (Config) – Configuration used to setup the client properly

Return type:

None