plato-learn 1.1__py3-none-any.whl

plato/servers/fedavg_he.py

@@ -0,0 +1,106 @@
+"""
+A federated learning server using federated averaging to aggregate updates after homomorphic encryption.
+"""
+
+from functools import reduce
+from plato.servers import fedavg
+from plato.utils import homo_enc
+
+
+class Server(fedavg.Server):
+    """
+    Federated learning server using federated averaging to aggregate updates after homomorphic
+    encryption.
+    """
+
+    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.context = homo_enc.get_ckks_context()
+        self.encrypted_model = None
+        self.weight_shapes = {}
+        self.para_nums = {}
+
+    def configure(self) -> None:
+        """Configure the model information like weight shapes and parameter numbers."""
+        super().configure()
+
+        extract_model = self.trainer.model.cpu().state_dict()
+
+        for key in extract_model.keys():
+            self.weight_shapes[key] = extract_model[key].size()
+            self.para_nums[key] = extract_model[key].numel()
+
+        self.encrypted_model = homo_enc.encrypt_weights(
+            extract_model, True, self.context, []
+        )
+
+    def customize_server_payload(self, payload):
+        """Server can only send the encrypted aggreagtion result to clients."""
+        return self.encrypted_model
+
+    # pylint: disable=unused-argument
+    async def aggregate_weights(self, updates, baseline_weights, weights_received):
+        """Aggregate the model updates and decrypt the result for evaluation purpose."""
+        self.encrypted_model = self._fedavg_hybrid(updates)
+
+        # Decrypt model weights for test accuracy
+        decrypted_weights = homo_enc.decrypt_weights(
+            self.encrypted_model, self.weight_shapes, self.para_nums
+        )
+        # Serialize the encrypted weights after decryption
+        self.encrypted_model["encrypted_weights"] = self.encrypted_model[
+            "encrypted_weights"
+        ].serialize()
+
+        return decrypted_weights
+
+    def _fedavg_hybrid(self, updates):
+        """Aggregate the model updates in the hybrid form of encrypted and unencrypted weights."""
+        weights_received = [
+            homo_enc.deserialize_weights(update.payload, self.context)
+            for update in updates
+        ]
+        unencrypted_weights = [
+            homo_enc.extract_encrypted_model(x)[0] for x in weights_received
+        ]
+        encrypted_weights = [
+            homo_enc.extract_encrypted_model(x)[1] for x in weights_received
+        ]
+        # Assert the encrypted weights from all clients are aligned
+        indices = [homo_enc.extract_encrypted_model(x)[2] for x in weights_received]
+        for i in range(1, len(indices)):
+            assert indices[i] == indices[0]
+        encrypt_indices = indices[0]
+
+        # Extract the total number of samples
+        self.total_samples = sum(update.report.num_samples for update in updates)
+
+        # Perform weighted averaging on unencrypted weights
+        unencrypted_avg_update = self.trainer.zeros(unencrypted_weights[0].size)
+        encrypted_avg_update = self.trainer.zeros(encrypted_weights[0].size())
+
+        for i, (unenc_w, enc_w) in enumerate(
+            zip(unencrypted_weights, encrypted_weights)
+        ):
+            report = updates[i].report
+            num_samples = report.num_samples
+
+            unencrypted_avg_update += unenc_w * (num_samples / self.total_samples)
+            encrypted_avg_update += enc_w * (num_samples / self.total_samples)
+
+        if len(encrypt_indices) == 0:
+            # No weights are encrypted, set to None
+            encrypted_avg_update = None
+
+        return homo_enc.wrap_encrypted_model(
+            unencrypted_avg_update, encrypted_avg_update, encrypt_indices
+        )

plato/servers/fedavg_personalized.py

@@ -0,0 +1,57 @@
+"""
+A personalized federated learning server that starts from a number of regular
+rounds of federated learning. In these regular rounds, only a subset of the
+total clients can be selected to perform the local update (the ratio of which is
+a configuration setting). After all regular rounds are completed, it starts a
+final round of personalization, where a selected subset of clients perform local
+training using their local dataset.
+"""
+
+from plato.servers import fedavg
+from plato.config import Config
+
+
+class Server(fedavg.Server):
+    """
+    A personalzed FL server that controls how many clients will participate in
+    the training process, and that adds a final personalization round with all
+    clients sampled.
+    """
+
+    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,
+        )
+        # Personalization starts after the final regular round of training
+        self.personalization_started = False
+
+    def choose_clients(self, clients_pool, clients_count):
+        """Choose a subset of the clients to participate in each round."""
+        if self.current_round > Config().trainer.rounds:
+            # In the final personalization round, choose from all clients
+            return super().choose_clients(clients_pool, clients_count)
+        else:
+            ratio = Config().algorithm.personalization.participating_client_ratio
+
+            return super().choose_clients(
+                clients_pool[: int(self.total_clients * ratio)],
+                clients_count,
+            )
+
+    async def wrap_up(self) -> None:
+        """Wraps up when each round of training is done."""
+        if self.personalization_started:
+            await super().wrap_up()
+        else:
+            # If the target number of training rounds has been reached, start
+            # the final round of training for personalization on the clients
+            self.save_to_checkpoint()
+
+            if self.current_round >= Config().trainer.rounds:
+                self.personalization_started = True

