plato-learn 1.1__py3-none-any.whl

plato/__init__.py

@@ -0,0 +1 @@
+__version__ = "1.0"

plato/algorithms/base.py

@@ -0,0 +1,45 @@
+"""
+Base class for algorithms.
+"""
+
+import os
+from abc import ABC, abstractmethod
+
+from plato.trainers.base import Trainer
+
+
+class Algorithm(ABC):
+    """Base class for all the algorithms."""
+
+    def __init__(self, trainer: Trainer):
+        """Initializes the algorithm with the provided model and trainer.
+
+        Arguments:
+        trainer: The trainer for the model, which is a trainers.base.Trainer class.
+        model: The model to train.
+        """
+        super().__init__()
+        self.trainer = trainer
+        self.model = trainer.model
+        self.client_id = 0
+
+    def __repr__(self):
+        if self.client_id == 0:
+            return f"Server #{os.getpid()}"
+        else:
+            return f"Client #{self.client_id}"
+
+    def set_client_id(self, client_id):
+        """Sets the client ID."""
+        self.client_id = client_id
+
+    @abstractmethod
+    def extract_weights(self, model=None):
+        """Extracts weights from a model passed in as a parameter."""
+
+    @abstractmethod
+    def load_weights(self, weights):
+        """Loads the model weights passed in as a parameter."""
+
+    async def aggregate_weights(self, baseline_weights, weights_received, **kwargs):
+        """Aggregates the weights received into baseline weights (optional)."""

plato/algorithms/fedavg.py

@@ -0,0 +1,48 @@
+"""
+The federated averaging algorithm for PyTorch.
+"""
+
+from collections import OrderedDict
+
+from plato.algorithms import base
+
+
+class Algorithm(base.Algorithm):
+    """PyTorch-based federated averaging algorithm, used by both the client and the server."""
+
+    def compute_weight_deltas(self, baseline_weights, weights_received):
+        """Compute the deltas between baseline weights and weights received."""
+        # Calculate updates from the received weights
+        deltas = []
+        for weight in weights_received:
+            delta = OrderedDict()
+            for name, current_weight in weight.items():
+                baseline = baseline_weights[name]
+
+                # Calculate update
+                _delta = current_weight - baseline
+                delta[name] = _delta
+            deltas.append(delta)
+
+        return deltas
+
+    def update_weights(self, deltas):
+        """Updates the existing model weights from the provided deltas."""
+        baseline_weights = self.extract_weights()
+
+        updated_weights = OrderedDict()
+        for name, weight in baseline_weights.items():
+            updated_weights[name] = weight + deltas[name]
+
+        return updated_weights
+
+    def extract_weights(self, model=None):
+        """Extracts weights from the model."""
+        if model is None:
+            return self.model.cpu().state_dict()
+        else:
+            return model.cpu().state_dict()
+
+    def load_weights(self, weights):
+        """Loads the model weights passed in as a parameter."""
+        self.model.load_state_dict(weights, strict=True)

plato/algorithms/fedavg_gan.py

