Merge pull request #10 from VectorInstitute/weighted-dp-fedavg

Weighted dp fedavg

Author: jewelltaylor

examples/datasets/breast_cancer_data/hospital_0.csv

@@ -6,7 +6,7 @@ n_server_rounds: 20
 # much more noise can kill server side convergence.
 server_noise_multiplier: 0.01
 n_clients: 3
-client_sampling: 0.66
+client_sampling_rate: 0.66
 server_learning_rate: 1.0
 server_momentum: 0.2
 

examples/datasets/breast_cancer_data/hospital_1.csv

@@ -86,14 +86,16 @@ def main(config: Dict[str, Any]) -> None:
     client_manager = PoissonSamplingClientManager()
 
     # Accountant that computes the privacy through training
-    accountant = FlClientLevelAccountantPoissonSampling(config["client_sampling"], config["server_noise_multiplier"])
+    accountant = FlClientLevelAccountantPoissonSampling(
+        config["client_sampling_rate"], config["server_noise_multiplier"]
+    )
     target_delta = 1.0 / config["n_clients"]
     epsilon = accountant.get_epsilon(config["n_server_rounds"], target_delta)
     log(INFO, f"Model privacy after full training will be ({epsilon}, {target_delta})")
 
     # Server performs simple FedAveraging as it's server-side optimization strategy
     strategy = ClientLevelDPFedAvgM(
-        fraction_fit=config["client_sampling"],
+        fraction_fit=config["client_sampling_rate"],
         # Server waits for min_available_clients before starting FL rounds
         min_available_clients=config["n_clients"],
         fit_metrics_aggregation_fn=fit_metrics_aggregation_fn,

examples/datasets/breast_cancer_data/hospital_2.csv

@@ -0,0 +1,28 @@
+# Client Level Differential Privacy Federated Learning Example
+
+This example shows how to implement Differential Privacy into the Federated Learning framework. In this case we focus on *client level* privacy which is a more substantial version of instance level DP, where the participation of an entire client's set of data is protected from training dataset membership inference. This example uses the FedAvgM implementation with weighted averaging suggested in Learning Differentially Private Recurrent Language Models along with the adaptive clipping scheme proposed in Differentially Private Learning with Adaptive Clipping. The example uses an accountant specifically tailered to this approach. The clients are Poisson sampled by default.
+
+The example involves collaboratively learning a logistic regression model across multiple hospitals to classify breast cancer given 31 features. The dataset is sourced from [kaggle](https://www.kaggle.com/competitions/breast-cancer-classification/overview). A processed federated version of the dataset is available in the repository.
+
+# 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.dp_fed_examples.client_level_dp_weighted.server --config_path examples/dp_fed_examples/client_level_dp_weighted/config.yaml
+```
+
+## Starting Clients
+
+Once the server has started and logged "FL starting," the next step, in separate terminals, is to start the clients expected by the server. This is done by simply running (remembering to activate your environment)
+```
+python -m examples.dp_fed_examples.client_level_dp_weighted.client --dataset_path examples/datasets/breast_cancer_data/hospital_#.csv
+```
+After the minimum number of clients have been started federated learning should commence.

examples/dp_fed_examples/client_level_dp/config.yaml

@@ -0,0 +1,169 @@
+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
+from flwr.common.logger import log
+from flwr.common.typing import Config, NDArrays, Scalar
+from torch.utils.data import DataLoader
+
+from examples.dp_fed_examples.client_level_dp_weighted.data import load_data
+from examples.models.logistic_regression import LogisticRegression
+from fl4health.clients.clipping_client import NumpyClippingClient
+
+
+def train(net: nn.Module, train_loader: DataLoader, epochs: int, device: torch.device = torch.device("cpu")) -> float:
+
+    criterion = torch.nn.BCELoss()
+    optimizer = torch.optim.SGD(net.parameters(), lr=0.01, weight_decay=1e-4)
+
+    for epoch in range(epochs):
+        correct, total, running_loss = 0, 0, 0.0
+        n_batches = len(train_loader)
+        for features, labels in train_loader:
+            features, labels = features.to(device), labels.to(device)
+            optimizer.zero_grad()
+            preds = net(features)
+            loss = criterion(preds, labels)
+            loss.backward()
+            optimizer.step()
+
+            running_loss += loss.item()
+            predicted = preds.data >= 0.5
+
+            total += labels.size(0)
+            correct += (predicted.int() == labels.int()).sum().item()
+
+        accuracy = correct / total
+        # Local client logging.
+        log(
+            INFO,
+            f"Epoch: {epoch}, Client Training Loss: {running_loss/n_batches},"
+            f" Client Training Accuracy: {accuracy}",
+        )
+    return accuracy
+
+
+def validate(
+    net: nn.Module,
+    validation_loader: DataLoader,
+    device: torch.device = torch.device("cpu"),
+) -> Tuple[float, float]:
+    """Validate the network on the entire validation set."""
+    criterion = torch.nn.BCELoss()
+    correct, total, loss = 0, 0, 0.0
+    with torch.no_grad():
+        n_batches = len(validation_loader)
+        for features, labels in validation_loader:
+            features, labels = features.to(device), labels.to(device)
+            preds = net(features)
+            loss += criterion(preds, labels).item()
+            predicted = preds.data >= 0.5
+            total += labels.size(0)
+            correct += (predicted.int() == labels.int()).sum().item()
+    accuracy = correct / total
+    # Local client logging.
+    log(INFO, f"Client Validation Loss: {loss/n_batches} Client Validation Accuracy: {accuracy}")
+    return loss / n_batches, accuracy
+
+
+class HospitalClient(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
+
+        clipping_bound = parameters[-1]
+        self.clipping_bound = float(clipping_bound)
+
+    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"]
+        self.scaler_bytes = config["scaler"]
+
+        train_loader, validation_loader, num_examples = load_data(self.data_path, self.batch_size, self.scaler_bytes)
+
+        self.train_loader = train_loader
+        self.validation_loader = validation_loader
+        self.num_examples = num_examples
+        self.model = LogisticRegression(input_dim=31, output_dim=1).to(self.device)
+        self.initialized = True
+
+    def fit(self, parameters: NDArrays, config: Config) -> Tuple[NDArrays, int, Dict[str, Scalar]]:
+        # Expectation is that the last entry in the parameters NDArrays is a clipping bound
+        if not self.initialized:
+            self.setup_client(config)
+        self.set_parameters(parameters, config)
+        accuracy = train(
+            self.model,
+            self.train_loader,
+            self.local_epochs,
+            self.device,
+        )
+        # FitRes should contain local parameters, number of examples on client, and a dictionary holding metrics
+        # calculation results.
+        return (
+            self.get_parameters(config),
+            self.num_examples["train_set"],
+            {"accuracy": accuracy},
+        )
+
+    def evaluate(self, parameters: NDArrays, config: Config) -> Tuple[float, int, Dict[str, Scalar]]:
+        # Expectation is that the last entry in the parameters NDArrays is a clipping bound (even if it isn't used
+        # for evaluation)
+        if not self.initialized:
+            self.setup_client(config)
+        self.set_parameters(parameters, config)
+        loss, accuracy = validate(self.model, self.validation_loader, device=self.device)
+        # EvaluateRes should return the loss, number of examples on client, and a dictionary holding metrics
+        # calculation results.
+        return (
+            loss,
+            self.num_examples["validation_set"],
+            {"accuracy": accuracy},
+        )
+
+
+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()
+
+    # Load model and data
+    DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    data_path = Path(args.dataset_path)
+    client = HospitalClient(data_path, DEVICE)
+    fl.client.start_numpy_client(server_address="0.0.0.0:8080", client=client)

examples/dp_fed_examples/client_level_dp/server.py

@@ -0,0 +1,27 @@
+# Server parameters
+n_server_rounds: 25
+
+# NOTE: This multiplier is small, yielding a vacuous epsilon for privacy. It is set to this small value for this
+# example due to the small number of clients (3, see below), which, when combined with the clipping implies that
+# much more noise can kill server side convergence.
+server_noise_multiplier: 0.01
+n_clients: 3
+client_sampling_rate: 0.667
+server_learning_rate: 1.0
+server_momentum: 1.0
+weighted_averaging: True
+
+# Client training parameters
+local_epochs: 1
+batch_size: 32
+total_samples: 569
+
+# Clipping settings for update and optionally
+# adaptive clipping
+adaptive_clipping: True
+clipping_bound: 0.1
+clipping_learning_rate: 0.5
+# NOTE: The noise multiplier here is just picked for convenience. The recommended heuristic is
+# expected clients per round/20
+clipping_bit_noise_multiplier: 0.5
+clipping_quantile: 0.5

examples/dp_fed_examples/client_level_dp_weighted/README.md

@@ -0,0 +1,41 @@
+import pickle
+from pathlib import Path
+from typing import Dict, Tuple
+
+import numpy as np
+import pandas as pd
+import torch
+from sklearn.preprocessing import MinMaxScaler
+from torch.utils.data import DataLoader, TensorDataset
+
+
+class Scaler:
+    def __init__(self) -> None:
+        self.scaler = MinMaxScaler()
+
+    def __call__(self, train_x: np.ndarray, val_x: np.ndarray) -> np.ndarray:
+        scaled_train_x = self.scaler.fit_transform(train_x)
+        scaled_val_x = self.scaler.transform(val_x)
+        return scaled_train_x, scaled_val_x
+
+
+def load_data(data_dir: Path, batch_size: int, scaler_bytes: bytes) -> Tuple[DataLoader, DataLoader, Dict[str, int]]:
+    data = pd.read_csv(data_dir, index_col=False)
+    features = data.loc[:, data.columns != "label"].values
+    labels = data["label"].values
+    n_samples = data.shape[0]
+
+    scaler = pickle.loads(scaler_bytes)
+    train_samples = int(n_samples * 0.8)
+    train_features, train_labels = features[:train_samples, :], labels[:train_samples]
+    val_features, val_labels = features[train_samples:, :], labels[train_samples:]
+    train_features, val_features = scaler(train_features, val_features)
+    train_X, train_Y = torch.from_numpy(train_features).float(), torch.from_numpy(train_labels).float()
+    val_X, val_Y = torch.from_numpy(val_features).float(), torch.from_numpy(val_labels).float()
+    train_ds, val_ds = TensorDataset(train_X, train_Y), TensorDataset(val_X, val_Y)
+    train_loader = DataLoader(train_ds, batch_size=batch_size, drop_last=True)
+    val_loader = DataLoader(val_ds, batch_size=batch_size, drop_last=True)
+
+    num_examples = {"train_set": train_samples, "validation_set": n_samples - train_samples}
+
+    return train_loader, val_loader, num_examples