Customization — fedsim 0.1 documentation

# Customization

## DataManager

Any custome DataManager class should inherit from `fedsim.data_manager.data_manager.DataManager` (or its children) and implement its abstract methods. For example:

    from fedsim.data_manager.data_manager import DataManager

    class CustomDataManager(DataManager)
        def __init__(self, root, other_arg, ...):
            self.other_arg = other_arg
            # note that super should be called at the end of init \
            # because the abstract classes are called in its __init__
            super(BasicDataManager, self).__init__(root, seed, save_path=save_path)

        def make_datasets(self, root: str) -> Iterable[Dict[str, object]]:
            """Abstract method to be implemented by child class.

            Args:
                dataset_name (str): name of the dataset.
                root (str): directory to download and manipulate data.
                save_path (str): directory to store the data after partitioning.

            Raises:
                NotImplementedError: if the dataset_name is not defined

            Returns:
                Iterable[Dict[str, object]]: dict of local datasets [split:dataset]
                                             followed by global ones.
            """
            raise NotImplementedError

        def partition_local_data(self, datasets: Dict[str, object]) -> Dict[str, Iterable[Iterable[int]]]:
            raise NotImplementedError

        def get_identifiers(self) -> Sequence[str]:
            """ Returns identifiers
                to be used for saving the partition info.

            Raises:
                NotImplementedError: this abstract method should be
                    implemented by child classes

            Returns:
                Sequence[str]: a sequence of str identifing class instance
            """
            raise NotImplementedError

### Integration with included cli

To automatically include your custom data manager in the provided cli tool, you can place your class in a file under `fedsim/data_manager`. Then, call it using option `--data-manager`. To deliver arguments to the `__init__` method of your custom data manager, you can pass options in form of `--d-<arg-name>` where `<arg-name>` is the argument. Example

    fedsim fed-learn --data-manager CustomDataManager --d-other_arg <other_arg_value> ...

## FLAlgorithm

Any custome DataManager class should inherit from `fedsim.fl.fl_algorithm.FLAlgorithm` (or its children) and implement its abstract methods. For example:

    from typing import Optional, Hashable, Mapping, Dict, Any
    from fedsim.fl.fl_algorithm import FLAlgorithm

    class CustomFLAlgorithm(FLAlgorithm):
        def __init__(
            self, data_manager, 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, metric_logger, device, log_freq, other_arg, ... , *args, **kwargs,
        ):
            self.other_arg = other_arg

            super(FedAvg, self).__init__(
                data_manager, 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, metric_logger, device, log_freq,
            )
            # make mode and optimizer
            model = self.get_model_class()().to(self.device)
            params = deepcopy(
                parameters_to_vector(model.parameters()).clone().detach())
            optimizer = SGD(params=[params], lr=slr)
            # write model and optimizer to server
            self.write_server('model', model)
            self.write_server('cloud_params', params)
            self.write_server('optimizer', optimizer)
            ...

        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, 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

### Integration with included cli

To automatically include your custom algorithm by the provided cli tool, you can place your class in a file under fedsim/fl/algorithms. Then, call it using option –algorithm. To deliver arguments to the init method of your custom algorithm, you can pass options in form of –a-<arg-name> where <arg-name> is the argument. Example

    fedsim fed-learn --algorithm CustomFLAlgorithm --a-other_arg <other_arg_value> ...