plato-learn 1.1__py3-none-any.whl

plato/clients/edge.py

@@ -0,0 +1,103 @@
+"""
+A federated learning client at the edge server in a cross-silo training workload.
+"""
+
+import time
+from types import SimpleNamespace
+
+from plato.clients import simple
+from plato.config import Config
+from plato.processors import registry as processor_registry
+
+
+class Client(simple.Client):
+    """A federated learning client at the edge server in a cross-silo training workload."""
+
+    def __init__(
+        self,
+        server,
+        model=None,
+        datasource=None,
+        algorithm=None,
+        trainer=None,
+        callbacks=None,
+    ):
+        super().__init__(
+            model=model,
+            datasource=datasource,
+            algorithm=algorithm,
+            trainer=trainer,
+            callbacks=callbacks,
+        )
+        self.server = server
+
+    def configure(self) -> None:
+        """Prepare this edge client for training."""
+        super().configure()
+
+        # Pass inbound and outbound data payloads through processors for
+        # additional data processing
+        self.outbound_processor, self.inbound_processor = processor_registry.get(
+            "Client", client_id=self.client_id, trainer=self.server.trainer
+        )
+
+    def load_data(self) -> None:
+        """The edge client does not need to train models using local data."""
+
+    def _load_payload(self, server_payload) -> None:
+        """The edge client loads the model from the central server."""
+        self.server.algorithm.load_weights(server_payload)
+
+    def process_server_response(self, server_response):
+        """Additional client-specific processing on the server response."""
+        if "current_global_round" in server_response:
+            self.server.current_global_round = server_response["current_global_round"]
+
+    async def _train(self):
+        """The aggregation workload on an edge client."""
+        training_start_time = time.perf_counter()
+        # Signal edge server to select clients to start a new round of local aggregation
+        self.server.new_global_round_begins.set()
+
+        # Wait for the edge server to finish model aggregation
+        await self.server.model_aggregated.wait()
+        self.server.model_aggregated.clear()
+
+        # Extract model weights and biases
+        weights = self.server.algorithm.extract_weights()
+
+        average_accuracy = self.server.average_accuracy
+        accuracy = self.server.accuracy
+
+        if (
+            hasattr(Config().clients, "sleep_simulation")
+            and Config().clients.sleep_simulation
+        ):
+            training_time = self.server.edge_training_time
+            self.server.edge_training_time = 0
+        else:
+            training_time = time.perf_counter() - training_start_time
+
+        comm_time = time.time()
+
+        edge_server_comm_time = self.server.edge_comm_time
+        self.server.edge_comm_time = 0
+
+        # Generate a report for the central server
+        report = SimpleNamespace(
+            client_id=self.client_id,
+            num_samples=self.server.total_samples,
+            accuracy=accuracy,
+            training_time=training_time,
+            comm_time=comm_time,
+            update_response=False,
+            average_accuracy=average_accuracy,
+            edge_server_comm_overhead=self.server.comm_overhead,
+            edge_server_comm_time=edge_server_comm_time,
+        )
+
+        self._report = self.customize_report(report)
+
+        self.server.comm_overhead = 0
+
+        return self._report, weights

plato/clients/fedavg_personalized.py

@@ -0,0 +1,40 @@
+"""
+A personalized federated learning client that saves its local layers before
+sending the shared global model to the server after local training.
+"""
+
+from collections import OrderedDict
+
+from plato.clients import simple
+from plato.config import Config
+
+
+class Client(simple.Client):
+    """
+    A personalized federated learning client that saves its local layers before sending the
+    shared global model to the server after local training.
+    """
+
+    def outbound_ready(self, report, outbound_processor):
+        super().outbound_ready(report, outbound_processor)
+        weights = self.algorithm.extract_weights()
+
+        # Save local layers before giving them to the outbound processor
+        if hasattr(Config().algorithm, "local_layer_names"):
+            # Extract weights of desired local layers
+            local_layers = OrderedDict(
+                [
+                    (name, param)
+                    for name, param in weights.items()
+                    if any(
+                        param_name in name.strip().split(".")
+                        for param_name in Config().algorithm.local_layer_names
+                    )
+                ]
+            )
+
+            model_path = Config().params["model_path"]
+            model_name = Config().trainer.model_name
+            filename = f"{model_path}/{model_name}_{self.client_id}_local_layers.pth"
+
+            self.algorithm.save_local_layers(local_layers, filename)