@@ -0,0 +1,79 @@
+"""
+The federated averaging algorithm for GAN model.
+"""
+
+from collections import OrderedDict
+
+from plato.algorithms import fedavg
+from plato.trainers.base import Trainer
+
+
+class Algorithm(fedavg.Algorithm):
+    """Federated averaging algorithm for GAN models, used by both the client and the server."""
+
+    def __init__(self, trainer: Trainer):
+        super().__init__(trainer=trainer)
+        self.generator = self.model.generator
+        self.discriminator = self.model.discriminator
+
+    def compute_weight_deltas(self, weights_received):
+        """Extract the weights received from a client and compute the updates."""
+        baseline_weights_gen, baseline_weights_disc = self.extract_weights()
+
+        deltas = []
+        for weight_gen, weight_disc in weights_received:
+            delta_gen = OrderedDict()
+            for name, current_weight in weight_gen.items():
+                baseline = baseline_weights_gen[name]
+
+                delta = current_weight - baseline
+                delta_gen[name] = delta
+
+            delta_disc = OrderedDict()
+            for name, current_weight in weight_disc.items():
+                baseline = baseline_weights_disc[name]
+
+                delta = current_weight - baseline
+                delta_disc[name] = delta
+
+            deltas.append((delta_gen, delta_disc))
+
+        return deltas
+
+    def update_weights(self, deltas):
+        """Update the existing model weights."""
+        baseline_weights_gen, baseline_weights_disc = self.extract_weights()
+        update_gen, update_disc = deltas
+
+        updated_weights_gen = OrderedDict()
+        for name, weight in baseline_weights_gen.items():
+            updated_weights_gen[name] = weight + update_gen[name]
+
+        updated_weights_disc = OrderedDict()
+        for name, weight in baseline_weights_disc.items():
+            updated_weights_disc[name] = weight + update_disc[name]
+
+        return updated_weights_gen, updated_weights_disc
+
+    def extract_weights(self, model=None):
+        """Extract weights from the model."""
+        generator = self.generator
+        discriminator = self.discriminator
+        if model is not None:
+            generator = model.generator
+            discriminator = model.discriminator
+
+        gen_weight = generator.cpu().state_dict()
+        disc_weight = discriminator.cpu().state_dict()
+
+        return gen_weight, disc_weight
+
+    def load_weights(self, weights):
+        """Load the model weights passed in as a parameter."""
+        weights_gen, weights_disc = weights
+        # The client might only receive one or none of the Generator
+        # and Discriminator model weight.
+        if weights_gen is not None:
+            self.generator.load_state_dict(weights_gen, strict=True)
+        if weights_disc is not None:
+            self.discriminator.load_state_dict(weights_disc, strict=True)

plato/algorithms/fedavg_personalized.py

@@ -0,0 +1,48 @@
+"""
+A personalized federate learning algorithm that loads and saves local layers of a model.
+"""
+
+import os
+import logging
+
+import torch
+from plato.algorithms import fedavg
+from plato.config import Config
+
+
+class Algorithm(fedavg.Algorithm):
+    """
+    A personalized federate learning algorithm that loads and saves local layers
+    of a model.
+    """
+
+    def load_weights(self, weights):
+        """
+        Loads local layers included in `local_layer_names` to the received weights which
+        will be loaded to the model
+        """
+        if hasattr(Config().algorithm, "local_layer_names"):
+            # Get the filename of the previous saved local layer
+            model_path = Config().params["model_path"]
+            model_name = Config().trainer.model_name
+            filename = f"{model_path}/{model_name}_{self.client_id}_local_layers.pth"
+
+            # Load local layers to the weights when the file exists
+            if os.path.exists(filename):
+                local_layers = torch.load(filename, map_location=torch.device("cpu"))
+
+                # Update the received weights with the loaded local layers
+                weights.update(local_layers)
+
+                logging.info(
+                    "[Client #%d] Replaced portions of the global model with local layers.",
+                    self.trainer.client_id,
+                )
+
+        self.model.load_state_dict(weights, strict=True)
+
+    def save_local_layers(self, local_layers, filename):
+        """
+        Save local layers to a file with the filename provided.
+        """
+        torch.save(local_layers, filename)

plato/algorithms/mistnet.py

