FedYogi

src/py/flwr/server/strategy/__init__.py

@@ -23,6 +23,7 @@
 from .fedavg import FedAvg as FedAvg
 from .fedfs_v0 import FedFSv0 as FedFSv0
 from .fedfs_v1 import FedFSv1 as FedFSv1
+from .fedyogi import FedYogi as FedYogi
 from .qfedavg import QFedAvg as QFedAvg
 from .qfedavg import QffedAvg as QffedAvg  # Deprecated
 from .strategy import Strategy as Strategy
@@ -36,6 +37,7 @@
     "FedAvg",
     "FedFSv0",
     "FedFSv1",
+    "FedYogi",
     "QFedAvg",
     "QffedAvg",  # Deprecated
     "Strategy",

src/py/flwr/server/strategy/fedyogi.py

@@ -0,0 +1,166 @@
+# Copyright 2020 Adap GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Adaptive Federated Optimization using Yogi (FedYogi) [Reddi et al., 2020]
+strategy.
+
+Paper: https://arxiv.org/abs/2003.00295
+"""
+
+
+from typing import Callable, Dict, List, Optional, Tuple
+
+import numpy as np
+
+from flwr.common import (
+    FitRes,
+    Parameters,
+    Scalar,
+    Weights,
+    parameters_to_weights,
+    weights_to_parameters,
+)
+from flwr.server.client_proxy import ClientProxy
+
+from .fedopt import FedOpt
+
+
+class FedYogi(FedOpt):
+    """Adaptive Federated Optimization using Yogi (FedYogi) [Reddi et al.,
+    2020] strategy.
+
+    Paper: https://arxiv.org/abs/2003.00295
+    """
+
+    # pylint: disable=too-many-arguments,too-many-instance-attributes,too-many-locals
+    def __init__(
+        self,
+        *,
+        fraction_fit: float = 0.1,
+        fraction_eval: float = 0.1,
+        min_fit_clients: int = 2,
+        min_eval_clients: int = 2,
+        min_available_clients: int = 2,
+        eval_fn: Optional[
+            Callable[[Weights], Optional[Tuple[float, Dict[str, Scalar]]]]
+        ] = None,
+        on_fit_config_fn: Optional[Callable[[int], Dict[str, Scalar]]] = None,
+        on_evaluate_config_fn: Optional[Callable[[int], Dict[str, Scalar]]] = None,
+        accept_failures: bool = True,
+        initial_parameters: Parameters,
+        eta: float = 1e-2,
+        eta_l: float = 0.0316,
+        beta_1: float = 0.9,
+        beta_2: float = 0.99,
+        tau: float = 1e-3,
+    ) -> None:
+        """Federated learning strategy using Yogi on server-side.
+
+        Implementation based on https://arxiv.org/abs/2003.00295
+
+        Args:
+            fraction_fit (float, optional): Fraction of clients used during
+                training. Defaults to 0.1.
+            fraction_eval (float, optional): Fraction of clients used during
+                validation. Defaults to 0.1.
+            min_fit_clients (int, optional): Minimum number of clients used
+                during training. Defaults to 2.
+            min_eval_clients (int, optional): Minimum number of clients used
+                during validation. Defaults to 2.
+            min_available_clients (int, optional): Minimum number of total
+                clients in the system. Defaults to 2.
+            eval_fn (Callable[[Weights], Optional[Tuple[float, float]]], optional):
+                Function used for validation. Defaults to None.
+            on_fit_config_fn (Callable[[int], Dict[str, str]], optional):
+                Function used to configure training. Defaults to None.
+            on_evaluate_config_fn (Callable[[int], Dict[str, str]], optional):
+                Function used to configure validation. Defaults to None.
+            accept_failures (bool, optional): Whether or not accept rounds
+                containing failures. Defaults to True.
+            initial_parameters (Parameters): Initial set of parameters from the server.
+            eta (float, optional): Server-side learning rate. Defaults to 1e-1.
+            eta_l (float, optional): Client-side learning rate. Defaults to 1e-1.
+            beta_1 (float, optional): Momentum parameter. Defaults to 0.9.
+            beta_2 (float, optional): Second moment parameter. Defaults to 0.99.
+            tau (float, optional): Controls the algorithm's degree of adaptability.
+                Defaults to 1e-9.
+        """
+        super().__init__(
+            fraction_fit=fraction_fit,
+            fraction_eval=fraction_eval,
+            min_fit_clients=min_fit_clients,
+            min_eval_clients=min_eval_clients,
+            min_available_clients=min_available_clients,
+            eval_fn=eval_fn,
+            on_fit_config_fn=on_fit_config_fn,
+            on_evaluate_config_fn=on_evaluate_config_fn,
+            accept_failures=accept_failures,
+            initial_parameters=initial_parameters,
+            eta=eta,
+            eta_l=eta_l,
+            beta_1=beta_1,
+            beta_2=beta_2,
+            tau=tau,
+        )
+        self.delta_t: Optional[Weights] = None
+        self.v_t: Optional[Weights] = None
+
+    def __repr__(self) -> str:
+        rep = f"FedYogi(accept_failures={self.accept_failures})"
+        return rep
+
+    def aggregate_fit(
+        self,
+        rnd: int,
+        results: List[Tuple[ClientProxy, FitRes]],
+        failures: List[BaseException],
+    ) -> Tuple[Optional[Parameters], Dict[str, Scalar]]:
+        """Aggregate fit results using weighted average."""
+        fedavg_parameters_aggregated, metrics_aggregated = super().aggregate_fit(
+            rnd=rnd, results=results, failures=failures
+        )
+        if fedavg_parameters_aggregated is None:
+            return None, {}
+
+        fedavg_weights_aggregate = parameters_to_weights(fedavg_parameters_aggregated)
+        aggregated_updates = [
+            x - y for x, y in zip(fedavg_weights_aggregate, self.current_weights)
+        ]
+
+        # Yogi
+
+        if not self.delta_t:
+            self.delta_t = [np.zeros_like(x) for x in self.current_weights]
+
+        self.delta_t = [
+            self.beta_1 * x + (1.0 - self.beta_1) * y
+            for x, y in zip(self.delta_t, aggregated_updates)
+        ]
+
+        if not self.v_t:
+            self.v_t = [np.zeros_like(x) for x in self.delta_t]
+
+        self.v_t = [
+            x + (1.0 - self.beta_2) * np.multiply(y, y) * np.sign(x - np.multiply(y, y))
+            for x, y in zip(self.v_t, self.delta_t)
+        ]
+
+        new_weights = [
+            x + self.eta * y / (np.sqrt(z) + self.tau)
+            for x, y, z in zip(self.current_weights, self.delta_t, self.v_t)
+        ]
+
+        self.current_weights = new_weights
+
+        return weights_to_parameters(self.current_weights), metrics_aggregated