plato/clients/mistnet.py

@@ -0,0 +1,49 @@
+"""
+A federated learning client for MistNet.
+
+Reference:
+
+P. Wang, et al. "MistNet: Towards Private Neural Network Training with Local
+Differential Privacy," found in docs/papers.
+"""
+
+import logging
+import time
+from types import SimpleNamespace
+
+from plato.config import Config
+from plato.clients import simple
+
+
+class Client(simple.Client):
+    """A federated learning client for MistNet."""
+
+    async def _train(self):
+        """A MistNet client only uses the first several layers in a forward pass."""
+        logging.info("Training on MistNet client #%d", self.client_id)
+
+        # Since training is performed on the server, the client should not be doing
+        # its own testing for the model accuracy
+        assert not Config().clients.do_test
+
+        tic = time.perf_counter()
+
+        # Perform a forward pass till the cut layer in the model
+        features = self.algorithm.extract_features(self.trainset, self.sampler)
+
+        training_time = time.perf_counter() - tic
+
+        # Generate a report for the server, performing model testing if applicable
+        comm_time = time.time()
+        return (
+            SimpleNamespace(
+                client_id=self.client_id,
+                num_samples=self.sampler.num_samples(),
+                accuracy=0,
+                training_time=training_time,
+                comm_time=comm_time,
+                update_response=False,
+                payload_length=len(features),
+            ),
+            features,
+        )

plato/clients/registry.py

@@ -0,0 +1,43 @@
+"""
+The registry that contains all available federated learning clients.
+
+Having a registry of all available classes is convenient for retrieving an instance based
+on a configuration at run-time.
+"""
+
+import logging
+
+from plato.config import Config
+from plato.clients import (
+    self_supervised_learning,
+    simple,
+    mistnet,
+    fedavg_personalized,
+    split_learning,
+)
+
+registered_clients = {
+    "simple": simple.Client,
+    "mistnet": mistnet.Client,
+    "fedavg_personalized": fedavg_personalized.Client,
+    "self_supervised_learning": self_supervised_learning.Client,
+    "split_learning": split_learning.Client,
+}
+
+
+def get(model=None, datasource=None, algorithm=None, trainer=None):
+    """Get an instance of the server."""
+    if hasattr(Config().clients, "type"):
+        client_type = Config().clients.type
+    else:
+        client_type = Config().algorithm.type
+
+    if client_type in registered_clients:
+        logging.info("Client: %s", client_type)
+        registered_client = registered_clients[client_type](
+            model=model, datasource=datasource, algorithm=algorithm, trainer=trainer
+        )
+    else:
+        raise ValueError(f"No such client: {client_type}")
+
+    return registered_client

plato/clients/self_supervised_learning.py

@@ -0,0 +1,51 @@
+"""
+A self-supervised learning (SSL) client prepares a personalized datasource for
+the personalization process, which will be performed after finishing the FL
+training process with SSL.
+
+Specifically, the conventional FL training process with SSL will train the model
+with the datasource and objective function of SSL. Yet, the datasource used in
+personalization should be one of supervised learning. Therefore, a client needs
+to prepare the personalized datasource.
+"""
+
+from plato.datasources import registry as datasources_registry
+from plato.clients import simple
+
+
+class Client(simple.Client):
+    """An SSL client to prepare the datasource for personalization."""
+
+    def __init__(
+        self,
+        model=None,
+        datasource=None,
+        algorithm=None,
+        trainer=None,
+        callbacks=None,
+        trainer_callbacks=None,
+    ):
+        super().__init__(
+            model=model,
+            datasource=datasource,
+            algorithm=algorithm,
+            trainer=trainer,
+            callbacks=callbacks,
+            trainer_callbacks=trainer_callbacks,
+        )
+        # The datasource used in personalization
+        self.personalized_datasource = None
+
+    def configure(self) -> None:
+        """Prepare this client for training."""
+        super().configure()
+
+        # Get the personalized datasource
+        if self.personalized_datasource is None:
+            personalized_datasource = datasources_registry.get()
+
+        # Set the train and the test set for the trainer
+        self.trainer.set_personalized_datasets(
+            personalized_datasource.get_train_set(),
+            personalized_datasource.get_test_set(),
+        )

