examples/async and model_pruing passed ty checks; server aggregation migrated to the new strategy API in examples/async.

Author: baochunli

examples/async/fedasync/fedasync_server.py

@@ -13,6 +13,7 @@
 
 from plato.config import Config
 from plato.servers import fedavg
+from plato.servers.strategies import FedAsyncAggregationStrategy
 
 
 class Server(fedavg.Server):
@@ -26,99 +27,45 @@ def __init__(
         trainer=None,
         callbacks=None,
     ):
+        aggregation_strategy = FedAsyncAggregationStrategy()
+
         super().__init__(
             model=model,
             datasource=datasource,
             algorithm=algorithm,
             trainer=trainer,
             callbacks=callbacks,
+            aggregation_strategy=aggregation_strategy,
         )
 
-        # The hyperparameter of FedAsync with a range of (0, 1)
-        self.mixing_hyperparam = 1
-
-        # Whether adjust mixing hyperparameter after each round
-        self.adaptive_mixing = False
-
     def configure(self) -> None:
         """Configure the mixing hyperparameter for the server, as well as
         other parameters from the configuration file.
         """
         super().configure()
 
-        # Configuring the mixing hyperparameter for FedAsync
-        self.adaptive_mixing = (
-            hasattr(Config().server, "adaptive_mixing")
-            and Config().server.adaptive_mixing
-        )
-
         if not hasattr(Config().server, "mixing_hyperparameter"):
             logging.warning(
                 "FedAsync: Variable mixing hyperparameter is required for the FedAsync server."
             )
         else:
-            self.mixing_hyperparam = Config().server.mixing_hyperparameter
+            try:
+                mixing_hyperparam = float(Config().server.mixing_hyperparameter)
+            except (TypeError, ValueError):
+                logging.warning(
+                    "FedAsync: Invalid mixing hyperparameter. "
+                    "Unable to cast %s to float.",
+                    Config().server.mixing_hyperparameter,
+                )
+                return
 
-            if 0 < self.mixing_hyperparam < 1:
+            if 0 < mixing_hyperparam < 1:
                 logging.info(
                     "FedAsync: Mixing hyperparameter is set to %s.",
-                    self.mixing_hyperparam,
+                    mixing_hyperparam,
                 )
             else:
                 logging.warning(
                     "FedAsync: Invalid mixing hyperparameter. "
                     "The hyperparameter needs to be between 0 and 1 (exclusive)."
                 )
-
-    async def aggregate_weights(self, updates, baseline_weights, weights_received):
-        """Process the client reports by aggregating their weights."""
-        # Calculate the new mixing hyperparameter with client's staleness
-        client_staleness = updates[0].staleness
-
-        if self.adaptive_mixing:
-            self.mixing_hyperparam *= self._staleness_function(client_staleness)
-
-        return await self.algorithm.aggregate_weights(
-            baseline_weights, weights_received, mixing=self.mixing_hyperparam
-        )
-
-    @staticmethod
-    def _staleness_function(staleness) -> float:
-        """Staleness function used to adjust the mixing hyperparameter"""
-        if hasattr(Config().server, "staleness_weighting_function"):
-            staleness_func_param = Config().server.staleness_weighting_function
-            func_type = staleness_func_param.type.lower()
-            if func_type == "constant":
-                return Server._constant_function()
-            elif func_type == "polynomial":
-                a = staleness_func_param.a
-                return Server._polynomial_function(staleness, a)
-            elif func_type == "hinge":
-                a = staleness_func_param.a
-                b = staleness_func_param.b
-                return Server._hinge_function(staleness, a, b)
-            else:
-                logging.warning(
-                    "FedAsync: Unknown staleness weighting function type. "
-                    "Type needs to be constant, polynomial, or hinge."
-                )
-        else:
-            return Server.constant_function()
-
-    @staticmethod
-    def _constant_function() -> float:
-        """Constant staleness function as proposed in Sec. 5.2, Evaluation Setup."""
-        return 1
-
-    @staticmethod
-    def _polynomial_function(staleness, a) -> float:
-        """Polynomial staleness function as proposed in Sec. 5.2, Evaluation Setup."""
-        return (staleness + 1) ** -a
-
-    @staticmethod
-    def _hinge_function(staleness, a, b) -> float:
-        """Hinge staleness function as proposed in Sec. 5.2, Evaluation Setup."""
-        if staleness <= b:
-            return 1
-        else:
-            return 1 / (a * (staleness - b) + 1)

examples/async/port/port_server.py