@@ -0,0 +1,52 @@
+"""
+The PyTorch-based MistNet algorithm, used by both the client and the server.
+
+Reference:
+
+P. Wang, et al. "MistNet: Towards Private Neural Network Training with Local
+Differential Privacy," found in docs/papers.
+"""
+
+import logging
+import time
+
+import torch
+from plato.algorithms import fedavg
+from plato.datasources import feature_dataset
+
+
+class Algorithm(fedavg.Algorithm):
+    """The PyTorch-based MistNet algorithm, used by both the client and the
+    server.
+    """
+
+    def extract_features(self, dataset, sampler):
+        """Extracting features using layers before the cut_layer.
+
+        dataset: The training or testing dataset.
+        """
+        self.model.eval()
+
+        data_loader = self.trainer.get_train_loader(
+            batch_size=1, trainset=dataset, sampler=sampler.get(), extract_features=True
+        )
+
+        tic = time.perf_counter()
+
+        features_dataset = []
+
+        for inputs, targets, *__ in data_loader:
+            with torch.no_grad():
+                logits = self.model.forward_to(inputs)
+            features_dataset.append((logits, targets))
+
+        toc = time.perf_counter()
+        logging.info("[Client #%s] Time used: %.2f seconds.", self.client_id, toc - tic)
+
+        return features_dataset
+
+    def train(self, trainset, sampler):
+        """Train the neural network model after the cut layer."""
+        self.trainer.train(
+            feature_dataset.FeatureDataset(trainset.feature_dataset), sampler
+        )

plato/algorithms/registry.py

@@ -0,0 +1,39 @@
+"""
+The registry for algorithms that contains framework-specific implementations.
+
+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.algorithms import (
+    fedavg,
+    mistnet,
+    fedavg_gan,
+    fedavg_personalized,
+    split_learning,
+)
+
+registered_algorithms = {
+    "fedavg": fedavg.Algorithm,
+    "mistnet": mistnet.Algorithm,
+    "fedavg_gan": fedavg_gan.Algorithm,
+    "fedavg_personalized": fedavg_personalized.Algorithm,
+    "split_learning": split_learning.Algorithm,
+}
+
+
+def get(trainer=None):
+    """Get the algorithm with the provided type."""
+    algorithm_type = Config().algorithm.type
+
+    if algorithm_type in registered_algorithms:
+        logging.info("Algorithm: %s", algorithm_type)
+        registered_alg = registered_algorithms[algorithm_type](trainer)
+        return registered_alg
+    else:
+        raise ValueError(f"No such algorithm: {algorithm_type}")

plato/algorithms/split_learning.py

@@ -0,0 +1,89 @@
+"""
+A federated learning algorithm 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 plato.algorithms import fedavg
+from plato.config import Config
+from plato.datasources import feature_dataset
+
+
+class Algorithm(fedavg.Algorithm):
+    """The PyTorch-based split learning algorithm, used by both the client and the
+    server.
+    """
+
+    def extract_features(self, dataset, sampler):
+        """Extracting features using layers before the cut_layer."""
+        self.model.to(self.trainer.device)
+        self.model.eval()
+
+        tic = time.perf_counter()
+
+        features_dataset = []
+
+        inputs, targets = self.trainer.get_train_samples(
+            Config().trainer.batch_size, dataset, sampler
+        )
+        inputs = inputs.to(self.trainer.device)
+        targets = targets.to(self.trainer.device)
+        outputs, targets = self.trainer.forward_to_intermediate_feature(inputs, targets)
+        features_dataset.append((outputs, targets))
+
+        toc = time.perf_counter()
+        logging.warning(
+            "[Client #%d] Features extracted from %s examples in %.2f seconds.",
+            self.client_id,
+            Config().trainer.batch_size,
+            toc - tic,
+        )
+
+        return features_dataset, toc - tic
+
+    def complete_train(self, gradients):
+        """Update the model on the client/device with the gradients received
+        from the server.
+        """
+        tic = time.perf_counter()
+
+        # Retrieve the training samples and let trainer do the training
+        samples, sampler = self.trainer.retrieve_train_samples()
+        self.trainer.load_gradients(gradients)
+        self.train(samples, sampler)
+
+        toc = time.perf_counter()
+        logging.warning(
+            "[Client #%d] Training completed in %.2f seconds.",
+            self.client_id,
+            toc - tic,
+        )
+
+        return toc - tic
+
+    def train(self, trainset, sampler):
+        """General training method that trains model with provided trainset and sampler."""
+        self.trainer.train(
+            feature_dataset.FeatureDataset(trainset.feature_dataset), sampler
+        )
+
+    def update_weights_before_cut(self, weights):
+        """Update the weights before cut layer, called when testing accuracy."""
+        current_weights = self.extract_weights()
+        current_weights = self.trainer.update_weights_before_cut(
+            current_weights, weights
+        )
+        self.load_weights(current_weights)

