Merge pull request #16 from VectorInstitute/apfl

Apfl

Author: jewelltaylor

.pre-commit-config.yaml

@@ -26,7 +26,7 @@ repos:
       - id: black
 
   - repo: https://github.com/PyCQA/isort
-    rev: 5.5.2
+    rev: 5.11.5
     hooks:
       - id: isort
 

examples/apfl_example/README.md

@@ -0,0 +1,39 @@
+# APFL Federated Learning Example
+This is an example of [Adaptive Personalized Federated Learning](https://arxiv.org/pdf/2003.13461.pdf) (APFL). APFL is a popular method for achieving personalization in the federated learning setting which is important in real world application where client distributions are heterogenous. APFL extends FedAVG by having each client train a local model distinct from the global model that is jointly learned across clients. A personalized model is generated as a convex combination of the local and global models with a per client mixing parameter that is learned. The local model is updated to minmize the loss of the personalized model on the local data.
+
+In this demo, APFL is applied to an augmented version of the MNIST dataset that is non--IID. The FL server expects three clients to be spun up (i.e. it will wait until three clients report in before starting training). Each client has a modified version of the MNIST dataset. This modification essentially subsamples a certain number from the original training and validation sets of MNIST in order to synthetically induce local variations in the statistical properties of the clients training/validation data. In theory, the models should be able to perform well on their local data while learning from other clients data that has different statistical properties. The proportion of labels at each client is determined by dirichlet distribtuion across the classes. The lower the beta parameter is for each class, the higher the degree of the label heterogeneity.
+
+The server has some custom metrics aggregation and uses Federated Averaging as its server-side optimization. The implementation uses a special type of weight exchange based on named-layer identification.
+
+# Running the Example
+In order to run the example, first ensure you have the virtual env of your choice activated and run
+```
+pip install --upgrade pip
+pip install -r requirements.txt
+```
+to install all of the dependencies for this project.
+
+## Starting Server
+
+The next step is to start the server by running
+```
+python -m examples.apfl_example.server  --config_path /path/to/config.yaml
+```
+from the FL4Health directory. The following arguments must be present in the specified config file:
+* `n_clients`: number of clients the server waits for in order to run the FL training
+* `local_epochs`: number of epochs each client will train for locally
+* `batch_size`: size of the batches each client will train on
+* `n_server_rounds`: The number of rounds to run FL
+
+## Starting Clients
+
+Once the server has started and logged "FL starting," the next step, in separate terminals, is to start the two
+clients. This is done by simply running (remembering to activate your environment)
+```
+python -m examples.apfl_example.client --dataset_path /path/to/data
+```
+**NOTE**: The argument `dataset_path` has two functions, depending on whether the dataset exists locally or not. If
+the dataset already exists at the path specified, it will be loaded from there. Otherwise, the dataset will be
+automatically downloaded to the path specified and used in the run.
+
+After all three clients have been started, federated learning should commence.

examples/apfl_example/client.py

@@ -0,0 +1,51 @@
+import argparse
+from pathlib import Path
+from typing import List
+
+import flwr as fl
+import torch
+from flwr.common.typing import Config
+
+from examples.models.cnn_model import MnistNet
+from fl4health.clients.apfl_client import ApflClient
+from fl4health.model_bases.apfl_base import APFLModule
+from fl4health.parameter_exchange.layer_exchanger import FixedLayerExchanger
+from fl4health.utils.load_data import load_mnist_data
+from fl4health.utils.metrics import Accuracy, Metric
+from fl4health.utils.sampler import DirichletLabelBasedSampler
+
+
+class MnistApflClient(ApflClient):
+    def __init__(
+        self,
+        data_path: Path,
+        metrics: List[Metric],
+        device: torch.device,
+    ) -> None:
+        super().__init__(data_path=data_path, metrics=metrics, device=device)
+
+    def setup_client(self, config: Config) -> None:
+        batch_size = self.narrow_config_type(config, "batch_size", int)
+        self.model: APFLModule = APFLModule(MnistNet()).to(self.device)
+        self.criterion = torch.nn.CrossEntropyLoss()
+        self.local_optimizer = torch.optim.AdamW(self.model.local_model.parameters(), lr=0.01)
+        self.global_optimizer = torch.optim.AdamW(self.model.global_model.parameters(), lr=0.01)
+        sampler = DirichletLabelBasedSampler(list(range(10)), sample_percentage=0.75)
+
+        self.train_loader, self.val_loader, self.num_examples = load_mnist_data(self.data_path, batch_size, sampler)
+        self.parameter_exchanger = FixedLayerExchanger(self.model.layers_to_exchange())
+
+        super().setup_client(config)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="FL Client Main")
+    parser.add_argument("--dataset_path", action="store", type=str, help="Path to the local dataset")
+
+    args = parser.parse_args()
+
+    DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    data_path = Path(args.dataset_path)
+
+    client = MnistApflClient(data_path, [Accuracy()], DEVICE)
+    fl.client.start_numpy_client(server_address="0.0.0.0:8080", client=client)

