"""
fl_sim.nodes
=============
:class:`Node` s are the core of the simulation framework.
:class:`Node` has two subclasses: :class:`Server` and :class:`Client`.
The :class:`Server` class is the base class for all servers,
which acts as the coordinator of the training process, as well as maintainer of status variables.
The :class:`Client` class is the base class for all clients.
.. contents::
:depth: 2
:local:
:backlinks: top
.. currentmodule:: fl_sim.nodes
Base Node class
---------------
.. autosummary::
:toctree: generated/
Node
Server classes
--------------
.. autosummary::
:toctree: generated/
Server
ServerConfig
Client classes
--------------
.. autosummary::
:toctree: generated/
Client
ClientConfig
"""
import json
import os
import types
import warnings
from abc import ABC, abstractmethod
from collections import defaultdict
from copy import deepcopy
from itertools import repeat
from numbers import Number
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union
import numpy as np
import torch
import torch.nn as nn
import yaml
from bib_lookup import CitationMixin
from torch import Tensor
from torch.nn.parameter import Parameter
from torch.optim import SGD
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader
from torch_ecg.cfg import CFG
from torch_ecg.utils import ReprMixin, add_docstring, get_kwargs
from tqdm.auto import tqdm
from .data_processing.fed_dataset import FedDataset
from .models import reset_parameters
from .optimizers import get_optimizer
from .utils.loggers import LoggerManager
from .utils.misc import default_dict_to_dict, get_scheduler, set_seed
from .utils.torch_compat import torch_norm
__all__ = [
"Server",
"Client",
"ServerConfig",
"ClientConfig",
"ClientMessage",
]
[docs]class ServerConfig(ReprMixin):
"""Configs for the Server.
Parameters
----------
algorithm : str
The algorithm name.
num_iters : int
The number of (outer) iterations.
num_clients : int
The number of clients.
clients_sample_ratio : float
The ratio of clients to sample for each iteration.
log_dir : str or pathlib.Path, optional
The log directory.
If not specified, will use the default log directory.
If not absolute, will be relative to the default log directory.
txt_logger : bool, default True
Whether to use txt logger.
json_logger : bool, default True
Whether to use json logger.
eval_every : int, default 1
The number of iterations to evaluate the model.
visiable_gpus : Sequence[int], optional
Visable GPU IDs for allocating devices for clients.
Defaults to use all GPUs if available.
extra_observes : List[str], optional
Extra attributes to observe during training.
seed : int, default 0
The random seed.
tag : str, optional
The tag of the experiment.
verbose : int, default 1
The verbosity level.
gpu_proportion : float, default 0.2
The proportion of clients to use GPU.
Used to similate the system heterogeneity of the clients.
Not used in the current version, reserved for future use.
**kwargs : dict, optional
The other arguments,
will be set as attributes of the class.
"""
__name__ = "ServerConfig"
def __init__(
self,
algorithm: str,
num_iters: int,
num_clients: int,
clients_sample_ratio: float,
log_dir: Optional[Union[str, Path]] = None,
txt_logger: bool = True,
json_logger: bool = True,
eval_every: int = 1,
visiable_gpus: Optional[Sequence[int]] = None,
extra_observes: Optional[List[str]] = None,
seed: int = 0,
tag: Optional[str] = None,
verbose: int = 1,
gpu_proportion: float = 0.2,
**kwargs: Any,
) -> None:
self.algorithm = algorithm
self.num_iters = num_iters
self.num_clients = num_clients
self.clients_sample_ratio = clients_sample_ratio
self.log_dir = log_dir
self.txt_logger = txt_logger
self.json_logger = json_logger
self.eval_every = eval_every
self.visiable_gpus = visiable_gpus
default_gpus = list(range(torch.cuda.device_count()))
if self.visiable_gpus is None:
self.visiable_gpus = default_gpus
if not set(self.visiable_gpus).issubset(set(default_gpus)):
warnings.warn(f"GPU(s) {set(self.visiable_gpus) - set(default_gpus)} " "are not available.")
self.visiable_gpus = [item for item in self.visiable_gpus if item in default_gpus]
self.extra_observes = extra_observes or []
self.seed = seed
self.tag = tag
self.verbose = verbose
self.gpu_proportion = gpu_proportion
for k, v in kwargs.items():
setattr(self, k, v)
set_seed(self.seed)
[docs]class ClientConfig(ReprMixin):
"""Configs for the Client.
Parameters
----------
algorithm : str
The algorithm name.
optimizer : str
The name of the optimizer to solve the local (inner) problem.
batch_size : int
The batch size.
num_epochs : int
The number of epochs.
lr : float
The learning rate.
scheduler : dict, optional
The scheduler config.
None for no scheduler, using constant learning rate.
extra_observes : List[str], optional
Extra attributes to observe during training,
which would be recorded in evaluated metrics,
sent to the server, and written to the log file.
verbose : int, default 1
The verbosity level.
latency : float, default 0.0
The latency of the client.
Not used in the current version, reserved for future use.
**kwargs : dict, optional
The other arguments,
will be set as attributes of the class.
"""
__name__ = "ClientConfig"
def __init__(
self,
algorithm: str,
optimizer: str,
batch_size: int,
num_epochs: int,
lr: float,
scheduler: Optional[dict] = None,
extra_observes: Optional[List[str]] = None,
verbose: int = 1,
latency: float = 0.0,
**kwargs: Any,
) -> None:
self.algorithm = algorithm
self.optimizer = optimizer
self.batch_size = batch_size
self.num_epochs = num_epochs
self.lr = lr
self.scheduler = scheduler or {"name": "none"}
self.extra_observes = extra_observes or []
self.verbose = verbose
self.latency = latency
for k, v in kwargs.items():
setattr(self, k, v)
[docs]class Node(ReprMixin, ABC):
"""An abstract base class for the server and client nodes."""
__name__ = "Node"
[docs] @abstractmethod
def communicate(self, target: "Node") -> None:
"""Communicate with the target node.
The current node communicates model parameters, gradients, etc. to `target` node.
For example, a client node communicates its local model parameters to server node via
.. code-block:: python
target._received_messages.append(
ClientMessage(
{
"client_id": self.client_id,
"parameters": self.get_detached_model_parameters(),
"train_samples": self.config.num_epochs * self.config.num_steps * self.config.batch_size,
"metrics": self._metrics,
}
)
)
For a server node, global model parameters are communicated to clients via
.. code-block:: python
target._received_messages = {"parameters": self.get_detached_model_parameters()}
"""
raise NotImplementedError
[docs] @abstractmethod
def update(self) -> None:
"""Update model parameters, gradients, etc.
according to `self._reveived_messages`.
"""
raise NotImplementedError
def _post_init(self) -> None:
"""Check if all required field in the config are set."""
assert all(
[hasattr(self.config, k) for k in self.required_config_fields]
), f"missing required config fields: {list(set(self.required_config_fields) - set(self.config.__dict__))}"
@property
@abstractmethod
def required_config_fields(self) -> List[str]:
"""The list of required fields in the config."""
raise NotImplementedError
@property
@abstractmethod
def is_convergent(self) -> bool:
"""Whether the training process on the node is convergent."""
raise NotImplementedError
[docs] def get_detached_model_parameters(self) -> List[Tensor]:
"""Get the detached model parameters."""
return [p.detach().clone() for p in self.model.parameters()]
[docs] def compute_gradients(
self,
at: Optional[Sequence[Tensor]] = None,
dataloader: Optional[DataLoader] = None,
) -> List[Tensor]:
"""Compute the gradients of the model on the node.
The gradients are computed on the model parameters `at` or the current model parameters,
as the average of the gradients on the mini-batches from `dataloader` or `self.train_loader`.
Parameters
----------
at : list of torch.Tensor, optional
The model parameters to compute the gradients.
None for the current model parameters.
dataloader : torch.utils.data.DataLoader, optional
The dataloader to compute the gradients.
None for `self.train_loader`.
"""
prev_training = self.model.training
prev_model_state_dict = self.model.state_dict()
if at is not None:
self.set_parameters(at)
if dataloader is None:
assert hasattr(self, "train_loader") and self.train_loader is not None, "train_loader is not set"
dataloader = self.train_loader
assert len(dataloader) > 0, "empty dataloader"
self.model.train()
self.optimizer.zero_grad()
total_samples = len(dataloader.dataset)
for X, y in dataloader:
X, y = X.to(self.device), y.to(self.device)
y_pred = self.model(X)
loss = self.criterion(y_pred, y) * len(X) / total_samples
# gradients are accumulated after each backward pass
# one should **NOT** call `optimizer.zero_grad()` before the backward pass
# and should **NOT** call `optimizer.step()` after the backward pass
# since the model parameters should not be updated
# otherwise, the result is not the gradients **at** ``at``
loss.backward()
mini_batch_grads = [p.grad.detach().clone() for p in self.model.parameters()]
# set the model state back and clear the gradients
self.model.load_state_dict(prev_model_state_dict)
self.optimizer.zero_grad()
self.model.train(prev_training)
# garbage collection
del prev_model_state_dict
return mini_batch_grads
[docs] def get_gradients(
self,
norm: Optional[Union[str, int, float]] = None,
model: Optional[torch.nn.Module] = None,
) -> Union[float, List[Tensor]]:
"""Get the gradients or norm of the gradients
of the model on the node.
Parameters
----------
norm : str or int or float, optional
The norm of the gradients to compute.
None for the raw gradients (list of tensors).
Refer to :func:`torch.linalg.norm` for more details.
model : torch.nn.Module, optional
The model to get the gradients,
default to `self.model`.
Returns
-------
float or List[torch.Tensor]
The gradients or norm of the gradients.
"""
if model is None:
model = self.model
grads = []
for param in model.parameters():
if param.grad is None:
grads.append(torch.zeros_like(param.data))
else:
grads.append(param.grad.data)
if norm is not None:
if len(grads) == 0:
grads = 0.0
warnings.warn("No gradients available. Set to 0.0 by default.", RuntimeWarning)
else:
grads = torch_norm(torch.cat([grad.view(-1) for grad in grads]), norm).item()
return grads
[docs] @staticmethod
def get_norm(
tensor: Union[
Number,
Tensor,
np.ndarray,
Parameter,
types.GeneratorType,
Sequence[Union[np.ndarray, Tensor, Parameter, types.GeneratorType]],
],
norm: Union[str, int, float] = "fro",
) -> float:
"""Get the norm of a tensor.
Parameters
----------
tensor : torch.Tensor or np.ndarray or torch.nn.parameter.Parameter or generator or list
The tensor (array, parameter, etc.) to compute the norm.
norm : str or int or float, default "fro"
The norm to compute.
Refer to :func:`torch.linalg.norm` for more details.
Returns
-------
float
The norm of the tensor.
"""
if tensor is None:
return float("nan")
if isinstance(tensor, Number):
return tensor
if isinstance(tensor, Tensor):
tensor = [tensor]
elif isinstance(tensor, np.ndarray):
tensor = [torch.from_numpy(tensor)]
elif isinstance(tensor, Parameter):
tensor = [tensor.data]
elif isinstance(tensor, (list, tuple)):
tensor = [torch.from_numpy(t) if isinstance(t, np.ndarray) else t.detach().clone() for t in tensor]
elif isinstance(tensor, types.GeneratorType):
return Node.get_norm(list(tensor), norm)
else:
raise TypeError(f"Unsupported type: {type(tensor)}")
return torch_norm(torch.cat([t.view(-1) for t in tensor]), norm).item()
[docs] def set_parameters(self, params: Iterable[Parameter], model: Optional[torch.nn.Module] = None) -> None:
"""Set the parameters of the model on the node.
Parameters
----------
params : Iterable[torch.nn.parameter.Parameter]
The parameters to set.
model : torch.nn.Module, optional
The model to set the parameters,
default to `self.model`.
Returns
-------
None
"""
if model is None:
model = self.model
for node_param, param in zip(model.parameters(), params):
node_param.data = param.data.detach().clone().to(self.device)
[docs] @staticmethod
def aggregate_results_from_json_log(d: Union[dict, str, Path], part: str = "val", metric: str = "acc") -> np.ndarray:
"""Aggregate the federated results from csv log.
Parameters
----------
d : dict or str or pathlib.Path
The dict of the json/yaml log,
or the path to the json/yaml log file.
part : str, default "train"
The part of the log to aggregate.
metric : str, default "acc"
The metric to aggregate.
Returns
-------
numpy.ndarray
The aggregated results (curve).
NOTE
----
The parameter `d` should be a dict similar to the following structure:
.. dropdown::
:animate: fade-in-slide-down
.. code-block:: json
{
"train": {
"client0": [
{
"epoch": 1,
"step": 1,
"time": "2020-01-01 00:00:00",
"loss": 0.1,
"acc": 0.2,
"top3_acc": 0.3,
"top5_acc": 0.4,
"num_samples": 100
}
]
},
"val": {
"client0": [
{
"epoch": 1,
"step": 1,
"time": "2020-01-01 00:00:00",
"loss": 0.1,
"acc": 0.2,
"top3_acc": 0.3,
"top5_acc": 0.4,
"num_samples": 100
}
]
}
}
"""
if isinstance(d, (str, Path)):
d = Path(d)
if d.suffix == ".json":
d = json.loads(d.read_text())
elif d.suffix in [".yaml", ".yml"]:
d = yaml.safe_load(d.read_text())
else:
raise ValueError(f"unsupported file type: {d.suffix}")
epochs = list(sorted(np.unique([item["epoch"] for _, v in d[part].items() for item in v])))
metric_curve = [[] for _ in range(len(epochs))]
num_samples = [0 for _ in range(len(epochs))]
for _, v in tqdm(
d[part].items(),
mininterval=1,
desc="Aggregating results",
total=len(d[part]),
unit="client",
leave=False,
disable=int(os.environ.get("FLSIM_VERBOSE", "1")) < 1,
):
for item in v:
idx = epochs.index(item["epoch"])
metric_curve[idx].append(item[metric] * item["num_samples"])
num_samples[idx] += item["num_samples"]
return np.array([sum(v) / num_samples[i] for i, v in enumerate(metric_curve)])
[docs]class Server(Node, CitationMixin):
"""The class to simulate the server node.
The server node is responsible for communicating with clients,
and perform the aggregation of the local model parameters (and/or gradients),
and update the global model parameters.
Parameters
----------
model : torch.nn.Module
The model to be trained (optimized).
dataset : FedDataset
The dataset to be used for training.
config : ServerConfig
The configs for the server.
client_config : ClientConfig
The configs for the clients.
lazy : bool, default False
Whether to use lazy initialization
for the client nodes. This is useful
when one wants to do centralized training
for verification.
TODO
----
1. Run clients training in parallel.
2. Use the attribute `_is_convergent` to control the termination of the training.
This perhaps can be achieved by comparing part of the items in `self._cached_models`
"""
__name__ = "Server"
def __init__(
self,
model: nn.Module,
dataset: FedDataset,
config: ServerConfig,
client_config: ClientConfig,
lazy: bool = False,
) -> None:
self.model = model
self.dataset = dataset
self.criterion = deepcopy(dataset.criterion)
assert isinstance(config, self.config_cls["server"]), (
f"(server) config should be an instance of " f"{self.config_cls['server']}, but got {type(config)}."
)
self.config = config
if not hasattr(self.config, "verbose"):
self.config.verbose = get_kwargs(ServerConfig)["verbose"]
warnings.warn(
"The `verbose` attribute is not found in the config, " f"set it to the default value {self.config.verbose}.",
RuntimeWarning,
)
if self.config.num_clients is None:
self.config.num_clients = self.dataset.DEFAULT_TRAIN_CLIENTS_NUM
self.device = torch.device("cpu")
assert isinstance(client_config, self.config_cls["client"]), (
f"client_config should be an instance of " f"{self.config_cls['client']}, but got {type(client_config)}."
)
self._client_config = client_config
if not hasattr(self._client_config, "verbose"):
# self._client_config.verbose = get_kwargs(ClientConfig)["verbose"]
self._client_config.verbose = self.config.verbose # set to server's verbose
warnings.warn(
"The `verbose` attribute is not found in the client_config, "
f"set it to the default value {self._client_config.verbose}.",
RuntimeWarning,
)
logger_config = dict(
log_dir=self.config.log_dir,
log_suffix=self.config.tag,
txt_logger=self.config.txt_logger,
json_logger=self.config.json_logger,
algorithm=self.config.algorithm,
model=self.model.__class__.__name__,
dataset=self.dataset.__class__.__name__,
verbose=self.config.verbose,
)
self._logger_manager = LoggerManager.from_config(logger_config)
# set batch_size, in case of centralized training
setattr(self.config, "batch_size", client_config.batch_size)
# echo the configs to stdout and txt log file
self._logger_manager.log_message(f"Server Configs:\n{str(self.config)}")
self._logger_manager.log_message(f"Client Configs:\n{str(self._client_config)}")
self._received_messages = []
self._num_communications = 0
self.n_iter = 0
self._metrics = {} # aggregated metrics
self._cached_metrics = [] # container for caching aggregated metrics
self._is_convergent = False # status variable for convergence
# flag for completing the experiment
# which will be checked before calling
# `self.train`, `self.train_centralized`, `self.train_federated`
self._complete_experiment = False
self._clients = None
if not lazy:
self._setup_clients(self.dataset, self._client_config)
self._post_init()
# checks that the client has all the required attributes in config.extra_observes
for attr in self.config.extra_observes:
assert hasattr(self, attr), f"{self.__name__} should have attribute {attr} for extra observes."
def _setup_clients(
self,
dataset: Optional[FedDataset] = None,
client_config: Optional[ClientConfig] = None,
force: bool = False,
) -> None:
"""Setup the clients.
Parameters
----------
dataset : FedDataset, optional
The dataset to be used for training the local models,
defaults to `self.dataset`.
client_config : ClientConfig, optional
The configs for the clients,
defaults to `self._client_config`.
force : bool, default False
Whether to force setup the clients.
If set to True, the clients will be setup
even if they have been setup before.
Returns
-------
None
"""
if self._clients is not None and not force:
return
elif force:
# reset the clients
del self._clients
self._logger_manager.log_message("Setup clients...")
dataset = dataset or self.dataset
client_config = client_config or self._client_config
self._clients = [
self.client_cls(client_id, device, deepcopy(self.model), dataset, client_config)
for client_id, device in tqdm(
zip(range(self.config.num_clients), self._allocate_devices()),
desc="Allocating devices",
total=self.config.num_clients,
unit="client",
mininterval=1,
leave=False,
)
]
def _allocate_devices(self) -> List[torch.device]:
"""Allocate devices for clients, can be used in :meth:`_setup_clients`."""
self._logger_manager.log_message("Allocate devices...")
# num_gpus = torch.cuda.device_count()
num_gpus = len(self.config.visiable_gpus)
if num_gpus == 0:
return list(repeat(torch.device("cpu"), self.config.num_clients))
return [torch.device(f"cuda:{self.config.visiable_gpus[i%num_gpus]}") for i in range(self.config.num_clients)]
def _sample_clients(
self,
subset: Optional[Sequence[int]] = None,
clients_sample_ratio: Optional[float] = None,
) -> List[int]:
"""Sample clients for each iteration.
Parameters
----------
subset : Sequence[int], optional
The subset of clients to be sampled from,
defaults to `None`, which means all clients.
clients_sample_ratio : float, optional
The ratio of clients to be sampled,
defaults to `None`, which means `self.config.clients_sample_ratio`.
Returns
-------
List[int]
The list of selected client ids.
"""
if subset is None:
subset = range(self.config.num_clients)
if clients_sample_ratio is None:
clients_sample_ratio = self.config.clients_sample_ratio
k = min(len(subset), max(1, round(clients_sample_ratio * len(subset))))
# return random.sample(subset, k=k)
# random.sample does not support numpy array input
return np.random.choice(subset, k, replace=False).tolist()
def _communicate(self, target: "Client") -> None:
"""Broadcast to target client, and maintain state variables."""
self.communicate(target)
self._num_communications += 1
def _update(self) -> None:
"""Server update, and clear cached messages from clients of the previous iteration."""
self._logger_manager.log_message("Server update...")
if len(self._received_messages) == 0:
warnings.warn(
"No message received from the clients, unable to update server model.",
RuntimeWarning,
)
return
assert all([isinstance(m, ClientMessage) for m in self._received_messages]), (
"received messages must be of type `ClientMessage`, "
f"but got {[type(m) for m in self._received_messages if not isinstance(m, ClientMessage)]}"
)
self.update()
# free the memory of the received messages
del self._received_messages
self._received_messages = [] # clear messages received in the previous iteration
self._logger_manager.log_message("Server update finished...")
[docs] def train(self, mode: str = "federated", extra_configs: Optional[dict] = None) -> None:
"""The main training loop.
Parameters
----------
mode : {"federated", "centralized", "local"}, optional
The mode of training, by default "federated", case-insensitive.
extra_configs : dict, optional
The extra configs for the training `mode`.
Returns
-------
None
"""
if self._complete_experiment:
# reset before training if a previous experiment is completed
self._reset(reset_clients=(mode.lower() != "centralized"))
self._complete_experiment = False
if mode.lower() == "federated":
self.train_federated(extra_configs)
elif mode.lower() == "centralized":
self.train_centralized(extra_configs)
elif mode.lower() == "local":
self.train_local(extra_configs)
else:
raise ValueError(f"mode {mode} is not supported")
self._complete_experiment = True
[docs] def train_centralized(self, extra_configs: Optional[dict] = None) -> None:
"""Centralized training, conducted only on the server node.
This is used as a baseline for comparison.
Parameters
----------
extra_configs : dict, optional
The extra configs for centralized training.
Returns
-------
None
"""
if self._complete_experiment:
# reset before training if a previous experiment is completed
# no need to reset clients for centralized training
self._reset(reset_clients=False)
self._logger_manager.log_message("Training centralized...")
extra_configs = CFG(extra_configs or {})
batch_size = extra_configs.get("batch_size", self.config.batch_size)
train_loader, val_loader = self.dataset.get_dataloader(batch_size, batch_size, None)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.train()
self.model.to(device)
criterion = deepcopy(self.dataset.criterion)
lr = extra_configs.get("lr", 1e-2)
optimizer = extra_configs.get("optimizer", SGD(self.model.parameters(), lr))
scheduler = extra_configs.get("scheduler", LambdaLR(optimizer, lambda epoch: 1 / (0.01 * epoch + 1)))
self._complete_experiment = False
epoch_losses = []
self.n_iter, global_step = 0, 0
for self.n_iter in range(self.config.num_iters):
with tqdm(
total=len(train_loader.dataset),
desc=f"Epoch {self.n_iter+1}/{self.config.num_iters}",
unit="sample",
mininterval=1 if torch.cuda.is_available() else 10,
) as pbar:
epoch_loss = []
batch_losses = []
for data, target in train_loader:
data, target = data.to(device), target.to(device)
output = self.model(data)
loss = criterion(output, target)
optimizer.zero_grad()
loss.backward()
batch_losses.append(loss.item())
optimizer.step()
global_step += 1
if self.config.verbose >= 2:
pbar.set_postfix(
**{
"loss (batch)": loss.item(),
"lr": scheduler.get_last_lr()[0],
}
)
pbar.update(data.shape[0])
# free memory
del data, target, output, loss
epoch_loss.append(sum(batch_losses) / len(batch_losses))
if (self.n_iter + 1) % self.config.eval_every == 0:
self._logger_manager.log_message("Evaluating...")
metrics = self.evaluate_centralized(val_loader)
self._logger_manager.log_metrics(
None,
metrics,
step=global_step,
epoch=self.n_iter + 1,
part="val",
)
metrics = self.evaluate_centralized(train_loader)
self._logger_manager.log_metrics(
None,
metrics,
step=global_step,
epoch=self.n_iter + 1,
part="train",
)
scheduler.step()
self.model.to(self.device) # move to the original device
self._logger_manager.log_message("Centralized training finished...")
self._logger_manager.flush()
# self._logger_manager.reset()
self._complete_experiment = True
[docs] def train_federated(self, extra_configs: Optional[dict] = None) -> None:
"""Federated (distributed) training,
conducted on the clients and the server.
Parameters
----------
extra_configs : dict, optional
The extra configs for federated training.
Returns
-------
None
TODO
----
Run clients training in parallel.
"""
if self._clients is None:
self._clients = self._setup_clients(self.dataset, self._client_config)
if self._complete_experiment:
# reset before training if a previous experiment is completed
self._reset()
self._complete_experiment = False
self._logger_manager.log_message("Training federated...")
self.n_iter = 0
with tqdm(
range(self.config.num_iters),
total=self.config.num_iters,
desc=f"{self.config.algorithm} training",
unit="iter",
mininterval=1, # usually useless since a full iteration takes a very long time
) as outer_pbar:
for self.n_iter in outer_pbar:
selected_clients = self._sample_clients()
with tqdm(
total=len(selected_clients),
desc=f"Iter {self.n_iter+1}/{self.config.num_iters}",
unit="client",
mininterval=max(1, len(selected_clients) // 50),
disable=self.config.verbose < 1,
leave=False,
) as pbar:
for client_id in selected_clients:
client = self._clients[client_id]
# server communicates with client
# typically broadcasting the global model to the client
self._communicate(client)
if self.n_iter > 0 and (self.n_iter + 1) % self.config.eval_every == 0:
for part in self.dataset.data_parts:
# NOTE: one should execute `client.evaluate`
# before `client._update`,
# otherwise the evaluation would be done
# on the ``personalized`` (locally fine-tuned) model.
metrics = client.evaluate(part)
self._logger_manager.log_metrics(
client_id,
metrics,
step=self.n_iter,
epoch=self.n_iter,
part=part,
)
# client trains the model
# and perhaps updates other local variables
client._update()
# client communicates with server
# typically sending the local model
# and evaluated local metrics to the server
# and perhaps other local variables (e.g. gradients, etc.)
client._communicate(self)
pbar.update(1)
if self.n_iter > 0 and (self.n_iter + 1) % self.config.eval_every == 0:
# server aggregates the metrics from clients
self.aggregate_client_metrics()
# server updates the global model
# and perhaps other global variables
self._update()
self._logger_manager.log_message("Federated training finished...")
self._logger_manager.flush()
self._complete_experiment = True
[docs] def train_local(self, extra_configs: Optional[dict] = None) -> None:
"""Local training, conducted on the clients,
without any communication with the server. Used for comparison.
Parameters
----------
extra_configs : dict, optional
The extra configs for local training.
Returns
-------
None
"""
if self._clients is None:
self._clients = self._setup_clients(self.dataset, self._client_config)
if self._complete_experiment:
# reset before training if a previous experiment is completed
self._reset()
self._complete_experiment = False
self._logger_manager.log_message("Training local...")
self.n_iter = 0
with tqdm(
range(self.config.num_iters),
total=self.config.num_iters,
desc=f"{self.config.algorithm} training",
unit="iter",
mininterval=1, # usually useless since a full iteration takes a very long time
) as outer_pbar:
for self.n_iter in outer_pbar:
with tqdm(
total=len(self._clients),
desc=f"Iter {self.n_iter+1}/{self.config.num_iters}",
unit="client",
mininterval=max(1, len(self._clients) // 10),
disable=self.config.verbose < 1,
leave=False,
) as pbar:
for client_id in range(len(self._clients)):
client = self._clients[client_id]
client.train()
if self.n_iter > 0 and (self.n_iter + 1) % self.config.eval_every == 0:
for part in self.dataset.data_parts:
metrics = client.evaluate(part)
self._logger_manager.log_metrics(
client_id,
metrics,
step=self.n_iter,
epoch=self.n_iter,
part=part,
)
pbar.update(1)
self._logger_manager.log_message("Local training finished...")
self._logger_manager.flush()
self._complete_experiment = True
[docs] def evaluate_centralized(self, dataloader: DataLoader) -> Dict[str, float]:
"""Evaluate the model on the given dataloader on the server node.
Parameters
----------
dataloader : DataLoader
The dataloader for evaluation.
Returns
-------
metrics : dict
The metrics of the model on the given dataloader.
"""
metrics = []
for (X, y) in dataloader:
X, y = X.to(self.model.device), y.to(self.model.device)
probs = self.model(X)
metrics.append(self.dataset.evaluate(probs, y))
num_samples = sum([m["num_samples"] for m in metrics])
metrics_names = [k for k in metrics[0] if k != "num_samples"]
metrics = {k: sum([m[k] * m["num_samples"] for m in metrics]) / num_samples for k in metrics_names}
metrics["num_samples"] = num_samples
# free memory
del X, y, probs
return metrics
[docs] def aggregate_client_metrics(self, ignore: Sequence[str] = None) -> None:
"""Aggregate the metrics transmitted from the clients.
Parameters
----------
ignore : Sequence[str], optional
The metrics to ignore.
Returns
-------
None
"""
ignore = ignore or []
if not any(["metrics" in m for m in self._received_messages]):
raise ValueError("no metrics received from clients")
# cache metrics stored in self._metrics into self._cached_metrics
if self._metrics: # not empty
self._cached_metrics.append(self._metrics.copy())
new_metrics = defaultdict(lambda: defaultdict(float))
for part in self.dataset.data_parts:
for m in self._received_messages:
if "metrics" not in m:
continue
for k, v in m["metrics"][part].items():
if k != "num_samples":
new_metrics[part][k] += m["metrics"][part][k] * m["metrics"][part]["num_samples"]
elif k in ignore:
continue
else: # num_samples
new_metrics[part][k] += m["metrics"][part][k]
for k in new_metrics[part]:
if k != "num_samples":
new_metrics[part][k] /= new_metrics[part]["num_samples"]
self._logger_manager.log_metrics(
None,
default_dict_to_dict(new_metrics[part]),
step=self.n_iter + 1,
epoch=self.n_iter + 1,
part=part,
)
# move self._metrics to self._cached_metrics and refresh self._metrics
self._cached_metrics.append(self._metrics.copy())
self._metrics = default_dict_to_dict(new_metrics)
del new_metrics
# TODO: use the cached sequence `self._cached_metrics` of aggregated metrics
# to update the status varibles `self._is_convergent`
[docs] def add_parameters(self, params: Iterable[Parameter], ratio: float) -> None:
"""Update the server's parameters with the given parameters.
Parameters
----------
params : Iterable[torch.nn.parameter.Parameter]
The parameters to be added.
ratio : float
The ratio of the parameters to be added.
Returns
-------
None
"""
for server_param, param in zip(self.model.parameters(), params):
server_param.data.add_(param.data.detach().clone().to(self.device), alpha=ratio)
[docs] def avg_parameters(self, size_aware: bool = False, inertia: float = 0.0) -> None:
"""Update the server's parameters via
averaging the parameters received from the clients.
Parameters
----------
size_aware : bool, default False
Whether to use the size-aware averaging,
which is the weighted average of the parameters,
where the weight is the number of training samples.
From the view of optimization theory,
this is recommended to be set `False`.
inertia : float, default 0.0
The weight of the previous parameters,
should be in the range [0, 1).
Returns
-------
None
"""
assert 0.0 <= inertia < 1.0, "`inertia` should be in [0, 1)"
if len(self._received_messages) == 0:
return
for param in self.model.parameters():
param.data.mul_(inertia)
total_samples = sum([m["train_samples"] for m in self._received_messages])
for m in self._received_messages:
ratio = (m["train_samples"] / total_samples if size_aware else 1 / len(self._received_messages)) * (1 - inertia)
self.add_parameters(m["parameters"], ratio)
[docs] def update_gradients(self) -> None:
"""Update the server's gradients."""
if len(self._received_messages) == 0:
return
assert all(["gradients" in m for m in self._received_messages]), "some clients have not sent gradients yet"
# self.model.zero_grad()
for mp, gd in zip(self.model.parameters(), self._received_messages[0]["gradients"]):
mp.grad = torch.zeros_like(gd).to(self.device)
total_samples = sum([m["train_samples"] for m in self._received_messages])
for rm in self._received_messages:
# TODO: consider alpha = 1 / len(self._received_messages)
for mp, gd in zip(self.model.parameters(), rm["gradients"]):
mp.grad.add_(
gd.detach().clone().to(self.device),
alpha=rm["train_samples"] / total_samples,
)
[docs] def get_client_data(self, client_idx: int) -> Tuple[Tensor, Tensor]:
"""Get all the data of the given client.
This method is a helper function for fast access
to the data of the given client.
Parameters
----------
client_idx : int
The index of the client.
Returns
-------
Tuple[Tensor, Tensor]
Input data and labels of the given client.
"""
if client_idx >= len(self._clients):
raise ValueError(f"client_idx should be less than {len(self._clients)}")
return self._clients[client_idx].get_all_data()
[docs] def get_client_model(self, client_idx: int) -> torch.nn.Module:
"""Get the model of the given client.
This method is a helper function for fast access
to the model of the given client.
Parameters
----------
client_idx : int
The index of the client.
Returns
-------
torch.nn.Module
The model of the given client.
"""
if client_idx >= len(self._clients):
raise ValueError(f"client_idx should be less than {len(self._clients)}")
return self._clients[client_idx].model
[docs] def get_cached_metrics(self, client_idx: Optional[int] = None) -> List[Dict[str, float]]:
"""Get the cached metrics of the given client,
or the cached aggregated metrics stored on the server.
Parameters
----------
client_idx : int, optional
The index of the client. If `None`,
returns the cached aggregated metrics stored on the server.
Returns
-------
List[Dict[str, float]]
The cached metrics of the given client,
or the cached aggregated metrics stored on the server.
"""
if client_idx is None:
return self._cached_metrics
if client_idx >= len(self._clients):
raise ValueError(f"client_idx should be less than {len(self._clients)}")
return self._clients[client_idx]._cached_metrics
def _reset(self, reset_clients: bool = True) -> None:
"""Reset the server.
State variables are reset to the initial state,
and the logger manager is reset, which outputs
to new log files.
Parameters
----------
reset_clients : bool, default True
Whether to reset the clients.
Returns
-------
None
"""
self._received_messages = []
self._metrics = {}
self._cached_metrics = []
self._is_convergent = False
self.n_iters = 0
self._num_communications = 0
self._logger_manager.reset()
self._complete_experiment = False
# reset the global model
reset_parameters(self.model)
# reset the clients
if reset_clients:
for c in tqdm(self._clients, desc="Resetting clients", mininterval=1, leave=False):
c._reset()
@property
@abstractmethod
def client_cls(self) -> type:
"""Class of the client node."""
raise NotImplementedError
@property
@abstractmethod
def config_cls(self) -> Dict[str, type]:
"""Class of the client node config and server node config.
Keys are "client" and "server".
"""
raise NotImplementedError
@property
@abstractmethod
def doi(self) -> Union[str, List[str]]:
raise NotImplementedError
@property
def is_convergent(self) -> bool:
"""Whether the training process is convergent."""
return self._is_convergent
[docs]class Client(Node):
"""The class to simulate the client node.
The client node is responsible for training the local models,
and communicating with the server node.
Parameters
----------
client_id : int
The id of the client.
device : torch.device
The device to train the model on.
model : torch.nn.Module
The model to train.
dataset : FedDataset
The dataset to train on.
config : ClientConfig
The config for the client.
"""
__name__ = "Client"
def __init__(
self,
client_id: int,
device: torch.device,
model: nn.Module,
dataset: FedDataset,
config: ClientConfig,
) -> None:
self.client_id = client_id
self.device = device
self.model = model
self.model.to(self.device)
self.criterion = deepcopy(dataset.criterion)
self.dataset = dataset
self.config = config
self.train_loader, self.val_loader = self.dataset.get_dataloader(
self.config.batch_size, self.config.batch_size, self.client_id
)
self.optimizer = get_optimizer(
optimizer_name=self.config.optimizer,
params=self.model.parameters(),
config=self.config,
)
scheduler_config = {k: v for k, v in self.config.scheduler.items() if k != "name"}
self.lr_scheduler = get_scheduler(
scheduler_name=self.config.scheduler["name"],
optimizer=self.optimizer,
config=scheduler_config,
)
self._cached_parameters = None
self._received_messages = {}
self._metrics = {}
self._cached_metrics = [] # container for caching metrics
self._is_convergent = False
self._post_init()
# checks that the client has all the required attributes in config.extra_observes
for attr in self.config.extra_observes:
assert hasattr(self, attr), f"{self.__name__} should have attribute {attr} for extra observes."
def _communicate(self, target: "Server") -> None:
"""Check validity and send messages to the server,
and maintain state variables.
Parameters
----------
target : Server
The server to communicate with.
Returns
-------
None
"""
# check validity of self._metrics
for part, metrics in self._metrics.items():
assert isinstance(metrics, dict), f"metrics for {part} should be a dict, " f"but got {type(metrics).__name__}"
assert "num_samples" in metrics, (
"In order to let the server aggregate the metrics, "
f"metrics for {part} should have key `num_samples`, "
f"but got {metrics.keys()}"
)
self.communicate(target)
target._num_communications += 1
# move self._metrics to self._cached_metrics
self._cached_metrics.append(self._metrics.copy())
self._metrics = {} # clear the metrics
def _update(self) -> None:
"""Client update, and clear cached messages
from the server of the previous iteration.
"""
self.update()
# free the memory in received_messages
del self._received_messages
self._received_messages = {} # clear the received messages
[docs] @abstractmethod
def train(self) -> None:
"""Main part of inner loop solver.
Basic example:
.. code-block:: python
self.model.train()
epoch_losses = []
for epoch in range(self.config.num_epochs):
batch_losses = []
for batch_idx, (data, target) in enumerate(self.train_loader):
data, target = data.to(self.device), target.to(self.device)
self.optimizer.zero_grad()
output = self.model(data)
loss = self.criterion(output, target)
loss.backward()
self.optimizer.step()
batch_losses.append(loss.item())
epoch_losses.append(sum(batch_losses) / len(batch_losses))
self.lr_scheduler.step()
"""
raise NotImplementedError
[docs] @add_docstring(train.__doc__)
def solve_inner(self) -> None:
"""alias of `train`"""
self.train()
[docs] def sample_data(self) -> Tuple[Tensor, Tensor]:
"""Sample data for training."""
return next(iter(self.train_loader))
[docs] @torch.no_grad()
def evaluate(self, part: str) -> Dict[str, float]:
"""Evaluate the model on the given part of the dataset.
Parameters
----------
part : str
The part of the dataset to evaluate on,
can be either "train" or "val".
Returns
-------
Dict[str, float]
The metrics of the evaluation.
"""
assert part in self.dataset.data_parts, f"Invalid part name, should be one of {self.dataset.data_parts}."
self.model.eval()
# _metrics = []
data_loader = self.val_loader if part == "val" else self.train_loader
if data_loader is None:
self._metrics[part] = {"num_samples": 0}
return self._metrics[part]
all_logits, all_labels = [], []
for X, y in data_loader:
X, y = X.to(self.device), y.to(self.device)
logits = self.model(X)
all_logits.append(logits)
all_labels.append(y)
# _metrics.append(self.dataset.evaluate(logits, y))
all_logits = torch.cat(all_logits, dim=0)
all_labels = torch.cat(all_labels, dim=0)
self._metrics[part] = {"num_samples": len(all_labels)}
self._metrics[part].update(self.dataset.evaluate(all_logits, all_labels))
# get the gradient norm of the local model
self._metrics[part]["grad_norm"] = self.get_gradients(norm="fro")
# record items in config.extra_observes
for attr in self.config.extra_observes:
self._metrics[part][attr] = Node.get_norm(getattr(self, attr))
# free the memory
del all_logits, all_labels, X, y
return self._metrics[part]
[docs] def get_all_data(self) -> Tuple[Tensor, Tensor]:
"""Get all the data on the client.
This method is a helper function for fast access
to the data on the client,
including both training and validation data;
both features and labels.
"""
feature_train, label_train = self.train_loader.dataset[:]
if self.val_loader is None:
return feature_train, label_train
feature_val, label_val = self.val_loader.dataset[:]
feature = torch.cat([feature_train, feature_val], dim=0)
label = torch.cat([label_train, label_val], dim=0)
return feature, label
def _reset(self) -> None:
"""Reset the model and optimizer."""
# reset the model
reset_parameters(self.model)
# reset the optimizer and clear gradients
del self.optimizer
self.optimizer = get_optimizer(
optimizer_name=self.config.optimizer,
params=self.model.parameters(),
config=self.config,
)
self.optimizer.zero_grad()
@property
def is_convergent(self) -> bool:
"""Whether the training process is convergent."""
return self._is_convergent
class ClientMessage(dict):
"""A class used to specify required fields
for a message from client to server.
Parameters
----------
client_id : int
The id of the client.
train_samples : int
The number of samples used for training on the client.
metrics : dict
The metrics evaluated on the client.
**kwargs : dict, optional
Extra message to be sent to the server.
"""
__name__ = "ClientMessage"
def __init__(self, client_id: int, train_samples: int, metrics: dict, **kwargs) -> None:
super().__init__(client_id=client_id, train_samples=train_samples, metrics=metrics, **kwargs)