plato/clients/simple.py

@@ -0,0 +1,218 @@
+"""
+A basic federated learning client who sends weight updates to the server.
+"""
+
+import logging
+import time
+from types import SimpleNamespace
+
+from plato.algorithms import registry as algorithms_registry
+from plato.clients import base
+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.trainers import registry as trainers_registry
+from plato.utils import fonts
+
+
+class Client(base.Client):
+    """A basic federated learning client who sends simple weight updates."""
+
+    def __init__(
+        self,
+        model=None,
+        datasource=None,
+        algorithm=None,
+        trainer=None,
+        callbacks=None,
+        trainer_callbacks=None,
+    ):
+        super().__init__(callbacks=callbacks)
+        # Save the callbacks that will be passed to trainer later
+        self.trainer_callbacks = trainer_callbacks
+
+        self.custom_model = model
+        self.model = None
+
+        self.custom_datasource = datasource
+        self.datasource = None
+
+        self.custom_algorithm = algorithm
+        self.algorithm = None
+
+        self.custom_trainer = trainer
+        self.trainer = None
+
+        self.trainset = None  # Training dataset
+        self.testset = None  # Testing dataset
+        self.sampler = None
+        self.testset_sampler = None  # Sampler for the test set
+
+        self._report = None
+
+    def configure(self) -> None:
+        """Prepares this client for training."""
+        super().configure()
+
+        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, callbacks=self.trainer_callbacks
+            )
+        elif self.trainer is None and self.custom_trainer is not None:
+            self.trainer = self.custom_trainer(
+                model=self.model, callbacks=self.trainer_callbacks
+            )
+
+        self.trainer.set_client_id(self.client_id)
+
+        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)
+
+        self.algorithm.set_client_id(self.client_id)
+
+        # Pass inbound and outbound data payloads through processors for
+        # additional data processing
+        self.outbound_processor, self.inbound_processor = processor_registry.get(
+            "Client", client_id=self.client_id, trainer=self.trainer
+        )
+
+        # Setting up the data sampler
+        if self.datasource:
+            self.sampler = samplers_registry.get(self.datasource, self.client_id)
+
+            if (
+                hasattr(Config().clients, "do_test")
+                and Config().clients.do_test
+                and hasattr(Config().data, "testset_sampler")
+            ):
+                # Set the sampler for test set
+                self.testset_sampler = samplers_registry.get(
+                    self.datasource, self.client_id, testing=True
+                )
+
+    def _load_data(self) -> None:
+        """Generates data and loads them onto this client."""
+        # The only case where Config().data.reload_data is set to true is
+        # when clients with different client IDs need to load from different datasets,
+        # such as in the pre-partitioned Federated EMNIST dataset. We do not support
+        # reloading data from a custom datasource at this time.
+        if (
+            self.datasource is None
+            or hasattr(Config().data, "reload_data")
+            and Config().data.reload_data
+        ):
+            logging.info("[%s] Loading its data source...", self)
+
+            if self.custom_datasource is None:
+                self.datasource = datasources_registry.get(client_id=self.client_id)
+            elif self.custom_datasource is not None:
+                self.datasource = self.custom_datasource()
+
+            logging.info(
+                "[%s] Dataset size: %s", self, self.datasource.num_train_examples()
+            )
+
+    def _allocate_data(self) -> None:
+        """Allocate training or testing dataset of this client."""
+        # PyTorch uses samplers when loading data with a data loader
+        self.trainset = self.datasource.get_train_set()
+
+        if hasattr(Config().clients, "do_test") and Config().clients.do_test:
+            # Set the testset if local testing is needed
+            self.testset = self.datasource.get_test_set()
+
+    def _load_payload(self, server_payload) -> None:
+        """Loads the server model onto this client."""
+        self.algorithm.load_weights(server_payload)
+
+    async def _train(self):
+        """The machine learning training workload on a client."""
+        logging.info(
+            fonts.colourize(
+                f"[{self}] Started training in communication round #{self.current_round}."
+            )
+        )
+
+        # Perform model training
+        try:
+            if hasattr(self.trainer, "current_round"):
+                self.trainer.current_round = self.current_round
+            training_time = self.trainer.train(self.trainset, self.sampler)
+
+        except ValueError as exc:
+            logging.info(
+                fonts.colourize(f"[{self}] Error occurred during training: {exc}")
+            )
+            await self.sio.disconnect()
+
+        # Extract model weights and biases
+        weights = self.algorithm.extract_weights()
+
+        # Generate a report for the server, performing model testing if applicable
+        if (hasattr(Config().clients, "do_test") and Config().clients.do_test) and (
+            not hasattr(Config().clients, "test_interval")
+            or self.current_round % Config().clients.test_interval == 0
+        ):
+            accuracy = self.trainer.test(self.testset, self.testset_sampler)
+
+            if accuracy == -1:
+                # The testing process failed, disconnect from the server
+                logging.info(
+                    fonts.colourize(
+                        f"[{self}] Accuracy is -1 when testing. Disconnecting from the server."
+                    )
+                )
+                await self.sio.disconnect()
+
+            if hasattr(Config().trainer, "target_perplexity"):
+                logging.info("[%s] Test perplexity: %.2f", self, accuracy)
+            else:
+                logging.info("[%s] Test accuracy: %.2f%%", self, 100 * accuracy)
+        else:
+            accuracy = 0
+
+        comm_time = time.time()
+
+        if (
+            hasattr(Config().clients, "sleep_simulation")
+            and Config().clients.sleep_simulation
+        ):
+            sleep_seconds = Config().client_sleep_times[self.client_id - 1]
+            avg_training_time = Config().clients.avg_training_time
+
+            training_time = (
+                avg_training_time + sleep_seconds
+            ) * Config().trainer.epochs
+
+        report = SimpleNamespace(
+            client_id=self.client_id,
+            num_samples=self.sampler.num_samples(),
+            accuracy=accuracy,
+            training_time=training_time,
+            comm_time=comm_time,
+            update_response=False,
+        )
+
+        self._report = self.customize_report(report)
+
+        return self._report, weights
+
+    async def _obtain_model_update(self, client_id, requested_time):
+        """Retrieves a model update corresponding to a particular wall clock time."""
+        model = self.trainer.obtain_model_update(client_id, requested_time)
+        weights = self.algorithm.extract_weights(model)
+        self._report.comm_time = time.time()
+        self._report.client_id = client_id
+        self._report.update_response = True
+
+        return self._report, weights
+
+    def customize_report(self, report: SimpleNamespace) -> SimpleNamespace:
+        """Customizes the report with any additional information."""
+        return report