examples/apfl_example/config.yaml

@@ -0,0 +1,7 @@
+# Parameters that describe server
+n_server_rounds: 25 # The number of rounds to run FL
+
+# Parameters that describe clients
+n_clients: 3 # The number of clients in the FL experiment
+local_epochs: 1 # The number of epochs to complete for client
+batch_size: 128 # The batch size for client training

examples/apfl_example/server.py

@@ -0,0 +1,86 @@
+import argparse
+from functools import partial
+from typing import Any, Dict, List, Tuple
+
+import flwr as fl
+from flwr.common.parameter import ndarrays_to_parameters
+from flwr.common.typing import Config, Metrics, Parameters
+from flwr.server.strategy import FedAvg
+
+from examples.models.cnn_model import MnistNet
+from examples.simple_metric_aggregation import metric_aggregation, normalize_metrics
+from fl4health.model_bases.apfl_base import APFLModule
+from fl4health.utils.config import load_config
+
+
+def fit_metrics_aggregation_fn(all_client_metrics: List[Tuple[int, Metrics]]) -> Metrics:
+    # This function is run by the server to aggregate metrics returned by each clients fit function
+    # NOTE: The first value of the tuple is number of examples for FedAvg
+    total_examples, aggregated_metrics = metric_aggregation(all_client_metrics)
+    return normalize_metrics(total_examples, aggregated_metrics)
+
+
+def evaluate_metrics_aggregation_fn(all_client_metrics: List[Tuple[int, Metrics]]) -> Metrics:
+    # This function is run by the server to aggregate metrics returned by each clients evaluate function
+    # NOTE: The first value of the tuple is number of examples for FedAvg
+    total_examples, aggregated_metrics = metric_aggregation(all_client_metrics)
+    return normalize_metrics(total_examples, aggregated_metrics)
+
+
+def get_initial_model_parameters() -> Parameters:
+    # Initializing the model parameters on the server side.
+    # Currently uses the Pytorch default initialization for the model parameters.
+    initial_model = APFLModule(MnistNet())
+    return ndarrays_to_parameters([val.cpu().numpy() for _, val in initial_model.state_dict().items()])
+
+
+def fit_config(local_epochs: int, batch_size: int, n_server_rounds: int, current_round: int) -> Config:
+    return {
+        "local_epochs": local_epochs,
+        "batch_size": batch_size,
+        "n_server_rounds": n_server_rounds,
+    }
+
+
+def main(config: Dict[str, Any]) -> None:
+    # This function will be used to produce a config that is sent to each client to initialize their own environment
+    fit_config_fn = partial(
+        fit_config,
+        config["local_epochs"],
+        config["batch_size"],
+        config["n_server_rounds"],
+    )
+
+    # Server performs simple FedAveraging as its server-side optimization strategy
+    strategy = FedAvg(
+        min_fit_clients=config["n_clients"],
+        min_evaluate_clients=config["n_clients"],
+        # Server waits for min_available_clients before starting FL rounds
+        min_available_clients=config["n_clients"],
+        on_fit_config_fn=fit_config_fn,
+        fit_metrics_aggregation_fn=fit_metrics_aggregation_fn,
+        evaluate_metrics_aggregation_fn=evaluate_metrics_aggregation_fn,
+        initial_parameters=get_initial_model_parameters(),
+    )
+
+    fl.server.start_server(
+        server_address="0.0.0.0:8080",
+        config=fl.server.ServerConfig(num_rounds=config["n_server_rounds"]),
+        strategy=strategy,
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="FL Server Main")
+    parser.add_argument(
+        "--config_path",
+        action="store",
+        type=str,
+        help="Path to configuration file.",
+        default="config.yaml",
+    )
+    args = parser.parse_args()
+
+    config = load_config(args.config_path)
+
+    main(config)

