fedsim.fl.fl_algorithm — fedsim 0.1 documentation

# Source code for fedsim.fl.fl_algorithm

    from http import client
    from tqdm import trange
    import random
    import math
    import torch
    from torch import nn
    from torch.utils.data import DataLoader
    from typing import Any, Dict, Hashable, Iterable, Mapping, Optional, Union
    import logging
    import inspect
    from pprint import pformat

    from .aggregators import SerialAggregator
    from fedsim.utils import search_in_submodules

    class FLAlgorithm(object):

        def __init__(
            self,
            data_manager,
            metric_logger,
            num_clients,
            sample_scheme,
            sample_rate,
            model_class,
            epochs,
            loss_fn,
            batch_size,
            test_batch_size,
            local_weight_decay,
            slr,
            clr,
            clr_decay,
            clr_decay_type,
            min_clr,
            clr_step_size,
            device,
            log_freq,
            *args,
            **kwargs,
        ):
            grandpa = inspect.getmro(self.__class__)[-2]
            cls = self.__class__

            grandpa_args = set(inspect.signature(grandpa).parameters.keys())
            self_args = set(inspect.signature(cls).parameters.keys())

            added_args = self_args - grandpa_args
            added_args_dict = {key: getattr(self, key) for key in added_args}
            if len(added_args_dict) > 0:
                logging.info('algorithm arguments: ' + pformat(added_args_dict))

            self._data_manager = data_manager
            self.num_clients = num_clients
            self.sample_scheme = sample_scheme
            self.sample_count = int(sample_rate * num_clients)
            if not 1 <= self.sample_count <= num_clients:
                raise Exception(
                    'invalid client sample size for {}% of {} clients'.format(
                        sample_rate, num_clients))

            if isinstance(model_class, str):
                self.model_class = search_in_submodules('fedsim.models',
                                                        model_class)
            elif issubclass(model_class, nn.Module):
                self.model_class = model_class
            else:
                raise Exception("incompatiple model!")
            self.epochs = epochs
            if loss_fn == 'ce':
                self.loss_fn = nn.CrossEntropyLoss()
            elif loss_fn == 'mse':
                self.loss_fn = nn.MSELoss()
            else:
                raise NotImplementedError
            self.batch_size = batch_size
            self.test_batch_size = test_batch_size
            self.local_weight_decay = local_weight_decay
            self.slr = slr
            self.clr = clr
            self.clr_decay = clr_decay
            self.clr_decay_type = clr_decay_type
            self.min_clr = min_clr
            self.clr_step_size = clr_step_size

            self.metric_logger = metric_logger
            self.device = device
            self.log_freq = log_freq

            self._server_memory: Dict[Hashable, object] = dict()
            self._client_memory: Dict[int, Dict[object]] = {
                k: dict()
                for k in range(num_clients)
            }

            self.global_dataloaders = {
                key: DataLoader(
                    dataset,
                    batch_size=self.test_batch_size,
                    pin_memory=True,
                )
                for key, dataset in
                self._data_manager.get_global_dataset().items()
            }

            self.oracle_dataset = self._data_manager.get_oracle_dataset()
            self.rounds = 0

            # for internal use only
            self._last_client_sampled: int = None

        def write_server(self, key, obj):
            self._server_memory[key] = obj

        def write_client(self, client_id, key, obj):
            self._client_memory[client_id][key] = obj

        def read_server(self, key):
            return self._server_memory[key] if key in self._server_memory else None

        def read_client(self, client_id, key):
            if client_id >= self.num_clients:
                raise Exception("invalid client id {} >=".format(
                    id, self.num_clients))
            if key in self._client_memory[client_id]:
                return self._client_memory[client_id][key]
            return None

        def _sample_clients(self):
            if self.sample_scheme == 'uniform':
                clients = random.sample(range(self.num_clients), self.sample_count)
            elif self.sample_scheme == 'sequential':
                last_sampled = -1 if self._last_client_sampled is None else self._last_client_sampled
                clients = [
                    (i + 1) % self.num_clients
                    for i in range(last_sampled, last_sampled + self.sample_count)
                ]
                self._last_client_sampled = clients[-1]
            else:
                raise NotImplementedError
            return clients

        def _send_to_client(self, client_id):
            return self.send_to_client(client_id=client_id)

        def _send_to_server(self, client_id):
            if self.clr_decay_type == 'step':
                decayed_clr = self.clr * \
                    (self.clr_decay ** (self.rounds // self.clr_step_size))
            elif self.clr_decay_type == 'cosine':
                T_i = self.clr_step_size
                T_cur = self.rounds % T_i
                decayed_clr = self.min_clr + 0.5 * \
                                (self.clr - self.min_clr) * \
                                    (1 + math.cos(math.pi * T_cur / T_i))

            client_ctx = self.send_to_server(
                client_id,
                self._data_manager.get_local_dataset(client_id),
                self.epochs,
                self.loss_fn,
                self.batch_size,
                decayed_clr,
                self.local_weight_decay,
                self.device,
                ctx=self._send_to_client(client_id))
            if not isinstance(client_ctx, dict):
                raise Exception('client should only return a dict!')
            return {**client_ctx, 'client_id': client_id}

        def _receive_from_client(self, client_msg, aggregator):
            client_id = client_msg.pop('client_id')
            return self.receive_from_client(client_id, client_msg, aggregator)

        def _optimize(self, aggregator):
            reports = self.optimize(aggregator)
            # purge aggregated results
            del aggregator
            return reports

        def _report(self, optimize_reports=None, deployment_points=None):
            self.report(self.global_dataloaders, self.metric_logger, self.device,
                        optimize_reports, deployment_points)

        def _train(self, rounds):
            for self.rounds in trange(rounds):
                aggregator = SerialAggregator()
                for client_id in self._sample_clients():
                    client_msg = self._send_to_server(client_id)
                    self._receive_from_client(client_msg, aggregator)
                opt_reports = self._optimize(aggregator)
                if self.rounds % self.log_freq == 0:
                    deploy_poiont = self.deploy()
                    self._report(opt_reports, deploy_poiont)
            # one last report
            if self.rounds % self.log_freq > 0:
                deploy_poiont = self.deploy()
                self._report(opt_reports, deploy_poiont)
            return 0

        def train(self, rounds):
            return self._train(rounds=rounds)

        def get_model_class(self):
            return self.model_class

        # we do not do type hinting, however, the hints for avstract
        # methods are provided to help clarity for users

        def send_to_client(self, client_id: int) -> Mapping[Hashable, Any]:
            """ returns context to send to the client corresponding to client_id.
                Do not send shared objects like server model if you made any
                before you deepcopy it.

            Args:
                client_id (int): id of the receiving client

            Raises:
                NotImplementedError: abstract class to be implemented by child

            Returns:
                Mapping[Hashable, Any]: the context to be sent in form of a Mapping
            """
            raise NotImplementedError

        def send_to_server(
            self,
            client_id: int,
            datasets: Dict[str, Iterable],
            epochs: int,
            loss_fn: nn.Module,
            batch_size: int,
            lr: float,
            weight_decay: float = 0,
            device: Union[int, str] = 'cuda',
            ctx: Optional[Dict[Hashable, Any]] = None,
            *args,
            **kwargs,
        ) -> Mapping[str, Any]:
            """ client operation on the recieved information.

            Args:
                client_id (int): id of the client
                datasets (Dict[str, Iterable]): this comes from Data Manager
                epochs (int): number of epochs to train
                loss_fn (nn.Module): either 'ce' (for cross-entropy) or 'mse'
                batch_size (int): training batch_size
                lr (float): client learning rate
                weight_decay (float, optional): weight decay for SGD. Defaults to 0.
                device (Union[int, str], optional): Defaults to 'cuda'.
                ctx (Optional[Dict[Hashable, Any]], optional): context reveived from server. Defaults to None.

            Raises:
                NotImplementedError: abstract class to be implemented by child

            Returns:
                Mapping[str, Any]: client context to be sent to the server
            """
            raise NotImplementedError

        def receive_from_client(self, client_id: int, client_msg: Mapping[Hashable,
                                                                          Any],
                                aggregator: Any):
            """receive and aggregate info from selected clients

            Args:
                client_id (int): id of the sender (client)
                client_msg (Mapping[Hashable, Any]): client context that is sent
                aggregator (Any): aggregator instance to collect info

            Raises:
                NotImplementedError: abstract class to be implemented by child
            """
            raise NotImplementedError

        def optimize(self, aggregator: Any) -> Mapping[Hashable, Any]:
            """ optimize server mdoel(s) and return metrics to be reported

            Args:
                aggregator (Any): Aggregator instance

            Raises:
                NotImplementedError: abstract class to be implemented by child

            Returns:
                Mapping[Hashable, Any]: context to be reported
            """
            raise NotImplementedError

        def deploy(self) -> Optional[Mapping[Hashable, Any]]:
            """ return Mapping of name -> parameters_set to test the model

            Raises:
                NotImplementedError: abstract class to be implemented by child
            """
            raise NotImplementedError

        def report(
            self,
            dataloaders: Dict[str, Any],
            metric_logger: Any,
            device: str,
            optimize_reports: Mapping[Hashable, Any],
            deployment_points: Optional[Mapping[Hashable, torch.Tensor]] = None
        ) -> None:
            """test on global data and report info

            Args:
                dataloaders (Any): dict of data loaders to test the global model(s)
                metric_logger (Any): the logging object (e.g., SummaryWriter)
                device (str): 'cuda', 'cpu' or gpu number
                optimize_reports (Mapping[Hashable, Any]): dict returned by optimzier
                deployment_points (Mapping[Hashable, torch.Tensor], optional): output of deploy method

            Raises:
                NotImplementedError: abstract class to be implemented by child
            """
            raise NotImplementedError