Update FedProx baseline

baselines/fedprox/README.md

@@ -1,15 +1,15 @@
 ---
 title: Federated Optimization in Heterogeneous Networks
 url: https://arxiv.org/abs/1812.06127
-labels: [image classification, cross-device, stragglers] # please add between 4 and 10 single-word (maybe two-words) labels (e.g. "system heterogeneity", "image classification", "asynchronous", "weight sharing", "cross-silo")
-dataset: [mnist] # list of datasets you include in your baseline
+labels: [image classification, cross-device, stragglers]
+dataset: [mnist]
 ---
 
 # FedProx: Federated Optimization in Heterogeneous Networks
 
 > Note: If you use this baseline in your work, please remember to cite the original authors of the paper as well as the Flower paper.
 
-**Paper:** https://arxiv.org/abs/1812.06127
+**Paper:** [arxiv.org/abs/1812.06127](https://arxiv.org/abs/1812.06127)
 
 **Authors:** Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar and Virginia Smith.
 
@@ -34,7 +34,7 @@ dataset: [mnist] # list of datasets you include in your baseline
 * A logistic regression model used in the FedProx paper for MNIST (see `models/LogisticRegression`). This is the model used by default.
 * A two-layer CNN network as used in the FedAvg paper (see `models/Net`)
 
-**Dataset:** This baseline only includes the MNIST dataset. By default it will be partitioned into 1000 clients following a pathological split where each client has examples of two (out of ten) class labels. The number of examples in each client is derived by sampling from a powerlaw distribution. The settings are as follow:
+**Dataset:** This baseline only includes the MNIST dataset. By default, it will be partitioned into 1000 clients following a pathological split where each client has examples of two (out of ten) class labels. The number of examples in each client is derived by sampling from a powerlaw distribution. The settings are as follows:
 
 | Dataset | #classes | #partitions | partitioning method | partition settings |
 | :------ | :---: | :---: | :---: | :---: |
@@ -56,17 +56,10 @@ The following table shows the main hyperparameters for this baseline with their
 
 ## Environment Setup
 
-To construct the Python environment follow these steps:
+To construct the Python environment, simply run:
 
 ```bash
-# install the base Poetry environment
 poetry install
-
-# activate the environment
-poetry shell
-
-# install PyTorch with GPU support. Please note this baseline is very lightweight so it can run fine on a CPU.
-pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116
 ```
 
 ## Running the Experiments
@@ -96,7 +89,7 @@ python -m fedprox.main --config-name fedavg
 
 ## Expected results
 
-With the following command we run both FedProx and FedAvg configurations while iterating through different values of `mu` and `stragglers_fraction`. We ran each experiment five times (this is achieved by artificially adding an extra element to the config but that it doesn't have an impact on the FL setting `'+repeat_num=range(5)'`)
+With the following command, we run both FedProx and FedAvg configurations while iterating through different values of `mu` and `stragglers_fraction`. We ran each experiment five times (this is achieved by artificially adding an extra element to the config but it doesn't have an impact on the FL setting `'+repeat_num=range(5)'`)
 
 ```bash
 python -m fedprox.main --multirun mu=0.0,2.0 stragglers_fraction=0.0,0.5,0.9 '+repeat_num=range(5)'

baselines/fedprox/fedprox/__init__.py

@@ -0,0 +1 @@
+"""FedProx package."""

baselines/fedprox/fedprox/client.py

@@ -12,10 +12,10 @@
 from omegaconf import DictConfig
 from torch.utils.data import DataLoader
 
-from fedprox.dataset import load_datasets
 from fedprox.models import test, train
 
 
+# pylint: disable=too-many-arguments
 class FlowerClient(
     fl.client.NumPyClient
 ):  # pylint: disable=too-many-instance-attributes
@@ -40,19 +40,19 @@ def __init__(
         self.straggler_schedule = straggler_schedule
 
     def get_parameters(self, config: Dict[str, Scalar]) -> NDArrays:
-        """Returns the parameters of the current net."""
+        """Return the parameters of the current net."""
         return [val.cpu().numpy() for _, val in self.net.state_dict().items()]
 
     def set_parameters(self, parameters: NDArrays) -> None:
-        """Changes the parameters of the model using the given ones."""
+        """Change the parameters of the model using the given ones."""
         params_dict = zip(self.net.state_dict().keys(), parameters)
         state_dict = OrderedDict({k: torch.Tensor(v) for k, v in params_dict})
         self.net.load_state_dict(state_dict, strict=True)
 
     def fit(
         self, parameters: NDArrays, config: Dict[str, Scalar]
     ) -> Tuple[NDArrays, int, Dict]:
-        """Implements distributed fit function for a given client."""
+        """Implement distributed fit function for a given client."""
         self.set_parameters(parameters)
 
         # At each round check if the client is a straggler,
@@ -88,15 +88,15 @@ def fit(
             self.device,
             epochs=num_epochs,
             learning_rate=self.learning_rate,
-            proximal_mu=config["proximal_mu"],
+            proximal_mu=float(config["proximal_mu"]),
         )
 
         return self.get_parameters({}), len(self.trainloader), {"is_straggler": False}
 
     def evaluate(
         self, parameters: NDArrays, config: Dict[str, Scalar]
     ) -> Tuple[float, int, Dict]:
-        """Implements distributed evaluation for a given client."""
+        """Implement distributed evaluation for a given client."""
         self.set_parameters(parameters)
         loss, accuracy = test(self.net, self.valloader, self.device)
         return float(loss), len(self.valloader), {"accuracy": float(accuracy)}
@@ -111,10 +111,8 @@ def gen_client_fn(
     learning_rate: float,
     stragglers: float,
     model: DictConfig,
-) -> Tuple[
-    Callable[[str], FlowerClient], DataLoader
-]:  # pylint: disable=too-many-arguments
-    """Generates the client function that creates the Flower Clients.
+) -> Callable[[str], FlowerClient]:  # pylint: disable=too-many-arguments
+    """Generate the client function that creates the Flower Clients.
 
     Parameters
     ----------
@@ -139,13 +137,11 @@ def gen_client_fn(
 
     Returns
     -------
-    Tuple[Callable[[str], FlowerClient], DataLoader]
-        A tuple containing the client function that creates Flower Clients and
-        the DataLoader that will be used for testing
+    Callable[[str], FlowerClient]
+        A client function that creates Flower Clients.
     """
-
-    # Defines a staggling schedule for each clients, i.e at which round will they
-    # be a straggler. This is done so at each round the proportion of staggling
+    # Defines a straggling schedule for each clients, i.e at which round will they
+    # be a straggler. This is done so at each round the proportion of straggling
     # clients is respected
     stragglers_mat = np.transpose(
         np.random.choice(
@@ -155,7 +151,6 @@ def gen_client_fn(
 
     def client_fn(cid: str) -> FlowerClient:
         """Create a Flower client representing a single organization."""
-
         # Load model
         device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
         net = instantiate(model).to(device)

baselines/fedprox/fedprox/conf/config.yaml

@@ -29,12 +29,12 @@ model:
 
 strategy:
   _target_: flwr.server.strategy.FedProx
-  fraction_fit: 0.00001 # because we want the number of clients to sample on each roudn to be solely defined by min_fit_clients
+  fraction_fit: 0.00001 # because we want the number of clients to sample on each round to be solely defined by min_fit_clients
   fraction_evaluate: 0.0
   min_fit_clients: ${clients_per_round}
   min_evaluate_clients: 0
   min_available_clients: ${clients_per_round}
   evaluate_metrics_aggregation_fn:
     _target_: fedprox.strategy.weighted_average
-    _partial_: true # we dont' want this function to be evaluated when instantiating the strategy, we treat it as a partial and evaluate it when the strategy actuallly calls the function (in aggregate_evaluate())
+    _partial_: true # we dont' want this function to be evaluated when instantiating the strategy, we treat it as a partial and evaluate it when the strategy actually calls the function (in aggregate_evaluate())
   proximal_mu: ${mu}

baselines/fedprox/fedprox/conf/fedavg.yaml

@@ -34,11 +34,11 @@ model:
 
 strategy:
   _target_: fedprox.strategy.FedAvgWithStragglerDrop #! this points to FedAvgWithStragglerDrop class in strategy.py, Note that we need the full module path (including `fedprox`)
-  fraction_fit: 0.00001 # because we want the number of clients to sample on each roudn to be solely defined by min_fit_clients
+  fraction_fit: 0.00001 # because we want the number of clients to sample on each round to be solely defined by min_fit_clients
   fraction_evaluate: 0.0
   min_fit_clients: ${clients_per_round}
   min_available_clients: ${clients_per_round}
   min_evaluate_clients: 0
   evaluate_metrics_aggregation_fn:
     _target_: fedprox.strategy.weighted_average
-    _partial_: true # we dont' want this function to be evaluated when instantiating the strategy, we treat it as a partial and evaluate it when the strategy actuallly calls the function (in aggregate_evaluate())
+    _partial_: true # we dont' want this function to be evaluated when instantiating the strategy, we treat it as a partial and evaluate it when the strategy actually calls the function (in aggregate_evaluate())

baselines/fedprox/fedprox/dataset.py

@@ -17,7 +17,7 @@ def load_datasets(  # pylint: disable=too-many-arguments
     batch_size: Optional[int] = 32,
     seed: Optional[int] = 42,
 ) -> Tuple[DataLoader, DataLoader, DataLoader]:
-    """Creates the dataloaders to be fed into the model.
+    """Create the dataloaders to be fed into the model.
 
     Parameters
     ----------
@@ -36,7 +36,8 @@ def load_datasets(  # pylint: disable=too-many-arguments
     Returns
     -------
     Tuple[DataLoader, DataLoader, DataLoader]
-        The DataLoader for training, the DataLoader for validation, the DataLoader for testing.
+        The DataLoader for training, the DataLoader for validation, the DataLoader
+        for testing.
     """
     print(f"Dataset partitioning config: {config}")
     datasets, testset = _partition_data(

baselines/fedprox/fedprox/dataset_preparation.py

@@ -1,3 +1,4 @@
+"""Functions for dataset download and processing."""
 from typing import List, Optional, Tuple
 
 import numpy as np
@@ -8,7 +9,7 @@
 
 
 def _download_data() -> Tuple[Dataset, Dataset]:
-    """Downloads (if necessary) and returns the MNIST dataset.
+    """Download (if necessary) and returns the MNIST dataset.
 
     Returns
     -------
@@ -23,24 +24,27 @@ def _download_data() -> Tuple[Dataset, Dataset]:
     return trainset, testset
 
 
+# pylint: disable=too-many-locals
 def _partition_data(
     num_clients,
     iid: Optional[bool] = False,
     power_law: Optional[bool] = True,
     balance: Optional[bool] = False,
     seed: Optional[int] = 42,
 ) -> Tuple[List[Dataset], Dataset]:
-    """Split training set into iid or non iid partitions to simulate the
-    federated setting.
+    """Split training set into iid or non iid partitions to simulate the federated.
+
+    setting.
 
     Parameters
     ----------
     num_clients : int
         The number of clients that hold a part of the data
     iid : bool, optional
         Whether the data should be independent and identically distributed between
-        the clients or if the data should first be sorted by labels and distributed by chunks
-        to each client (used to test the convergence in a worst case scenario), by default False
+        the clients or if the data should first be sorted by labels and distributed
+        by chunks to each client (used to test the convergence in a worst case scenario)
+        , by default False
     power_law: bool, optional
         Whether to follow a power-law distribution when assigning number of samples
         for each client, defaults to True
@@ -53,13 +57,14 @@ def _partition_data(
     Returns
     -------
     Tuple[List[Dataset], Dataset]
-        A list of dataset for each client and a single dataset to be use for testing the model.
+        A list of dataset for each client and a single dataset to be use for testing
+        the model.
     """
     trainset, testset = _download_data()
 
     if balance:
         trainset = _balance_classes(trainset, seed)
-        
+
     partition_size = int(len(trainset) / num_clients)
     lengths = [partition_size] * num_clients
 
@@ -171,6 +176,7 @@ def _sort_by_class(
     return sorted_dataset
 
 
+# pylint: disable=too-many-locals, too-many-arguments
 def _power_law_split(
     sorted_trainset: Dataset,
     num_partitions: int,
@@ -179,9 +185,10 @@ def _power_law_split(
     mean: float = 0.0,
     sigma: float = 2.0,
 ) -> Dataset:
-    """Partitions the dataset following a power-law distribution. It follows
-    the implementation of Li et al 2020: https://arxiv.org/abs/1812.06127 with
-    default values set accordingly.
+    """Partition the dataset following a power-law distribution. It follows the.
+
+    implementation of Li et al 2020: https://arxiv.org/abs/1812.06127 with default
+    values set accordingly.
 
     Parameters
     ----------
@@ -205,15 +212,14 @@ def _power_law_split(
     Dataset
         The partitioned training dataset.
     """
-
     targets = sorted_trainset.targets
-    full_idx = range(len(targets))
+    full_idx = list(range(len(targets)))
 
     class_counts = np.bincount(sorted_trainset.targets)
     labels_cs = np.cumsum(class_counts)
     labels_cs = [0] + labels_cs[:-1].tolist()
 
-    partitions_idx = []
+    partitions_idx: List[List[int]] = []
     num_classes = len(np.bincount(targets))
     hist = np.zeros(num_classes, dtype=np.int32)
 
@@ -243,7 +249,9 @@ def _power_law_split(
         (num_classes, int(num_partitions / num_classes), num_labels_per_partition),
     )
     remaining_per_class = class_counts - hist
-    # obtain how many samples each partition should be assigned for each of the labels it contains
+    # obtain how many samples each partition should be assigned for each of the
+    # labels it contains
+    # pylint: disable=too-many-function-args
     probs = (
         remaining_per_class.reshape(-1, 1, 1)
         * probs

baselines/fedprox/fedprox/main.py

@@ -1,5 +1,7 @@
 """Runs CNN federated learning for MNIST dataset."""
 
+from typing import Dict, Union
+
 import flwr as fl
 import hydra
 from hydra.core.hydra_config import HydraConfig
@@ -10,17 +12,18 @@
 from fedprox.dataset import load_datasets
 from fedprox.utils import save_results_as_pickle
 
+FitConfig = Dict[str, Union[bool, float]]
+
 
 @hydra.main(config_path="conf", config_name="config", version_base=None)
 def main(cfg: DictConfig) -> None:
-    """Main function to run CNN federated learning on MNIST.
+    """Run CNN federated learning on MNIST.
 
     Parameters
     ----------
     cfg : DictConfig
         An omegaconf object that stores the hydra config.
     """
-
     # print config structured as YAML
     print(OmegaConf.to_yaml(cfg))
 
@@ -53,7 +56,9 @@ def main(cfg: DictConfig) -> None:
     def get_on_fit_config():
         def fit_config_fn(server_round: int):
             # resolve and convert to python dict
-            fit_config = OmegaConf.to_container(cfg.fit_config, resolve=True)
+            fit_config: FitConfig = OmegaConf.to_container(  # type: ignore
+                cfg.fit_config, resolve=True
+            )
             fit_config["curr_round"] = server_round  # add round info
             return fit_config
 

baselines/fedprox/fedprox/models.py

@@ -1,4 +1,4 @@
-"""CNN model architecutre, training, and testing functions for MNIST."""
+"""CNN model architecture, training, and testing functions for MNIST."""
 
 
 from typing import List, Tuple
@@ -61,7 +61,7 @@ class LogisticRegression(nn.Module):
 
     def __init__(self, num_classes: int) -> None:
         super().__init__()
-        self.fc = nn.Linear(28 * 28, num_classes)
+        self.linear = nn.Linear(28 * 28, num_classes)
 
     def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
         """Forward pass.
@@ -76,7 +76,7 @@ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
         torch.Tensor
             The resulting Tensor after it has passed through the network
         """
-        output_tensor = self.fc(torch.flatten(input_tensor, 1))
+        output_tensor = self.linear(torch.flatten(input_tensor, 1))
         return output_tensor