plato-learn 1.1__py3-none-any.whl

plato/servers/fedavg.py

@@ -0,0 +1,281 @@
+"""
+A simple federated learning server using federated averaging.
+"""
+
+import asyncio
+import logging
+import os
+
+from plato.algorithms import registry as algorithms_registry
+from plato.config import Config
+from plato.datasources import registry as datasources_registry
+from plato.processors import registry as processor_registry
+from plato.samplers import all_inclusive
+from plato.servers import base
+from plato.trainers import registry as trainers_registry
+from plato.utils import csv_processor, fonts
+
+
+class Server(base.Server):
+    """Federated learning server using federated averaging."""
+
+    def __init__(
+        self, model=None, datasource=None, algorithm=None, trainer=None, callbacks=None
+    ):
+        super().__init__(callbacks=callbacks)
+
+        self.custom_model = model
+        self.model = None
+
+        self.custom_algorithm = algorithm
+        self.algorithm = None
+
+        self.custom_trainer = trainer
+        self.trainer = None
+
+        self.custom_datasource = datasource
+        self.datasource = None
+
+        self.testset = None
+        self.testset_sampler = None
+        self.total_samples = 0
+
+        self.total_clients = Config().clients.total_clients
+        self.clients_per_round = Config().clients.per_round
+
+        logging.info(
+            "[Server #%d] Started training on %d clients with %d per round.",
+            os.getpid(),
+            self.total_clients,
+            self.clients_per_round,
+        )
+
+    def configure(self) -> None:
+        """
+        Booting the federated learning server by setting up the data, model, and
+        creating the clients.
+        """
+        super().configure()
+
+        total_rounds = Config().trainer.rounds
+        target_accuracy = None
+        target_perplexity = None
+
+        if hasattr(Config().trainer, "target_accuracy"):
+            target_accuracy = Config().trainer.target_accuracy
+        elif hasattr(Config().trainer, "target_perplexity"):
+            target_perplexity = Config().trainer.target_perplexity
+
+        if target_accuracy:
+            logging.info(
+                "Training: %s rounds or accuracy above %.1f%%\n",
+                total_rounds,
+                100 * target_accuracy,
+            )
+        elif target_perplexity:
+            logging.info(
+                "Training: %s rounds or perplexity below %.1f\n",
+                total_rounds,
+                target_perplexity,
+            )
+        else:
+            logging.info("Training: %s rounds\n", total_rounds)
+
+        self.init_trainer()
+
+        # Prepares this server for processors that processes outbound and inbound
+        # data payloads
+        self.outbound_processor, self.inbound_processor = processor_registry.get(
+            "Server", server_id=os.getpid(), trainer=self.trainer
+        )
+
+        if not (hasattr(Config().server, "do_test") and not Config().server.do_test):
+            if self.datasource is None and self.custom_datasource is None:
+                self.datasource = datasources_registry.get(client_id=0)
+            elif self.datasource is None and self.custom_datasource is not None:
+                self.datasource = self.custom_datasource()
+
+            self.testset = self.datasource.get_test_set()
+            if hasattr(Config().data, "testset_size"):
+                self.testset_sampler = all_inclusive.Sampler(
+                    self.datasource, testing=True
+                )
+
+        # Initialize the test accuracy csv file if clients compute locally
+        if (
+            hasattr(Config().clients, "do_test")
+            and Config().clients.do_test
+            and (
+                hasattr(Config(), "results")
+                and hasattr(Config().results, "record_clients_accuracy")
+                and Config().results.record_clients_accuracy
+            )
+        ):
+            accuracy_csv_file = (
+                f"{Config().params['result_path']}/{os.getpid()}_accuracy.csv"
+            )
+            accuracy_headers = ["round", "client_id", "accuracy"]
+            csv_processor.initialize_csv(
+                accuracy_csv_file, accuracy_headers, Config().params["result_path"]
+            )
+
+    def init_trainer(self) -> None:
+        """Setting up the global model, trainer, and algorithm."""
+        if self.model is None and self.custom_model is not None:
+            self.model = self.custom_model
+
+        if self.trainer is None and self.custom_trainer is None:
+            self.trainer = trainers_registry.get(model=self.model)
+        elif self.trainer is None and self.custom_trainer is not None:
+            self.trainer = self.custom_trainer(model=self.model)
+
+        if self.algorithm is None and self.custom_algorithm is None:
+            self.algorithm = algorithms_registry.get(trainer=self.trainer)
+        elif self.algorithm is None and self.custom_algorithm is not None:
+            self.algorithm = self.custom_algorithm(trainer=self.trainer)
+
+    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)
+
+        # Perform weighted averaging
+        avg_update = {
+            name: self.trainer.zeros(delta.shape)
+            for name, delta in deltas_received[0].items()
+        }
+
+        for i, update in enumerate(deltas_received):
+            report = updates[i].report
+            num_samples = report.num_samples
+
+            for name, delta in update.items():
+                # Use weighted average by the number of samples
+                avg_update[name] += delta * (num_samples / self.total_samples)
+
+            # Yield to other tasks in the server
+            await asyncio.sleep(0)
+
+        return avg_update
+
+    async def _process_reports(self):
+        """Process the client reports by aggregating their weights."""
+        weights_received = [update.payload for update in self.updates]
+
+        weights_received = self.weights_received(weights_received)
+        self.callback_handler.call_event("on_weights_received", self, weights_received)
+
+        # Extract the current model weights as the baseline
+        baseline_weights = self.algorithm.extract_weights()
+
+        if hasattr(self, "aggregate_weights"):
+            # Runs a server aggregation algorithm using weights rather than deltas
+            logging.info(
+                "[Server #%d] Aggregating model weights directly rather than weight deltas.",
+                os.getpid(),
+            )
+            updated_weights = await self.aggregate_weights(
+                self.updates, baseline_weights, weights_received
+            )
+
+            # Loads the new model weights
+            self.algorithm.load_weights(updated_weights)
+        else:
+            # Computes the weight deltas by comparing the weights received with
+            # the current global model weights
+            deltas_received = self.algorithm.compute_weight_deltas(
+                baseline_weights, weights_received
+            )
+            # Runs a framework-agnostic server aggregation algorithm, such as
+            # the federated averaging algorithm
+            logging.info("[Server #%d] Aggregating model weight deltas.", os.getpid())
+            deltas = await self.aggregate_deltas(self.updates, deltas_received)
+            # Updates the existing model weights from the provided deltas
+            updated_weights = self.algorithm.update_weights(deltas)
+            # Loads the new model weights
+            self.algorithm.load_weights(updated_weights)
+
+        # The model weights have already been aggregated, now calls the
+        # corresponding hook and callback
+        self.weights_aggregated(self.updates)
+        self.callback_handler.call_event("on_weights_aggregated", self, self.updates)
+
+        # Testing the global model accuracy
+        if hasattr(Config().server, "do_test") and not Config().server.do_test:
+            # Compute the average accuracy from client reports
+            self.accuracy, self.accuracy_std = self.get_accuracy_mean_std(self.updates)
+            logging.info(
+                "[%s] Average client accuracy: %.2f%%.", self, 100 * self.accuracy
+            )
+        else:
+            # Testing the updated model directly at the server
+            logging.info("[%s] Started model testing.", self)
+            self.accuracy = self.trainer.test(self.testset, self.testset_sampler)
+
+        if hasattr(Config().trainer, "target_perplexity"):
+            logging.info(
+                fonts.colourize(
+                    f"[{self}] Global model perplexity: {self.accuracy:.2f}\n"
+                )
+            )
+        else:
+            logging.info(
+                fonts.colourize(
+                    f"[{self}] Global model accuracy: {100 * self.accuracy:.2f}%\n"
+                )
+            )
+
+        self.clients_processed()
+        self.callback_handler.call_event("on_clients_processed", self)
+
+    def clients_processed(self) -> None:
+        """Additional work to be performed after client reports have been processed."""
+
+    def get_logged_items(self) -> dict:
+        """Get items to be logged by the LogProgressCallback class in a .csv file."""
+        return {
+            "round": self.current_round,
+            "accuracy": self.accuracy,
+            "accuracy_std": self.accuracy_std,
+            "elapsed_time": self.wall_time - self.initial_wall_time,
+            "processing_time": max(
+                update.report.processing_time for update in self.updates
+            ),
+            "comm_time": max(update.report.comm_time for update in self.updates),
+            "round_time": max(
+                update.report.training_time
+                + update.report.processing_time
+                + update.report.comm_time
+                for update in self.updates
+            ),
+            "comm_overhead": self.comm_overhead,
+        }
+
+    @staticmethod
+    def get_accuracy_mean_std(updates):
+        """Compute the accuracy mean and standard deviation across clients."""
+        # Get total number of samples
+        total_samples = sum(update.report.num_samples for update in updates)
+
+        # Perform weighted averaging
+        updates_accuracy = [update.report.accuracy for update in updates]
+        weights = [update.report.num_samples / total_samples for update in updates]
+
+        mean = sum(acc * weights[idx] for idx, acc in enumerate(updates_accuracy))
+        variance = sum(
+            (acc - mean) ** 2 * weights[idx] for idx, acc in enumerate(updates_accuracy)
+        )
+        std = variance**0.5
+
+        return mean, std
+
+    def weights_received(self, weights_received):
+        """
+        Method called after the updated weights have been received.
+        """
+        return weights_received
+
+    def weights_aggregated(self, updates):
+        """
+        Method called after the updated weights have been aggregated.
+        """