examples/fenda_example/README.md

@@ -1,14 +1,13 @@
 # FENDA Federated Learning Example
 This example provides an example of training a FENDA type model on a non-IID subset of the MNIST data. The FL server
 expects two clients to be spun up (i.e. it will wait until two clients report in before starting training). Each client
-has a modified version of the MNIST dataset. This modification essentially subsamples certain number from the original
+has a modified version of the MNIST dataset. This modification essentially subsamples a certain number from the original
 training and validation sets of MNIST in order to synthetically induce local variations in the statistical properties
 of the clients training/validation data. In theory, the models should be able to perform well on their local data
 while learning from other clients data that has different statistical properties. The subsampling is specified by
 sending a list of integers between 0-9 to the clients when they are run with the argument `--minority_numbers`.
 
-The server has some custom metrics aggregation and using Federated Averaging as its server-side optimization. The
-implementation uses a special type of weight exchange based on named-layer identification.
+The server has some custom metrics aggregation and uses Federated Averaging as its server-side optimization. The implementation uses a special type of weight exchange based on named-layer identification.
 
 # Running the Example
 In order to run the example, first ensure you have the virtual env of your choice activated and run

examples/fenda_example/client.py

@@ -10,11 +10,12 @@
 from flwr.common.typing import Config, NDArrays, Scalar
 from torch.utils.data import DataLoader
 
-from examples.fenda_example.client_data import load_data
 from examples.models.fenda_cnn import FendaClassifier, GlobalCnn, LocalCnn
 from fl4health.clients.numpy_fl_client import NumpyFlClient
 from fl4health.model_bases.fenda_base import FendaJoinMode, FendaModel
 from fl4health.parameter_exchange.layer_exchanger import FixedLayerExchanger
+from fl4health.utils.load_data import load_mnist_data
+from fl4health.utils.sampler import MinorityLabelBasedSampler
 
 
 def train(
@@ -120,9 +121,9 @@ def setup_client(self, config: Config) -> None:
         batch_size = self.narrow_config_type(config, "batch_size", int)
         downsample_percentage = self.narrow_config_type(config, "downsampling_ratio", float)
 
-        train_loader, validation_loader, num_examples = load_data(
-            self.data_path, batch_size, downsample_percentage, self.minority_numbers
-        )
+        sampler = MinorityLabelBasedSampler(list(range(10)), downsample_percentage, self.minority_numbers)
+
+        train_loader, validation_loader, num_examples = load_mnist_data(self.data_path, batch_size, sampler)
 
         self.train_loader = train_loader
         self.validation_loader = validation_loader

examples/fenda_example/client_data.py

@@ -21,3 +21,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = F.relu(self.fc2(x))
         x = self.fc3(x)
         return x
+
+
+class MnistNet(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 8, 5)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(8, 16, 5)
+        self.fc1 = nn.Linear(16 * 4 * 4, 120)
+        self.fc2 = nn.Linear(120, 10)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = x.view(-1, 16 * 4 * 4)
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        return x