Merge pull request #12 from VectorInstitute/dbe/fenda-implementation-other-fixes

FENDA model implementation and other fixes

Author: jewelltaylor

.gitignore

@@ -139,5 +139,6 @@ settings.json
 .DS_Store
 
 #datasets
-**/cifar_data/**
-**/news_classification/partitioned_datasets/**
+**/datasets/cifar_data/**
+**/datasets/news_classification/partitioned_datasets/**
+**/datasets/mnist_data/**

.pre-commit-config.yaml

@@ -31,9 +31,12 @@ repos:
       - id: isort
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.942
+    rev: v0.991
     hooks:
       - id: mypy
+        name: mypy
+        entry: ./run_mypy.sh
+        language: system
 
   - repo: https://github.com/nbQA-dev/nbQA
     rev: 1.3.1

examples/basic_example/README.md

@@ -32,4 +32,8 @@ clients. This is done by simply running (remembering to activate your environmen
 ```
 python -m examples.basic_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 both clients have been started federated learning should commence.

examples/basic_example/client.py

@@ -1,5 +1,4 @@
 import argparse
-from collections import OrderedDict
 from logging import INFO
 from pathlib import Path
 from typing import Dict, Tuple
@@ -14,6 +13,8 @@
 from torchvision.datasets import CIFAR10
 
 from examples.models.cnn_model import Net
+from fl4health.clients.numpy_fl_client import NumpyFlClient
+from fl4health.parameter_exchange.full_exchanger import FullParameterExchanger
 
 
 def load_data(data_dir: Path, batch_size: int) -> Tuple[DataLoader, DataLoader, Dict[str, int]]:
@@ -96,35 +97,19 @@ def validate(
     return loss / n_batches, accuracy
 
 
-class CifarClient(fl.client.NumPyClient):
-    def __init__(
-        self,
-        data_path: Path,
-        device: torch.device,
-    ) -> None:
-        self.data_path = data_path
-        self.device = device
-        self.initialized = False
-
-    def get_parameters(self, config: Config) -> NDArrays:
-        # Determines which weights are sent back to the server for aggregation.
-        # Currently sending all of them ordered by state_dict keys
-        # NOTE: Order matters, because it is relied upon by set_parameters below
-        return [val.cpu().numpy() for _, val in self.model.state_dict().items()]
-
-    def set_parameters(self, parameters: NDArrays, config: Config) -> None:
-        # Sets the local model parameters transfered from the server. The state_dict is
-        # reconstituted because parameters is simply a list of bytes
-        params_dict = zip(self.model.state_dict().keys(), parameters)
-        state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
-        self.model.load_state_dict(state_dict, strict=True)
+class CifarClient(NumpyFlClient):
+    def __init__(self, data_path: Path, device: torch.device) -> None:
+        super().__init__(data_path, device)
+        self.model = Net().to(self.device)
+        self.parameter_exchanger = FullParameterExchanger()
 
     def fit(self, parameters: NDArrays, config: Config) -> Tuple[NDArrays, int, Dict[str, Scalar]]:
         if not self.initialized:
             self.setup_client(config)
-
         self.set_parameters(parameters, config)
-        accuracy = train(self.model, self.train_loader, epochs=config["local_epochs"], device=self.device)
+
+        local_epochs = self.narrow_config_type(config, "local_epochs", int)
+        accuracy = train(self.model, self.train_loader, epochs=local_epochs, device=self.device)
         # FitRes should contain local parameters, number of examples on client, and a dictionary holding metrics
         # calculation results.
         return (
@@ -145,16 +130,15 @@ def evaluate(self, parameters: NDArrays, config: Config) -> Tuple[float, int, Di
         )
 
     def setup_client(self, config: Config) -> None:
+        super().setup_client(config)
+        batch_size = self.narrow_config_type(config, "batch_size", int)
 
-        train_loader, validation_loader, num_examples = load_data(self.data_path, config["batch_size"])
+        train_loader, validation_loader, num_examples = load_data(self.data_path, batch_size)
 
         self.train_loader = train_loader
         self.validation_loader = validation_loader
         self.num_examples = num_examples
 
-        model = Net().to(self.device)
-        self.model = model
-
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="FL Client Main")

examples/basic_example/server.py

@@ -8,30 +8,10 @@
 from flwr.server.strategy import FedAvg
 
 from examples.models.cnn_model import Net
+from examples.simple_metric_aggregation import metric_aggregation, normalize_metrics
 from fl4health.utils.config import load_config
 
 
-def metric_aggregation(all_client_metrics: List[Tuple[int, Metrics]]) -> Tuple[int, Metrics]:
-    aggregated_metrics: Metrics = {}
-    total_examples = 0
-    # Run through all of the metrics
-    for num_examples_on_client, client_metrics in all_client_metrics:
-        total_examples += num_examples_on_client
-        for metric_name, metric_value in client_metrics.items():
-            # Here we assume each metric is normalized by the number of examples on the client. So we scale up to
-            # get the "raw" value
-            if metric_name in aggregated_metrics:
-                aggregated_metrics[metric_name] += num_examples_on_client * metric_value
-            else:
-                aggregated_metrics[metric_name] = num_examples_on_client * metric_value
-    return total_examples, aggregated_metrics
-
-
-def normalize_metrics(total_examples: int, aggregated_metrics: Metrics) -> Metrics:
-    # Normalize all metric values by the total count of examples seen.
-    return {metric_name: metric_value / total_examples for metric_name, metric_value in aggregated_metrics.items()}
-
-
 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
@@ -56,9 +36,9 @@ def get_initial_model_parameters() -> Parameters:
 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}
+    return {"local_epochs": local_epochs, "batch_size": batch_size, "current_round": current_round}
 
 
 def main(config: Dict[str, Any]) -> None:

examples/docker_basic_example/fl_client/client.py

@@ -1,5 +1,4 @@
 import argparse
-from collections import OrderedDict
 from logging import INFO
 from pathlib import Path
 from typing import Dict, Tuple
@@ -14,6 +13,8 @@
 from torchvision.datasets import CIFAR10
 
 from examples.docker_basic_example.model import Net
+from fl4health.clients.numpy_fl_client import NumpyFlClient
+from fl4health.parameter_exchange.full_exchanger import FullParameterExchanger
 
 
 def load_data(data_dir: Path, batch_size: int) -> Tuple[DataLoader, DataLoader, Dict[str, int]]:
@@ -96,36 +97,19 @@ def validate(
     return loss / n_batches, accuracy
 
 
-class CifarClient(fl.client.NumPyClient):
-    def __init__(
-        self,
-        data_path: Path,
-        device: torch.device,
-    ) -> None:
-
-        self.data_path = data_path
-        self.device = device
-        self.initialized = False
-
-    def get_parameters(self, config: Config) -> NDArrays:
-        # Determines which weights are sent back to the server for aggregation.
-        # Currently sending all of them ordered by state_dict keys
-        # NOTE: Order matters, because it is relied upon by set_parameters below
-        return [val.cpu().numpy() for _, val in self.model.state_dict().items()]
-
-    def set_parameters(self, parameters: NDArrays, config: Config) -> None:
-        # Sets the local model parameters transfered from the server. The state_dict is
-        # reconstituted because parameters is simply a list of bytes
-        params_dict = zip(self.model.state_dict().keys(), parameters)
-        state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
-        self.model.load_state_dict(state_dict, strict=True)
+class CifarClient(NumpyFlClient):
+    def __init__(self, data_path: Path, device: torch.device) -> None:
+        super().__init__(data_path, device)
+        self.model = Net()
+        self.parameter_exchanger = FullParameterExchanger()
 
     def fit(self, parameters: NDArrays, config: Config) -> Tuple[NDArrays, int, Dict[str, Scalar]]:
         if not self.initialized:
             self.setup_client(config)
-
         self.set_parameters(parameters, config)
-        accuracy = train(self.model, self.train_loader, epochs=config["local_epochs"], device=self.device)
+
+        local_epochs = self.narrow_config_type(config, "local_epochs", int)
+        accuracy = train(self.model, self.train_loader, epochs=local_epochs, device=self.device)
         # FitRes should contain local parameters, number of examples on client, and a dictionary holding metrics
         # calculation results.
         return (
@@ -146,16 +130,14 @@ def evaluate(self, parameters: NDArrays, config: Config) -> Tuple[float, int, Di
         )
 
     def setup_client(self, config: Config) -> None:
-
-        train_loader, validation_loader, num_examples = load_data(self.data_path, config["batch_size"])
+        super().setup_client(config)
+        batch_size = self.narrow_config_type(config, "batch_size", int)
+        train_loader, validation_loader, num_examples = load_data(self.data_path, batch_size)
 
         self.train_loader = train_loader
         self.validation_loader = validation_loader
         self.num_examples = num_examples
 
-        model = Net()
-        self.model = model
-
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="FL Client Main")

examples/docker_basic_example/fl_server/server.py

@@ -8,30 +8,10 @@
 from flwr.server.strategy import FedAvg
 
 from examples.docker_basic_example.model import Net
+from examples.simple_metric_aggregation import metric_aggregation, normalize_metrics
 from fl4health.utils.config import load_config
 
 
-def metric_aggregation(all_client_metrics: List[Tuple[int, Metrics]]) -> Tuple[int, Metrics]:
-    aggregated_metrics: Metrics = {}
-    total_examples = 0
-    # Run through all of the metrics
-    for num_examples_on_client, client_metrics in all_client_metrics:
-        total_examples += num_examples_on_client
-        for metric_name, metric_value in client_metrics.items():
-            # Here we assume each metric is normalized by the number of examples on the client. So we scale up to
-            # get the "raw" value
-            if metric_name in aggregated_metrics:
-                aggregated_metrics[metric_name] += num_examples_on_client * metric_value
-            else:
-                aggregated_metrics[metric_name] = num_examples_on_client * metric_value
-    return total_examples, aggregated_metrics
-
-
-def normalize_metrics(total_examples: int, aggregated_metrics: Metrics) -> Metrics:
-    # Normalize all metric values by the total count of examples seen.
-    return {metric_name: metric_value / total_examples for metric_name, metric_value in aggregated_metrics.items()}
-
-
 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
@@ -49,9 +29,9 @@ def evaluate_metrics_aggregation_fn(all_client_metrics: List[Tuple[int, Metrics]
 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}
+    return {"local_epochs": local_epochs, "batch_size": batch_size, "current_round": current_round}
 
 
 def get_initial_model_parameters() -> Parameters:

examples/docker_basic_example/model.py

@@ -5,7 +5,7 @@
 
 class Net(nn.Module):
     def __init__(self) -> None:
-        super(Net, self).__init__()
+        super().__init__()
         self.conv1 = nn.Conv2d(3, 6, 5)
         self.pool = nn.MaxPool2d(2, 2)
         self.conv2 = nn.Conv2d(6, 16, 5)

examples/dp_fed_examples/client_level_dp/README.md

@@ -23,4 +23,8 @@ Once the server has started and logged "FL starting," the next step, in separate
 ```
 python -m examples.dp_fed_examples.client_level_dp.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 the minimum number of clients have been started federated learning should commence.

examples/dp_fed_examples/client_level_dp/client.py

@@ -1,11 +1,9 @@
 import argparse
-from collections import OrderedDict
 from logging import INFO
 from pathlib import Path
 from typing import Dict, Tuple
 
 import flwr as fl
-import numpy as np
 import torch
 import torch.nn as nn
 import torchvision.transforms as transforms
@@ -16,6 +14,7 @@
 
 from examples.models.cnn_model import Net
 from fl4health.clients.clipping_client import NumpyClippingClient
+from fl4health.parameter_exchange.full_exchanger import FullParameterExchanger
 
 
 def load_data(data_dir: Path, batch_size: int) -> Tuple[DataLoader, DataLoader, Dict[str, int]]:
@@ -98,57 +97,25 @@ def validate(
 
 
 class CifarClient(NumpyClippingClient):
-    def __init__(
-        self,
-        data_path: Path,
-        device: torch.device,
-    ) -> None:
-        super().__init__()
-        self.device = device
-        self.data_path = data_path
-        self.initialized = False
-        self.train_loader: DataLoader
-
-    def get_parameters(self, config: Config) -> NDArrays:
-        # Determines which weights are sent back to the server for aggregation.
-        # Currently sending all of them ordered by state_dict keys
-        # NOTE: Order matters, because it is relied upon by set_parameters below
-        model_weights = [val.cpu().numpy() for _, val in self.model.state_dict().items()]
-        # Clipped the weights and store clipping information in parameters
-        clipped_weight_update, clipping_bit = self.compute_weight_update_and_clip(model_weights)
-        return clipped_weight_update + [np.array([clipping_bit])]
-
-    def set_parameters(self, parameters: NDArrays, config: Config) -> None:
-        # Sets the local model parameters transfered from the server. The state_dict is
-        # reconstituted because parameters is simply a list of bytes
-        # The last entry in the parameters list is assumed to be a clipping bound (even if we're evaluating)
-        server_model_parameters = parameters[:-1]
-        params_dict = zip(self.model.state_dict().keys(), server_model_parameters)
-        state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
-        self.model.load_state_dict(state_dict, strict=True)
-
-        # Store the starting parameters without clipping bound before client optimization steps
-        self.current_weights = server_model_parameters
-
-        # Expectation is that the last entry in the parameters NDArrays is a clipping bound
-        clipping_bound = parameters[-1]
-        self.clipping_bound = float(clipping_bound)
+    def __init__(self, data_path: Path, device: torch.device) -> None:
+        super().__init__(data_path, device)
+        self.model = Net().to(self.device)
+        self.parameter_exchanger = FullParameterExchanger()
 
     def setup_client(self, config: Config) -> None:
-        self.batch_size = config["batch_size"]
-        self.local_epochs = config["local_epochs"]
-        self.adaptive_clipping = config["adaptive_clipping"]
+        super().setup_client(config)
+        self.batch_size = self.narrow_config_type(config, "batch_size", int)
+        self.local_epochs = self.narrow_config_type(config, "local_epochs", int)
+        # Server explicitly sets the clipping strategy
+        self.adaptive_clipping = self.narrow_config_type(config, "adaptive_clipping", bool)
 
         train_loader, validation_loader, num_examples = load_data(self.data_path, self.batch_size)
 
         self.train_loader = train_loader
         self.validation_loader = validation_loader
         self.num_examples = num_examples
-        self.model = Net().to(self.device)
-        self.initialized = True
 
     def fit(self, parameters: NDArrays, config: Config) -> Tuple[NDArrays, int, Dict[str, Scalar]]:
-
         self.set_parameters(parameters, config)
         accuracy = train(
             self.model,