refactor(examples) Update opacus example (#4262)

Co-authored-by: jafermarq <javier@flower.ai>

examples/opacus/README.md

@@ -1,5 +1,5 @@
 ---
-tags: [dp, security, fds]
+tags: [DP, DP-SGD, basic, vision, fds, privacy]
 dataset: [CIFAR-10]
 framework: [opacus, torch]
 ---
@@ -10,57 +10,55 @@ In this example, we demonstrate how to train a model with differential privacy (
 
 For more information about DP in Flower please refer to the [tutorial](https://flower.ai/docs/framework/how-to-use-differential-privacy.html). For additional information about Opacus, visit the official [website](https://opacus.ai/).
 
-## Environments Setup
+## Set up the project
 
-Start by cloning the example. We prepared a single-line command that you can copy into your shell which will checkout the example for you:
+### Clone the project
+
+Start by cloning the example project:
 
 ```shell
-git clone --depth=1 https://github.com/adap/flower.git && mv flower/examples/opacus . && rm -rf flower && cd opacus
+git clone --depth=1 https://github.com/adap/flower.git \
+        && mv flower/examples/opacus . \
+        && rm -rf flower \
+        && cd opacus
 ```
 
 This will create a new directory called `opacus` containing the following files:
 
 ```shell
--- pyproject.toml
--- client.py
--- server.py
--- README.md
+opacus
+├── opacus_fl
+│   ├── client_app.py   # Defines your ClientApp
+│   ├── server_app.py   # Defines your ServerApp
+│   └── task.py         # Defines your model, training, and data loading
+├── pyproject.toml      # Project metadata like dependencies and configs
+└── README.md
 ```
 
-### Installing dependencies
+### Install dependencies and project
 
-Project dependencies are defined in `pyproject.toml`. Install them with:
+Install the dependencies defined in `pyproject.toml` as well as the `opacus_fl` package.
 
 ```shell
-pip install .
+# From a new python environment, run:
+pip install -e .
 ```
 
-## Run Flower with Opacus and Pytorch
-
-### 1. Start the long-running Flower server (SuperLink)
-
-```bash
-flower-superlink --insecure
-```
+## Run the project
 
-### 2. Start the long-running Flower clients (SuperNodes)
+You can run your Flower project in both _simulation_ and _deployment_ mode without making changes to the code. If you are starting with Flower, we recommend you using the _simulation_ mode as it requires fewer components to be launched manually. By default, `flwr run` will make use of the Simulation Engine.
 
-Start 2 Flower `SuperNodes` in 2 separate terminal windows, using:
+### Run with the Simulation Engine
 
 ```bash
-flower-client-app client:appA --insecure
+flwr run .
 ```
 
-```bash
-flower-client-app client:appB --insecure
-```
-
-Opacus hyperparameters can be passed for each client in `ClientApp` instantiation (in `client.py`). In this example, `noise_multiplier=1.5` and `noise_multiplier=1` are used for the first and second client respectively.
-
-### 3. Run the Flower App
-
-With both the long-running server (SuperLink) and two clients (SuperNode) up and running, we can now run the actual Flower App:
+You can also override some of the settings for your `ClientApp` and `ServerApp` defined in `pyproject.toml`. For example:
 
 ```bash
-flower-server-app server:app --insecure
+flwr run . --run-config "max-grad-norm=1.0 num-server-rounds=5"
 ```
+
+> \[!NOTE\]
+> Please note that, at the current state, users cannot set `NodeConfig` for simulated `ClientApp`s. For this reason, the hyperparameter `noise_multiplier` is set in the `client_fn` method based on a condition check on `partition_id`. This will be modified in a future version of Flower to allow users to set `NodeConfig` for simulated `ClientApp`s.

examples/opacus/opacus_fl/__init__.py

@@ -0,0 +1 @@
+"""opacus: Training with Sample-Level Differential Privacy using Opacus Privacy Engine."""

examples/opacus/opacus_fl/client_app.py

@@ -0,0 +1,92 @@
+"""opacus: Training with Sample-Level Differential Privacy using Opacus Privacy Engine."""
+
+import warnings
+
+import torch
+from opacus import PrivacyEngine
+from opacus_fl.task import Net, get_weights, load_data, set_weights, test, train
+import logging
+
+from flwr.client import ClientApp, NumPyClient
+from flwr.common import Context
+
+warnings.filterwarnings("ignore", category=UserWarning)
+
+
+class FlowerClient(NumPyClient):
+    def __init__(
+        self,
+        train_loader,
+        test_loader,
+        target_delta,
+        noise_multiplier,
+        max_grad_norm,
+    ) -> None:
+        super().__init__()
+        self.model = Net()
+        self.train_loader = train_loader
+        self.test_loader = test_loader
+        self.target_delta = target_delta
+        self.noise_multiplier = noise_multiplier
+        self.max_grad_norm = max_grad_norm
+
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    def fit(self, parameters, config):
+        model = self.model
+        set_weights(model, parameters)
+
+        optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
+
+        privacy_engine = PrivacyEngine(secure_mode=False)
+        (
+            model,
+            optimizer,
+            self.train_loader,
+        ) = privacy_engine.make_private(
+            module=model,
+            optimizer=optimizer,
+            data_loader=self.train_loader,
+            noise_multiplier=self.noise_multiplier,
+            max_grad_norm=self.max_grad_norm,
+        )
+
+        epsilon = train(
+            model,
+            self.train_loader,
+            privacy_engine,
+            optimizer,
+            self.target_delta,
+            device=self.device,
+        )
+
+        if epsilon is not None:
+            print(f"Epsilon value for delta={self.target_delta} is {epsilon:.2f}")
+        else:
+            print("Epsilon value not available.")
+
+        return (get_weights(model), len(self.train_loader.dataset), {})
+
+    def evaluate(self, parameters, config):
+        set_weights(self.model, parameters)
+        loss, accuracy = test(self.model, self.test_loader, self.device)
+        return loss, len(self.test_loader.dataset), {"accuracy": accuracy}
+
+
+def client_fn(context: Context):
+    partition_id = context.node_config["partition-id"]
+    noise_multiplier = 1.0 if partition_id % 2 == 0 else 1.5
+
+    train_loader, test_loader = load_data(
+        partition_id=partition_id, num_partitions=context.node_config["num-partitions"]
+    )
+    return FlowerClient(
+        train_loader,
+        test_loader,
+        context.run_config["target-delta"],
+        noise_multiplier,
+        context.run_config["max-grad-norm"],
+    ).to_client()
+
+
+app = ClientApp(client_fn=client_fn)

examples/opacus/opacus_fl/server_app.py

@@ -0,0 +1,37 @@
+"""opacus: Training with Sample-Level Differential Privacy using Opacus Privacy Engine."""
+
+import logging
+from typing import List, Tuple
+
+from opacus_fl.task import Net, get_weights
+
+from flwr.common import Context, Metrics, ndarrays_to_parameters
+from flwr.server import ServerApp, ServerAppComponents, ServerConfig
+from flwr.server.strategy import FedAvg
+
+# Opacus logger seems to change the flwr logger to DEBUG level. Set back to INFO
+logging.getLogger("flwr").setLevel(logging.INFO)
+
+
+def weighted_average(metrics: List[Tuple[int, Metrics]]) -> Metrics:
+    accuracies = [num_examples * m["accuracy"] for num_examples, m in metrics]
+    examples = [num_examples for num_examples, _ in metrics]
+    return {"accuracy": sum(accuracies) / sum(examples)}
+
+
+def server_fn(context: Context) -> ServerAppComponents:
+    num_rounds = context.run_config["num-server-rounds"]
+
+    ndarrays = get_weights(Net())
+    parameters = ndarrays_to_parameters(ndarrays)
+
+    strategy = FedAvg(
+        evaluate_metrics_aggregation_fn=weighted_average,
+        initial_parameters=parameters,
+    )
+    config = ServerConfig(num_rounds=num_rounds)
+
+    return ServerAppComponents(config=config, strategy=strategy)
+
+
+app = ServerApp(server_fn=server_fn)