plato/servers/mistnet.py

@@ -0,0 +1,67 @@
+"""
+A federated learning server for MistNet.
+
+Reference:
+P. Wang, et al. "MistNet: Towards Private Neural Network Training with Local
+Differential Privacy," found in docs/papers.
+"""
+
+import logging
+import os
+
+from plato.config import Config
+from plato.datasources import feature
+from plato.samplers import all_inclusive
+from plato.servers import fedavg
+
+
+class Server(fedavg.Server):
+    """The MistNet server for federated learning."""
+
+    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,
+        )
+
+        # MistNet requires one round of client-server communication
+        assert Config().trainer.rounds == 1
+
+    def init_trainer(self) -> None:
+        """Setting up a pre-trained model to be loaded on the server."""
+        super().init_trainer()
+
+        model_path = Config().params["model_path"]
+        model_file_name = (
+            Config().trainer.pretrained_model
+            if hasattr(Config().trainer, "pretrained_model")
+            else f"{Config().trainer.model_name}.pth"
+        )
+        pretrained_model_path = f"{model_path}/{model_file_name}"
+
+        if os.path.exists(pretrained_model_path):
+            logging.info("[Server #%d] Loading a pre-trained model.", os.getpid())
+            self.trainer.load_model(filename=model_file_name)
+
+    async def _process_reports(self):
+        """Process the features extracted by the client and perform server-side training."""
+        features = [update.payload for update in self.updates]
+        feature_dataset = feature.DataSource(features)
+
+        # Training the model using all the features received from the client
+        sampler = all_inclusive.Sampler(feature_dataset)
+        self.algorithm.train(feature_dataset, sampler)
+
+        # Test the updated model
+        if not hasattr(Config().server, "do_test") or Config().server.do_test:
+            self.accuracy = self.trainer.test(self.testset)
+            logging.info(
+                "[%s] Global model accuracy: %.2f%%\n", self, 100 * self.accuracy
+            )
+
+        self.clients_processed()

plato/servers/registry.py

@@ -0,0 +1,52 @@
+"""
+The registry for servers that contains framework-agnostic implementations on a
+federated learning server.
+
+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.servers import (
+    fedavg,
+    fedavg_cs,
+    mistnet,
+    fedavg_gan,
+    fedavg_personalized,
+    split_learning,
+)
+
+if hasattr(Config().server, "type") and Config().server.type == "fedavg_he":
+    # FedAvg server with homomorphic encryption supports PyTorch only
+    from plato.servers import fedavg_he
+
+    registered_servers = {"fedavg_he": fedavg_he.Server}
+
+else:
+    registered_servers = {
+        "fedavg": fedavg.Server,
+        "fedavg_cross_silo": fedavg_cs.Server,
+        "mistnet": mistnet.Server,
+        "fedavg_gan": fedavg_gan.Server,
+        "fedavg_personalized": fedavg_personalized.Server,
+        "split_learning": split_learning.Server,
+    }
+
+
+def get(model=None, algorithm=None, trainer=None):
+    """Get an instance of the server."""
+    if hasattr(Config().server, "type"):
+        server_type = Config().server.type
+    else:
+        server_type = Config().algorithm.type
+
+    if server_type in registered_servers:
+        logging.info("Server: %s", server_type)
+        return registered_servers[server_type](
+            model=model, algorithm=algorithm, trainer=trainer
+        )
+    else:
+        raise ValueError(f"No such server: {server_type}")

plato/servers/split_learning.py

@@ -0,0 +1,109 @@
+"""
+A federated learning server 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
+
+from plato.config import Config
+from plato.datasources import feature
+from plato.samplers import all_inclusive
+from plato.servers import fedavg
+from plato.utils import fonts
+from plato.datasources import registry as datasources_registry
+
+
+# pylint:disable=too-many-instance-attributes
+class Server(fedavg.Server):
+    """The split learning server."""
+
+    # pylint:disable=too-many-arguments
+    def __init__(
+        self, model=None, datasource=None, algorithm=None, trainer=None, callbacks=None
+    ):
+        super().__init__(model, datasource, algorithm, trainer, callbacks)
+        # Split learning clients interact with server sequentially
+        assert Config().clients.per_round == 1
+        self.phase = "prompt"
+        self.clients_list = []
+        self.client_last = None
+        self.next_client = True
+        self.test_accuracy = 0.0
+
+        # Manually set up the testset since do_test is turned off in config
+        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()
+        self.testset_sampler = all_inclusive.Sampler(self.datasource, testing=True)
+
+    def choose_clients(self, clients_pool, clients_count):
+        """Shuffle the clients and sequentially select them when the previous one is done."""
+        if len(self.clients_list) == 0 and self.next_client:
+            # Shuffle the client list
+            self.clients_list = super().choose_clients(clients_pool, len(clients_pool))
+            logging.warning("Client order: %s", str(self.clients_list))
+
+        if self.next_client:
+            # Sequentially select clients
+            self.client_last = [self.clients_list.pop(0)]
+            self.next_client = False
+        return self.client_last
+
+    def customize_server_payload(self, payload):
+        """Wrap up generating the server payload with any additional information."""
+        if self.phase == "prompt":
+            # Split learning server doesn't send weights to client
+            return (None, "prompt")
+        return (self.trainer.get_gradients(), "gradients")
+
+    # pylint: disable=unused-argument
+    async def aggregate_weights(self, updates, baseline_weights, weights_received):
+        """Aggregate weight updates from the clients or train the model."""
+        update = updates[0]
+        report = update.report
+        if report.type == "features":
+            logging.warning("[%s] Features received, compute gradients.", self)
+            feature_dataset = feature.DataSource([update.payload])
+
+            # Training the model using all the features received from the client
+            sampler = all_inclusive.Sampler(feature_dataset)
+            self.algorithm.train(feature_dataset, sampler)
+
+            self.phase = "gradient"
+        elif report.type == "weights":
+            logging.warning("[%s] Weights received, start testing accuracy.", self)
+            weights = update.payload
+
+            # The weights after cut layer are not trained by clients
+            self.algorithm.update_weights_before_cut(weights)
+
+            self.test_accuracy = self.trainer.test(self.testset, self.testset_sampler)
+
+            logging.warning(
+                fonts.colourize(
+                    f"[{self}] Global model accuracy: {100 * self.test_accuracy:.2f}%\n"
+                )
+            )
+            self.phase = "prompt"
+            # Change client in next round
+            self.next_client = True
+
+        updated_weights = self.algorithm.extract_weights()
+        return updated_weights
+
+    def clients_processed(self):
+        # Replace the default accuracy by manually tested accuracy
+        self.accuracy = self.test_accuracy