plato/servers/fedavg_cs.py

@@ -0,0 +1,335 @@
+"""
+A cross-silo federated learning server using federated averaging, as either edge or central servers.
+"""
+
+import asyncio
+import logging
+import os
+import numpy as np
+
+from plato.config import Config
+from plato.datasources import registry as datasources_registry
+from plato.processors import registry as processor_registry
+from plato.samplers import registry as samplers_registry
+from plato.samplers import all_inclusive
+from plato.servers import fedavg
+from plato.utils import fonts
+
+
+class Server(fedavg.Server):
+    """Cross-silo federated learning server using federated averaging."""
+
+    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,
+        )
+
+        self.current_global_round = 0
+        self.average_accuracy = 0
+        self.std_accuracy = 0
+
+        if Config().is_edge_server():
+            # An edge client waits for the event that a certain number of
+            # aggregations are completed
+            self.model_aggregated = asyncio.Event()
+
+            # An edge client waits for the event that a new global round begins
+            # before starting the first round of local aggregation
+            self.new_global_round_begins = asyncio.Event()
+
+            edge_server_id = Config().args.id - Config().clients.total_clients
+
+            # Compute the total number of clients in each silo for edge servers
+            edges_total_clients = [
+                len(i)
+                for i in np.array_split(
+                    np.arange(Config().clients.total_clients),
+                    Config().algorithm.total_silos,
+                )
+            ]
+            self.total_clients = edges_total_clients[edge_server_id - 1]
+
+            self.clients_per_round = [
+                len(i)
+                for i in np.array_split(
+                    np.arange(Config().clients.per_round),
+                    Config().algorithm.total_silos,
+                )
+            ][edge_server_id - 1]
+
+            starting_client_id = sum(edges_total_clients[: edge_server_id - 1])
+            self.clients_pool = list(
+                range(
+                    starting_client_id + 1, starting_client_id + 1 + self.total_clients
+                )
+            )
+
+            logging.info(
+                "[Edge server #%d (#%d)] Started training on %d clients with %d per round.",
+                Config().args.id,
+                os.getpid(),
+                self.total_clients,
+                self.clients_per_round,
+            )
+
+            # The training time of a edge server in one global round
+            self.edge_training_time = 0
+
+            # The training time of a edge server with its clients in one global round
+            self.edge_comm_time = 0
+
+        # Compute the number of clients for the central server
+        if Config().is_central_server():
+            self.clients_per_round = Config().algorithm.total_silos
+            self.total_clients = self.clients_per_round
+
+            logging.info(
+                "The central server starts training with %s edge servers.",
+                self.total_clients,
+            )
+
+    def configure(self) -> None:
+        """
+        Booting the federated learning server by setting up the data, model, and
+        creating the clients.
+        """
+        super().configure()
+
+        if Config().is_edge_server():
+            logging.info(
+                "Configuring edge server #%d as a %s server.",
+                Config().args.id,
+                Config().algorithm.type,
+            )
+            logging.info(
+                "[Edge server #%d (#%d)] Training with %s local aggregation rounds.",
+                Config().args.id,
+                os.getpid(),
+                Config().algorithm.local_rounds,
+            )
+
+            self.init_trainer()
+            self.trainer.set_client_id(Config().args.id)
+
+            # Prepares this server for processors that processes outbound and inbound
+            # data payloads
+            self.outbound_processor, self.inbound_processor = processor_registry.get(
+                "Server", server_id=os.getpid(), trainer=self.trainer
+            )
+
+            if (
+                hasattr(Config().server, "edge_do_test")
+                and Config().server.edge_do_test
+            ):
+                self.datasource = datasources_registry.get(client_id=0)
+                self.testset = self.datasource.get_test_set()
+
+                if hasattr(Config().data, "testset_sampler"):
+                    # Set the sampler for test set
+                    self.testset_sampler = samplers_registry.get(
+                        self.datasource, Config().args.id, testing=True
+                    )
+                else:
+                    if hasattr(Config().data, "testset_size"):
+                        self.testset_sampler = all_inclusive.Sampler(
+                            self.datasource, testing=True
+                        )
+
+    async def _select_clients(self, for_next_batch=False):
+        if Config().is_edge_server() and not for_next_batch:
+            if self.current_round == 0:
+                # Wait until this edge server is selected by the central server
+                # to avoid the edge server selects clients and clients begin training
+                # before the edge server is selected
+                await self.new_global_round_begins.wait()
+                self.new_global_round_begins.clear()
+
+        await super()._select_clients(for_next_batch=for_next_batch)
+
+    def customize_server_response(self, server_response: dict, client_id) -> dict:
+        """Wrap up generating the server response with any additional information."""
+        if Config().is_central_server():
+            server_response["current_global_round"] = self.current_round
+        return server_response
+
+    async def _process_reports(self):
+        """Process the client reports by aggregating their weights."""
+        # To pass the client_id == 0 assertion during aggregation
+        self.trainer.set_client_id(0)
+
+        weights_received = [update.payload for update in self.updates]
+
+        weights_received = self.weights_received(weights_received)
+        self.callback_handler.call_event("on_weights_received", self, weights_received)
+
+        # Extract the current model weights as the baseline
+        baseline_weights = self.algorithm.extract_weights()
+
+        if hasattr(self, "aggregate_weights"):
+            # Runs a server aggregation algorithm using weights rather than deltas
+            logging.info(
+                "[Server #%d] Aggregating model weights directly rather than weight deltas.",
+                os.getpid(),
+            )
+            updated_weights = await self.aggregate_weights(
+                self.updates, baseline_weights, weights_received
+            )
+
+            # Loads the new model weights
+            self.algorithm.load_weights(updated_weights)
+        else:
+            # Computes the weight deltas by comparing the weights received with
+            # the current global model weights
+            deltas_received = self.algorithm.compute_weight_deltas(
+                baseline_weights, weights_received
+            )
+            # Runs a framework-agnostic server aggregation algorithm, such as
+            # the federated averaging algorithm
+            logging.info("[Server #%d] Aggregating model weight deltas.", os.getpid())
+            deltas = await self.aggregate_deltas(self.updates, deltas_received)
+            # Updates the existing model weights from the provided deltas
+            updated_weights = self.algorithm.update_weights(deltas)
+            # Loads the new model weights
+            self.algorithm.load_weights(updated_weights)
+
+        # The model weights have already been aggregated, now calls the
+        # corresponding hook and callback
+        self.weights_aggregated(self.updates)
+        self.callback_handler.call_event("on_weights_aggregated", self, self.updates)
+
+        if Config().is_edge_server():
+            self.trainer.set_client_id(Config().args.id)
+
+        # Testing the model accuracy
+        if (Config().is_edge_server() and Config().clients.do_test) or (
+            Config().is_central_server()
+            and hasattr(Config().server, "edge_do_test")
+            and Config().server.edge_do_test
+        ):
+            # Compute the average accuracy from client reports
+            (
+                self.average_accuracy,
+                self.std_accuracy,
+            ) = self.get_accuracy_mean_std(self.updates)
+            logging.info(
+                "[%s] Average client accuracy: %.2f%%.",
+                self,
+                100 * self.average_accuracy,
+            )
+        elif Config().is_central_server() and Config().clients.do_test:
+            # Compute the average accuracy from client reports
+            total_samples = sum(update.report.num_samples for update in self.updates)
+            self.average_accuracy = (
+                sum(
+                    update.report.average_accuracy * update.report.num_samples
+                    for update in self.updates
+                )
+                / total_samples
+            )
+
+            logging.info(
+                "[%s] Average client accuracy: %.2f%%.",
+                self,
+                100 * self.average_accuracy,
+            )
+
+        if (
+            Config().is_central_server()
+            and hasattr(Config().server, "do_test")
+            and Config().server.do_test
+        ):
+            # Testing the updated model directly at the server
+            logging.info("[%s] Started model testing.", self)
+            self.accuracy = self.trainer.test(self.testset, self.testset_sampler)
+
+            if hasattr(Config().trainer, "target_perplexity"):
+                logging.info(
+                    fonts.colourize(
+                        f"[{self}] Global model perplexity: {self.accuracy:.2f}\n"
+                    )
+                )
+            else:
+                logging.info(
+                    fonts.colourize(
+                        f"[{self}] Global model accuracy: {100 * self.accuracy:.2f}%\n"
+                    )
+                )
+        elif (
+            Config().is_edge_server()
+            and hasattr(Config().server, "edge_do_test")
+            and Config().server.edge_do_test
+        ):
+            # Test the aggregated model directly at the edge server
+            logging.info("[%s] Started model testing.", self)
+            self.accuracy = self.trainer.test(self.testset, self.testset_sampler)
+
+            if hasattr(Config().trainer, "target_perplexity"):
+                logging.info(
+                    fonts.colourize(
+                        f"[{self}] Global model perplexity: {self.accuracy:.2f}\n"
+                    )
+                )
+            else:
+                logging.info(
+                    fonts.colourize(
+                        f"[{self}] Global model accuracy: {100 * self.accuracy:.2f}%\n"
+                    )
+                )
+        else:
+            self.accuracy = self.average_accuracy
+            self.accuracy_std = self.std_accuracy
+
+        self.clients_processed()
+        self.callback_handler.call_event("on_clients_processed", self)
+
+    def clients_processed(self):
+        """Additional work to be performed after client reports have been processed."""
+        # Record results into a .csv file
+        if Config().is_central_server():
+            super().clients_processed()
+
+        if Config().is_edge_server():
+            logged_items = self.get_logged_items()
+            self.edge_training_time += logged_items["round_time"]
+            self.edge_comm_time += logged_items["comm_time"]
+
+            # When a certain number of aggregations are completed, an edge client
+            # needs to be signaled to send a report to the central server
+            if self.current_round == Config().algorithm.local_rounds:
+                logging.info(
+                    "[Server #%d] Completed %s rounds of local aggregation.",
+                    os.getpid(),
+                    Config().algorithm.local_rounds,
+                )
+                self.model_aggregated.set()
+
+                self.current_round = 0
+                self.current_global_round += 1
+
+    def get_logged_items(self) -> dict:
+        """Get items to be logged by the LogProgressCallback class in a .csv file."""
+        logged_items = super().get_logged_items()
+
+        logged_items["global_round"] = self.current_global_round
+        logged_items["average_accuracy"] = self.average_accuracy
+        logged_items["edge_agg_num"] = Config().algorithm.local_rounds
+        logged_items["local_epoch_num"] = Config().trainer.epochs
+
+        if Config().is_central_server():
+            logged_items["comm_time"] = max(
+                update.report.comm_time + update.report.edge_server_comm_time
+                for update in self.updates
+            )
+
+        return logged_items
+
+    async def wrap_up(self) -> None:
+        """Wrapping up when each round of training is done."""
+        if Config().is_central_server():
+            await super().wrap_up()