plato/callbacks/client.py

@@ -0,0 +1,56 @@
+"""
+Defines the ClientCallback class, which is the abstract base class to be subclassed
+when creating new client callbacks.
+
+Defines a default callback to print local training progress.
+"""
+
+from abc import ABC
+import logging
+
+
+class ClientCallback(ABC):
+    """
+    The abstract base class to be subclassed when creating new client callbacks.
+    """
+
+    def on_inbound_received(self, client, inbound_processor):
+        """
+        Event called before inbound processors start to process data.
+        """
+
+    def on_inbound_processed(self, client, data):
+        """
+        Event called when payload was processed by inbound processors.
+        """
+
+    def on_outbound_ready(self, client, report, outbound_processor):
+        """
+        Event called before outbound processors start to process data.
+        """
+
+
+class LogProgressCallback(ClientCallback):
+    """
+    A callback which prints a message when needed.
+    """
+
+    def on_inbound_received(self, client, inbound_processor):
+        """
+        Event called before inbound processors start to process data.
+        """
+        logging.info("[%s] Start to process inbound data.", client)
+
+    def on_inbound_processed(self, client, data):
+        """
+        Event called when payload was processed by inbound processors.
+        """
+        logging.info("[%s] Inbound data has been processed.", client)
+
+    def on_outbound_ready(self, client, report, outbound_processor):
+        """
+        Event called before outbound processors start to process data.
+        """
+        logging.info(
+            "[%s] Outbound data is ready to be sent after being processed.", client
+        )

plato/callbacks/handler.py

@@ -0,0 +1,78 @@
+"""
+Defines the :class:`CallbackHandler`, which is responsible for calling a list of callbacks.
+"""
+
+
+class CallbackHandler:
+    """
+    The :class:`CallbackHandler` is responsible for calling a list of callbacks.
+    This class calls the callbacks in the order that they are given.
+    """
+
+    def __init__(self, callbacks):
+        self.callbacks = []
+        self.add_callbacks(callbacks)
+
+    def add_callbacks(self, callbacks):
+        """
+        Adds a list of callbacks to the callback handler.
+
+        :param callbacks: a list of instances of a subclass of :class:`TrainerCallback`.
+        """
+        for callback in callbacks:
+            self.add_callback(callback)
+
+    def add_callback(self, callback):
+        """
+        Adds a callback to the callback handler.
+
+        :param callback: an instance of a subclass of :class:`TrainerCallback`.
+        """
+        _callback = callback() if isinstance(callback, type) else callback
+        _callback_class = callback if isinstance(callback, type) else callback.__class__
+
+        if _callback_class in {c.__class__ for c in self.callbacks}:
+            existing_callbacks = "\n".join(cb for cb in self.callback_list)
+
+            raise ValueError(
+                f"You attempted to add multiple instances of the callback "
+                f"{_callback_class}.\n"
+                f"The list of callbacks already present is: {existing_callbacks}"
+            )
+        self.callbacks.append(_callback)
+
+    def __iter__(self):
+        return self.callbacks
+
+    def clear_callbacks(self):
+        """
+        Clears all the callbacks in the current list.
+        """
+        self.callbacks = []
+
+    @property
+    def callback_list(self):
+        """
+        Retruns the names for the current list of callbacks.
+        """
+        return [cb.__class__.__name__ for cb in self.callbacks]
+
+    def call_event(self, event, *args, **kwargs):
+        """
+        For each callback which has been registered, sequentially call the method corresponding
+        to the given event.
+
+        :param event: The event corresponding to the method to call on each callback.
+        :param args: a list of arguments to be passed to each callback.
+        :param kwargs: a list of keyword arguments to be passed to each callback.
+        """
+        for callback in self.callbacks:
+            try:
+                getattr(callback, event)(
+                    *args,
+                    **kwargs,
+                )
+            except AttributeError as exc:
+                raise ValueError(
+                    "The callback method has not been implemented"
+                ) from exc