plato/trainers/base.py

@@ -0,0 +1,99 @@
+"""
+Base class for trainers.
+"""
+
+from abc import ABC, abstractmethod
+import os
+
+from plato.config import Config
+
+
+class Trainer(ABC):
+    """Base class for all the trainers."""
+
+    def __init__(self):
+        self.device = Config().device()
+        self.client_id = 0
+
+    def set_client_id(self, client_id):
+        """Setting the client ID."""
+        self.client_id = client_id
+
+    @abstractmethod
+    def save_model(self, filename=None, location=None):
+        """Saving the model to a file."""
+        raise TypeError("save_model() not implemented.")
+
+    @abstractmethod
+    def load_model(self, filename=None, location=None):
+        """Loading pre-trained model weights from a file."""
+        raise TypeError("load_model() not implemented.")
+
+    @staticmethod
+    def save_accuracy(accuracy, filename=None):
+        """Saving the test accuracy to a file."""
+        model_path = Config().params["model_path"]
+        model_name = Config().trainer.model_name
+
+        if not os.path.exists(model_path):
+            os.makedirs(model_path)
+
+        if filename is not None:
+            accuracy_path = f"{model_path}/{filename}"
+        else:
+            accuracy_path = f"{model_path}/{model_name}.acc"
+
+        with open(accuracy_path, "w", encoding="utf-8") as file:
+            file.write(str(accuracy))
+
+    @staticmethod
+    def load_accuracy(filename=None):
+        """Loading the test accuracy from a file."""
+        model_path = Config().params["model_path"]
+        model_name = Config().trainer.model_name
+
+        if filename is not None:
+            accuracy_path = f"{model_path}/{filename}"
+        else:
+            accuracy_path = f"{model_path}/{model_name}.acc"
+
+        with open(accuracy_path, "r", encoding="utf-8") as file:
+            accuracy = float(file.read())
+
+        return accuracy
+
+    def pause_training(self):
+        """Remove files of running trainers."""
+        if hasattr(Config().trainer, "max_concurrency"):
+            model_name = Config().trainer.model_name
+            model_path = Config().params["model_path"]
+            model_file = f"{model_path}/{model_name}_{self.client_id}_{Config().params['run_id']}.pth"
+            accuracy_file = f"{model_path}/{model_name}_{self.client_id}_{Config().params['run_id']}.acc"
+
+            if os.path.exists(model_file):
+                os.remove(model_file)
+                os.remove(model_file + ".pkl")
+
+            if os.path.exists(accuracy_file):
+                os.remove(accuracy_file)
+
+    @abstractmethod
+    def train(self, trainset, sampler, **kwargs) -> float:
+        """The main training loop in a federated learning workload.
+
+        Arguments:
+        trainset: The training dataset.
+        sampler: the sampler that extracts a partition for this client.
+
+        Returns:
+        float: The training time.
+        """
+
+    @abstractmethod
+    def test(self, testset, sampler=None, **kwargs) -> float:
+        """Testing the model using the provided test dataset.
+
+        Arguments:
+        testset: The test dataset.
+        sampler: The sampler that extracts a partition of the test dataset.
+        """