@@ -7,139 +7,35 @@
 
 """
 
-import asyncio
-import copy
-import logging
-import os
-
-import torch
-import torch.nn.functional as F
-
-from plato.config import Config
 from plato.servers import fedavg
+from plato.servers.strategies import PortAggregationStrategy
 
 
 class Server(fedavg.Server):
-    """A federated learning server using the FedAsync algorithm."""
-
-    async def cosine_similarity(self, update, staleness):
-        """Compute the cosine similarity of the received updates and the difference
-        between the current and a previous model according to client staleness."""
-        # Loading the global model from a previous round according to staleness
-        filename = f"model_{self.current_round - 2}.pth"
-        model_path = Config().params["model_path"]
-        model_path = f"{model_path}/{filename}"
-
-        similarity = 1.0
-
-        if staleness > 1 and os.path.exists(model_path):
-            previous_model = copy.deepcopy(self.trainer.model)
-            previous_model.load_state_dict(torch.load(model_path))
-
-            previous = torch.zeros(0)
-            for __, weight in previous_model.cpu().state_dict().items():
-                previous = torch.cat((previous, weight.view(-1)))
-
-            current = torch.zeros(0)
-            for __, weight in self.trainer.model.cpu().state_dict().items():
-                current = torch.cat((current, weight.view(-1)))
-
-            deltas = torch.zeros(0)
-            for __, delta in update.items():
-                deltas = torch.cat((deltas, delta.view(-1)))
-
-            similarity = F.cosine_similarity(current - previous, deltas, dim=0)
-
-        return similarity
-
-    async def aggregate_deltas(self, updates, deltas_received):
-        """Aggregate weight updates from the clients using federated averaging."""
-        # Extract the total number of samples
-        self.total_samples = sum(update.report.num_samples for update in updates)
-
-        # Constructing the aggregation weights to be used
-        aggregation_weights = []
-
-        for i, update in enumerate(deltas_received):
-            report = updates[i].report
-            staleness = updates[i].staleness
-            num_samples = report.num_samples
-
-            similarity = await self.cosine_similarity(update, staleness)
-            staleness_factor = Server.staleness_function(staleness)
-
-            similarity_weight = (
-                Config().server.similarity_weight
-                if hasattr(Config().server, "similarity_weight")
-                else 1
-            )
-            staleness_weight = (
-                Config().server.staleness_weight
-                if hasattr(Config().server, "staleness_weight")
-                else 1
-            )
-
-            logging.info("[Client %s] similarity: %s", i, (similarity + 1) / 2)
-            logging.info(
-                "[Client %s] staleness: %s, staleness factor: %s",
-                i,
-                staleness,
-                staleness_factor,
-            )
-            raw_weight = (
-                num_samples
-                / self.total_samples
-                * (
-                    (similarity + 1) / 2 * similarity_weight
-                    + staleness_factor * staleness_weight
-                )
-            )
-            logging.info("[Client %s] raw weight = %s", i, raw_weight)
-
-            aggregation_weights.append(raw_weight)
-
-        # Normalize so that the sum of aggregation weights equals 1
-        aggregation_weights = [
-            i / sum(aggregation_weights) for i in aggregation_weights
-        ]
-
-        logging.info(
-            "[Server #%s] normalized aggregation weights: %s",
-            os.getpid(),
-            aggregation_weights,
+    """A federated learning server using the Port aggregation strategy."""
+
+    def __init__(
+        self,
+        model=None,
+        datasource=None,
+        algorithm=None,
+        trainer=None,
+        callbacks=None,
+    ):
+        super().__init__(
+            model=model,
+            datasource=datasource,
+            algorithm=algorithm,
+            trainer=trainer,
+            callbacks=callbacks,
+            aggregation_strategy=PortAggregationStrategy(),
         )
 
-        # Perform weighted averaging
-        avg_update = {
-            name: self.trainer.zeros(weights.shape)
-            for name, weights in deltas_received[0].items()
-        }
-
-        for i, update in enumerate(deltas_received):
-            for name, delta in update.items():
-                avg_update[name] += delta * aggregation_weights[i]
-
-            # Yield to other tasks in the server
-            await asyncio.sleep(0)
-
-        return avg_update
-
     def weights_aggregated(self, updates):
         """
         Method called at the end of aggregating received weights.
         """
         # Save the current model for later retrieval when cosine similarity needs to be computed
         filename = f"model_{self.current_round}.pth"
-        self.trainer.save_model(filename)
-
-    @staticmethod
-    def staleness_function(staleness):
-        """The staleness_function."""
-        staleness_bound = 10
-
-        if hasattr(Config().server, "staleness_bound"):
-            staleness_bound = Config().server.staleness_bound
-
-        staleness_factor = staleness_bound / (staleness + staleness_bound)
-
-        return staleness_factor
+        trainer = self.require_trainer()
+        trainer.save_model(filename)