plato/callbacks/server.py

@@ -0,0 +1,139 @@
+"""
+Defines the ServerCallback class, which is the abstract base class to be subclassed
+when creating new server callbacks.
+
+Defines a default callback to print training progress.
+"""
+
+import logging
+import os
+from abc import ABC
+from plato.config import Config
+from plato.utils import csv_processor, fonts
+
+
+class ServerCallback(ABC):
+    """
+    The abstract base class to be subclassed when creating new server callbacks.
+    """
+
+    def __init__(self):
+        """
+        Initializer.
+        """
+
+    def on_weights_received(self, server, weights_received):
+        """
+        Event called after the updated weights have been received.
+        """
+
+    def on_weights_aggregated(self, server, updates):
+        """
+        Event called after the updated weights have been aggregated.
+        """
+
+    def on_clients_selected(self, server, selected_clients, **kwargs):
+        """
+        Event called after a new client arrived.
+        """
+
+    def on_clients_processed(self, server, **kwargs):
+        """Additional work to be performed after client reports have been processed."""
+
+    def on_training_will_start(self, server, **kwargs):
+        """
+        Event called before selecting clients for the first round of training.
+        """
+
+    def on_server_will_close(self, server, **kwargs):
+        """
+        Event called at the start of closing the server.
+        """
+
+
+class LogProgressCallback(ServerCallback):
+    """
+    A callback which prints a message when needed.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+        recorded_items = Config().params["result_types"]
+        self.recorded_items = [x.strip() for x in recorded_items.split(",")]
+
+        # Initialize the .csv file for logging runtime results
+        result_csv_file = f"{Config().params['result_path']}/{os.getpid()}.csv"
+        csv_processor.initialize_csv(
+            result_csv_file, self.recorded_items, Config().params["result_path"]
+        )
+
+        logging.info(
+            fonts.colourize(
+                f"[{os.getpid()}] Logging runtime results to: {result_csv_file}."
+            )
+        )
+
+    def on_weights_received(self, server, weights_received):
+        """
+        Event called after the updated weights have been received.
+        """
+        logging.info("[%s] Updated weights have been received.", server)
+
+    def on_weights_aggregated(self, server, updates):
+        """
+        Event called after the updated weights have been aggregated.
+        """
+        logging.info("[%s] Finished aggregating updated weights.", server)
+
+    def on_clients_selected(self, server, selected_clients):
+        """
+        Event called after clients have been selected in each round.
+        """
+
+    def on_clients_processed(self, server, **kwargs):
+        """Additional work to be performed after client reports have been processed."""
+        # Record results into a .csv file
+        new_row = []
+        for item in self.recorded_items:
+            item_value = server.get_logged_items()[item]
+            new_row.append(item_value)
+
+        result_csv_file = f"{Config().params['result_path']}/{os.getpid()}.csv"
+        csv_processor.write_csv(result_csv_file, new_row)
+
+        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
+            )
+        ):
+            # Updates the log for client test accuracies
+            accuracy_csv_file = (
+                f"{Config().params['result_path']}/{os.getpid()}_accuracy.csv"
+            )
+
+            for update in server.updates:
+                accuracy_row = [
+                    server.current_round,
+                    update.client_id,
+                    update.report.accuracy,
+                ]
+                csv_processor.write_csv(accuracy_csv_file, accuracy_row)
+
+        logging.info("[%s] All client reports have been processed.", server)
+
+    def on_training_will_start(self, server, **kwargs):
+        """
+        Event called before selecting clients for the first round of training.
+        """
+        logging.info("[%s] Starting training.", server)
+
+    def on_server_will_close(self, server, **kwargs):
+        """
+        Event called at the start of closing the server.
+        """
+        logging.info("[%s] Closing the server.", server)