blazefl.contrib.FedAvgThreadPoolClientTrainer

blazefl.contrib.FedAvgThreadPoolClientTrainer#

class blazefl.contrib.FedAvgThreadPoolClientTrainer(model_selector: ModelSelector, model_name: str, dataset: PartitionedDataset[FedAvgPartitionType], device: str, num_clients: int, epochs: int, batch_size: int, lr: float, seed: int, num_parallels: int)[source]#

Bases: ThreadPoolClientTrainer[FedAvgUplinkPackage, FedAvgDownlinkPackage]

__init__(model_selector: ModelSelector, model_name: str, dataset: PartitionedDataset[FedAvgPartitionType], device: str, num_clients: int, epochs: int, batch_size: int, lr: float, seed: int, num_parallels: int) → None[source]#

Methods

__init__(model_selector, model_name, ...)

get_client_device(cid)

Retrieve the device to use for processing a given client.

local_process(payload, cid_list)

Manage the parallel processing of clients using threads.

progress_fn(it)

A no-op progress function that can be overridden to provide custom progress tracking.

train(model, model_parameters, train_loader, ...)

uplink_package()

Prepare the data package to be sent from the client to the server.

worker(cid, device, payload, stop_event)

Process a single client's training task in a thread.

Attributes

num_parallels

device

device_count

cache

stop_event

progress_fn(it: list[Future[FedAvgUplinkPackage]]) → Iterable[Future[FedAvgUplinkPackage]][source]#

A no-op progress function that can be overridden to provide custom progress tracking.

Parameters:

it (list[Future[UplinkPackage]]) – A list of Future objects representing the results of client processing.

Returns:

The original iterable.

Return type:

Iterable[Future[UplinkPackage]]

train(model: Module, model_parameters: Tensor, train_loader: DataLoader, device: str, epochs: int, lr: float, stop_event: Event, cid: int) → FedAvgUplinkPackage[source]#

Prepare the data package to be sent from the client to the server.

Returns:

A list of data packages prepared for uplink transmission.

Return type:

list[UplinkPackage]

worker(cid: int, device: str, payload: FedAvgDownlinkPackage, stop_event: Event) → FedAvgUplinkPackage[source]#

Process a single client’s training task in a thread.

Parameters:

  • cid (int) – The client ID.

  • device (str) – The device to use for processing this client.

  • payload (DownlinkPackage) – The data package received from the server.

  • stop_event (threading.Event) – Event to signal stopping the worker.

Returns:

The uplink package containing the client’s results.

Return type:

UplinkPackage