plato/clients/split_learning.py

@@ -0,0 +1,150 @@
+"""
+A federated learning client using split learning.
+
+Reference:
+
+Vepakomma, et al., "Split Learning for Health: Distributed Deep Learning without Sharing
+Raw Patient Data," in Proc. AI for Social Good Workshop, affiliated with ICLR 2018.
+
+https://arxiv.org/pdf/1812.00564.pdf
+
+Chopra, Ayush, et al. "AdaSplit: Adaptive Trade-offs for Resource-constrained Distributed
+Deep Learning." arXiv preprint arXiv:2112.01637 (2021).
+
+https://arxiv.org/pdf/2112.01637.pdf
+"""
+
+import logging
+import time
+from types import SimpleNamespace
+
+from plato.clients import simple
+from plato.config import Config
+from plato.utils import fonts
+
+
+class Client(simple.Client):
+    """A split learning client."""
+
+    # pylint:disable=too-many-arguments
+    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,
+        )
+        assert not Config().clients.do_test
+
+        self.model_received = False
+        self.gradient_received = False
+        self.contexts = {}
+        self.original_weights = None
+
+        # Iteration control
+        self.iterations = Config().clients.iteration
+        self.iter_left = Config().clients.iteration
+
+        # Sampler cannot be reconfigured otherwise same training samples
+        # will be selected every round
+        self.static_sampler = None
+
+    async def inbound_processed(self, processed_inbound_payload):
+        """Extract features or complete the training using split learning."""
+        server_payload, info = processed_inbound_payload
+
+        # Preparing the client response
+        report, payload = None, None
+
+        if info == "prompt":
+            # Server prompts a new client to conduct split learning
+            self._load_context(self.client_id)
+            report, payload = self._extract_features()
+        elif info == "gradients":
+            # server sends the gradients of the features, i.e., complete training
+            logging.warning("[%s] Gradients received, complete training.", self)
+            training_time, weights = self._complete_training(server_payload)
+            self.iter_left -= 1
+
+            if self.iter_left == 0:
+                logging.warning(
+                    "[%s] Finished training, sending weights to the server.", self
+                )
+                # Send weights to server for evaluation
+                report = SimpleNamespace(
+                    client_id=self.client_id,
+                    num_samples=self.sampler.num_samples(),
+                    accuracy=0,
+                    training_time=training_time,
+                    comm_time=time.time(),
+                    update_response=False,
+                    type="weights",
+                )
+                payload = weights
+                self.iter_left = self.iterations
+            else:
+                # Continue feature extraction
+                report, payload = self._extract_features()
+                report.training_time += training_time
+
+            # Save the state of current client
+            self._save_context(self.client_id)
+        return report, payload
+
+    def _save_context(self, client_id):
+        """Saving the extracted weights and the data sampler for a given client."""
+        # Sampler needs to be saved otherwise same data samples will be selected every round
+        self.contexts[client_id] = (
+            self.algorithm.extract_weights(),
+            self.static_sampler,
+        )
+
+    def _load_context(self, client_id):
+        """Load client's model weights and the sampler from last selected round."""
+        if not client_id in self.contexts:
+            if self.original_weights is None:
+                self.original_weights = self.algorithm.extract_weights()
+            self.algorithm.load_weights(self.original_weights)
+            self.static_sampler = self.sampler.get()
+        else:
+            weights, sampler = self.contexts.pop(client_id)
+            self.algorithm.load_weights(weights)
+            self.static_sampler = sampler
+
+    def _extract_features(self):
+        """Extract the feature till the cut layer."""
+        round_number = self.iterations - self.iter_left + 1
+        logging.warning(
+            fonts.colourize(
+                f"[{self}] Started split learning in round #{round_number}/{self.iterations}"
+                + f" (Global round {self.current_round})."
+            )
+        )
+
+        features, training_time = self.algorithm.extract_features(
+            self.trainset, self.static_sampler
+        )
+        report = SimpleNamespace(
+            client_id=self.client_id,
+            num_samples=self.sampler.num_samples(),
+            accuracy=0,
+            training_time=training_time,
+            comm_time=time.time(),
+            update_response=False,
+            type="features",
+        )
+        return report, features
+
+    def _complete_training(self, payload):
+        """Complete the training based on the gradients from server."""
+        training_time = self.algorithm.complete_train(payload)
+        weights = self.algorithm.extract_weights()
+        return training_time, weights