plato/servers/fedavg_gan.py

@@ -0,0 +1,74 @@
+"""
+A federated learning server using federated averaging to train GAN models.
+"""
+
+import asyncio
+
+from plato.servers import fedavg
+from plato.config import Config
+
+
+class Server(fedavg.Server):
+    """Federated learning server using federated averaging to train GAN models."""
+
+    async def aggregate_deltas(self, updates, deltas_received):
+        """Aggregate weight updates from the clients using federated averaging."""
+        # Total sample is the same for both Generator and Discriminator
+        self.total_samples = sum(update.report.num_samples for update in updates)
+
+        # Perform weighted averaging for both Generator and Discriminator
+        gen_avg_update = {
+            name: self.trainer.zeros(weights.shape)
+            for name, weights in deltas_received[0][0].items()
+        }
+        disc_avg_update = {
+            name: self.trainer.zeros(weights.shape)
+            for name, weights in deltas_received[0][1].items()
+        }
+
+        for i, update in enumerate(deltas_received):
+            num_samples = updates[i].report.num_samples
+
+            update_from_gen, update_from_disc = update
+
+            for name, delta in update_from_gen.items():
+                gen_avg_update[name] += delta * (num_samples / self.total_samples)
+
+            for name, delta in update_from_disc.items():
+                disc_avg_update[name] += delta * (num_samples / self.total_samples)
+
+            # Yield to other tasks in the server
+            await asyncio.sleep(0)
+
+        return gen_avg_update, disc_avg_update
+
+    def customize_server_payload(self, payload):
+        """
+        Customize the server payload before sending to the client.
+
+        At the end of each round, the server can choose to only send the global Generator
+        or Discriminator (or both or neither) model to the clients next round.
+
+        Reference this paper for more detail:
+        https://deepai.org/publication/federated-generative-adversarial-learning
+
+        By default, both model will be sent to the clients.
+        """
+        if hasattr(Config().server, "network_to_sync"):
+            network = Config().server.network_to_sync.lower()
+        else:
+            network = "both"
+
+        weights_gen, weights_disc = payload
+        if network == "none":
+            server_payload = None, None
+        elif network == "generator":
+            server_payload = weights_gen, None
+        elif network == "discriminator":
+            server_payload = None, weights_disc
+        elif network == "both":
+            server_payload = payload
+        else:
+            raise ValueError(f"Unknown value to attribute network_to_sync: {network}")
+
+        return server_payload