plato-learn 1.1__py3-none-any.whl

plato/processors/model_quantize_qsgd.py

@@ -0,0 +1,82 @@
+"""
+Implements a Processor to quantize and compress upload models.
+
+In more detail, this processor first quantizes each upload parameter under
+the given quantization level. Next, compress and store each quantized value.
+Hence, these 32-bit parameters can be converted into 8-bit parameters.
+
+Reference:
+
+Alistarh, D., Grubic, D., Li, J., Tomioka, R., & Vojnovic, M. (2017).
+"QSGD: Communication-efficient SGD via gradient quantization and encoding."
+Advances in neural information processing systems.
+
+https://proceedings.neurips.cc/paper/2017/file/6c340f25839e6acdc73414517203f5f0-Paper.pdf
+"""
+
+import random
+from struct import pack, unpack
+from typing import Any
+
+import torch
+
+from plato.processors import model
+
+
+class Processor(model.Processor):
+    """
+    Implements a Processor to quantize model parameters with QSGD.
+    """
+
+    def __init__(self, quantization_level=64, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+        self.quantization_level = quantization_level  # must <= 128!
+
+    def _process_layer(self, layer: Any) -> Any:
+        """Quantizes each individual layer of the model with QSGD."""
+
+        def add_prob(prob: Any) -> Any:
+            """Adds 1 to the corresponding positions with given probability."""
+            size = prob.size()
+            prob = prob.reshape(-1)
+            random.seed()
+            for count, value in enumerate(prob):
+                if random.random() <= value:
+                    prob[count] = 1
+                else:
+                    prob[count] = 0
+            return torch.reshape(prob, size)
+
+        def handler(tensor: Any) -> Any:
+            """Handler function for the compression of quantized values."""
+            content = b""
+            tensor = tensor.reshape(-1)
+            for _, value in enumerate(tensor):
+                num = value.item()
+                if num < 0:
+                    num = abs(num) ^ unpack("!i", b"\x00\x00\x00\x80")[0]
+                content += pack("!I", num)[3:4]  # present each parameter in 1 byte
+            return content
+
+        # Step 1: quantization
+        tuning_param = self.quantization_level - 1  # tuning parameter
+        max_v = torch.max(abs(layer))  # max absolute value
+        neg = (-1) * layer.lt(0) + 1 * layer.ge(0)
+        ratio = abs(layer) / max_v  # |v_i| / ||v||
+        level = (ratio * tuning_param - 1).ceil()
+        zeta = level + add_prob(ratio * tuning_param - level)
+        zeta = zeta.mul(neg).to(int)
+
+        # Step 2: handle the header
+        output = pack("!f", max_v.item())  # ! represents for big-endian
+        output += pack("!I", zeta.numel())
+        dimensions = len(zeta.size())
+        output += pack("!h", dimensions)
+        for i in range(dimensions):
+            output += pack("!h", zeta.size(i))
+
+        # Step 3: handle the content, each consists of 1 sign bit followed by 7 bits
+        output += handler(zeta)
+
+        return output

plato/processors/model_randomized_response.py

@@ -0,0 +1,34 @@
+"""
+Implements a Processor for applying local differential privacy using randomized response.
+"""
+
+import torch
+
+from plato.config import Config
+from plato.processors import model
+from plato.utils import unary_encoding
+
+
+class Processor(model.Processor):
+    """
+    Implements a Processor for applying local differential privacy using randomized response.
+    """
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+    def _process_layer(self, layer: torch.Tensor) -> torch.Tensor:
+        if Config().algorithm.epsilon is None:
+            return layer
+
+        epsilon = Config().algorithm.epsilon
+
+        # Apply randomized response as the local differential privacy mechanism
+        layer = layer.detach().cpu().numpy()
+
+        layer = unary_encoding.encode(layer)
+        layer = unary_encoding.randomize(layer, epsilon)
+
+        layer = torch.tensor(layer, dtype=torch.float32)
+
+        return layer

plato/processors/outbound_feature_ndarrays.py

@@ -0,0 +1,38 @@
+"""
+Implements a Processor for converting MistNet features from PyTorch tensors to numpy ndarrays.
+"""
+
+import logging
+from typing import Any
+
+from plato.processors import base
+
+
+class Processor(base.Processor):
+    """
+    Implements a Processor for converting MistNet features from PyTorch tensors to numpy ndarrays.
+    This is used only by MistNet clients at this time.
+    """
+
+    def __init__(self, client_id=None, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+        self.client_id = client_id
+
+    def process(self, data: Any) -> Any:
+        """
+        Converts MistNet features from PyTorch tensors to numpy ndarrays.
+        """
+        feature_dataset = []
+
+        for logit, target in data:
+            feature_dataset.append(
+                (logit.detach().cpu().numpy(), target.detach().cpu().numpy())
+            )
+
+        logging.info(
+            "[Client #%d] Features converted from PyTorch tensors to ndarrays.",
+            self.client_id,
+        )
+
+        return feature_dataset

plato/processors/pipeline.py

@@ -0,0 +1,26 @@
+"""
+Implements a pipeline of processors for data payloads to pass through.
+"""
+
+from typing import Any, List
+
+from plato.processors import base
+
+
+class Processor(base.Processor):
+    """
+    Pipelining a list of Processors from the configuration file.
+    """
+
+    def __init__(self, processors: List[base.Processor], *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.processors = processors
+
+    def process(self, data: Any) -> Any:
+        """
+        Implementing a pipeline of Processors for data payloads.
+        """
+        for processor in self.processors:
+            data = processor.process(data)
+
+        return data

plato/processors/registry.py

@@ -0,0 +1,124 @@
+"""
+This registry for Processors contains framework-specific implementations of
+Processors for data payloads.
+
+Having a registry of all available classes is convenient for retrieving an instance
+based on a configuration at run-time.
+"""
+
+import logging
+from typing import Tuple
+
+from plato.config import Config
+from plato.processors import (
+    base,
+    compress,
+    decompress,
+    feature_dequantize,
+    feature_gaussian,
+    feature_laplace,
+    feature_quantize,
+    feature_randomized_response,
+    feature_unbatch,
+    inbound_feature_tensors,
+    model_compress,
+    model_decompress,
+    model_deepcopy,
+    model_dequantize,
+    model_dequantize_qsgd,
+    model_quantize,
+    model_quantize_qsgd,
+    model_randomized_response,
+    outbound_feature_ndarrays,
+    pipeline,
+    structured_pruning,
+    unstructured_pruning,
+)
+
+registered_processors = {
+    "base": base.Processor,
+    "compress": compress.Processor,
+    "decompress": decompress.Processor,
+    "feature_randomized_response": feature_randomized_response.Processor,
+    "feature_gaussian": feature_gaussian.Processor,
+    "feature_laplace": feature_laplace.Processor,
+    "feature_quantize": feature_quantize.Processor,
+    "feature_dequantize": feature_dequantize.Processor,
+    "feature_unbatch": feature_unbatch.Processor,
+    "inbound_feature_tensors": inbound_feature_tensors.Processor,
+    "outbound_feature_ndarrays": outbound_feature_ndarrays.Processor,
+    "model_deepcopy": model_deepcopy.Processor,
+    "model_quantize": model_quantize.Processor,
+    "model_dequantize": model_dequantize.Processor,
+    "model_compress": model_compress.Processor,
+    "model_quantize_qsgd": model_quantize_qsgd.Processor,
+    "model_decompress": model_decompress.Processor,
+    "model_dequantize_qsgd": model_dequantize_qsgd.Processor,
+    "model_randomized_response": model_randomized_response.Processor,
+    "structured_pruning": structured_pruning.Processor,
+    "unstructured_pruning": unstructured_pruning.Processor,
+}
+
+
+def register_he_processors():
+    """Register homomorphic encryption processors if needed."""
+
+
+def get(
+    user: str, processor_kwargs=None, **kwargs
+) -> Tuple[pipeline.Processor, pipeline.Processor]:
+    """Get an instance of the processor."""
+    outbound_processors = []
+    inbound_processors = []
+
+    assert user in ("Server", "Client")
+
+    if user == "Server":
+        config = Config().server
+    else:
+        config = Config().clients
+
+    if hasattr(config, "outbound_processors") and isinstance(
+        config.outbound_processors, list
+    ):
+        outbound_processors = config.outbound_processors
+
+    if hasattr(config, "inbound_processors") and isinstance(
+        config.inbound_processors, list
+    ):
+        inbound_processors = config.inbound_processors
+
+    for processor in outbound_processors:
+        logging.info("%s: Using Processor for sending payload: %s", user, processor)
+    for processor in inbound_processors:
+        logging.info("%s: Using Processor for receiving payload: %s", user, processor)
+
+    # Check if HE processors are needed based on server configuration
+    if hasattr(config, "type") and config.type == "fedavg_he":
+        # FedAvg server with homomorphic encryption needs to import tenseal,
+        # which is not available on all platforms such as macOS
+        from plato.processors import model_decrypt, model_encrypt
+
+        registered_processors.update(
+            {
+                "model_encrypt": model_encrypt.Processor,
+                "model_decrypt": model_decrypt.Processor,
+            }
+        )
+
+        logging.info("%s: Using homomorphic encryption processors.", user)
+
+    def map_f(name):
+        if processor_kwargs is not None and name in processor_kwargs:
+            this_kwargs = {**kwargs, **(processor_kwargs[name])}
+        else:
+            this_kwargs = kwargs
+
+        return registered_processors[name](name=name, **this_kwargs)
+
+    outbound_processors = list(map(map_f, outbound_processors))
+    inbound_processors = list(map(map_f, inbound_processors))
+
+    return pipeline.Processor(outbound_processors), pipeline.Processor(
+        inbound_processors
+    )

plato/processors/structured_pruning.py

@@ -0,0 +1,57 @@
+"""
+Processor for structured pruning of model weights.
+"""
+
+import logging
+
+import torch
+import torch.nn.utils.prune as prune
+
+from plato.processors import model
+
+
+class Processor(model.Processor):
+    """
+    A processor for the structured pruning of model weights.
+    """
+
+    def __init__(
+        self, pruning_method="ln", amount=0.2, norm=1, dim=-1, **kwargs
+    ) -> None:
+        super().__init__(**kwargs)
+
+        self.pruning_method = pruning_method
+        self.amount = amount
+        self.norm = norm
+        self.dim = dim
+        self.model = None
+
+    def process(self, data):
+        """
+        Processes structured pruning of model weights layer by layer.
+        """
+        self.model = self.trainer.model
+
+        for _, module in self.model.named_modules():
+            if isinstance(module, torch.nn.Conv2d) or isinstance(
+                module, torch.nn.Linear
+            ):
+                if self.pruning_method == "ln":
+                    prune.ln_structured(
+                        module, "weight", self.amount, n=self.norm, dim=self.dim
+                    )
+                elif self.pruning_method == "random":
+                    prune.random_structured(module, "weight", self.amount, dim=self.dim)
+                prune.remove(module, "weight")
+
+        output = self.model.cpu().state_dict()
+
+        if self.client_id is None:
+            logging.info("[Server #%d] Structured pruning applied.", self.server_id)
+        else:
+            logging.info("[Client #%d] Structured pruning applied.", self.client_id)
+
+        return output
+
+    def _process_layer(self, layer: torch.Tensor) -> torch.Tensor:
+        """No need to process individual layer of the model"""

plato/processors/unstructured_pruning.py

@@ -0,0 +1,73 @@
+"""
+Implements a Processor for global unstructured pruning of model weights.
+"""
+
+import logging
+from typing import Any
+
+import torch
+import torch.nn.utils.prune as prune
+
+from plato.processors import model
+
+
+class Processor(model.Processor):
+    """
+    Implements a Processor for global unstructured pruning of model weights.
+    """
+
+    def __init__(
+        self,
+        parameters_to_prune=None,
+        pruning_method=prune.L1Unstructured,
+        amount=0.2,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+
+        self.parameters_to_prune = parameters_to_prune
+        self.pruning_method = pruning_method
+        self.amount = amount
+        self.model = None
+
+    def process(self, data: Any) -> Any:
+        """
+        Proceesses global unstructured pruning on model weights.
+        """
+
+        self.model = self.trainer.model
+
+        if self.parameters_to_prune is None:
+            self.parameters_to_prune = []
+            for _, module in self.model.named_modules():
+                if isinstance(module, torch.nn.Conv2d) or isinstance(
+                    module, torch.nn.Linear
+                ):
+                    self.parameters_to_prune.append((module, "weight"))
+
+        prune.global_unstructured(
+            self.parameters_to_prune,
+            pruning_method=self.pruning_method,
+            amount=self.amount,
+        )
+
+        for module, name in self.parameters_to_prune:
+            prune.remove(module, name)
+
+        output = self.model.cpu().state_dict()
+
+        if self.client_id is None:
+            logging.info(
+                "[Server #%d] Global unstructured pruning applied.",
+                self.server_id,
+            )
+        else:
+            logging.info(
+                "[Client #%d] Global unstructured pruning applied.",
+                self.client_id,
+            )
+
+        return output
+
+    def _process_layer(self, layer: torch.Tensor) -> torch.Tensor:
+        return layer

plato/samplers/all_inclusive.py

@@ -0,0 +1,41 @@
+"""
+Samples all the data from a dataset. Applicable in cases where the dataset comes from
+local sources only. Used by the Federated EMNIST dataset and the MistNet server.
+"""
+
+import random
+
+from plato.samplers import base
+from plato.config import Config
+
+
+class Sampler(base.Sampler):
+    """Create a data sampler that samples all the data in the dataset.
+    Used by the MistNet server.
+    """
+
+    def __init__(self, datasource, client_id=0, testing=False):
+        super().__init__()
+        self.client_id = client_id
+
+        if testing:
+            all_inclusive = range(len(datasource.get_test_set()))
+            if hasattr(Config().data, "testset_size"):
+                self.data_samples = random.sample(
+                    all_inclusive, Config().data.testset_size
+                )
+            else:
+                self.data_samples = all_inclusive
+        else:
+            self.data_samples = range(len(datasource.get_train_set()))
+
+    def get(self):
+        import torch
+
+        gen = torch.Generator()
+        gen.manual_seed(self.random_seed)
+        return torch.utils.data.SubsetRandomSampler(self.data_samples, generator=gen)
+
+    def num_samples(self):
+        """Returns the length of the dataset after sampling."""
+        return len(self.data_samples)

plato/samplers/base.py

@@ -0,0 +1,31 @@
+"""
+Base class for sampling data so that a dataset can be divided across the clients.
+"""
+
+import os
+from abc import abstractmethod
+
+from plato.config import Config
+
+
+class Sampler:
+    """Base class for data samplers so that the dataset is divided into
+    partitions across the clients."""
+
+    def __init__(self):
+        if hasattr(Config().data, "random_seed"):
+            # Keeping random seed the same across the clients
+            # so that the experiments are reproducible
+            self.random_seed = Config().data.random_seed
+        else:
+            # The random seed will be different across different
+            # runs if it is not provided.
+            self.random_seed = os.getpid()
+
+    @abstractmethod
+    def get(self):
+        """Obtains an instance of the sampler."""
+
+    @abstractmethod
+    def num_samples(self):
+        """Returns the length of the dataset after sampling."""

plato/samplers/dirichlet.py

@@ -0,0 +1,81 @@
+"""
+Samples data from a dataset, biased across labels according to the Dirichlet distribution.
+"""
+
+import numpy as np
+import torch
+from torch.utils.data import WeightedRandomSampler, SubsetRandomSampler
+from plato.config import Config
+
+from plato.samplers import base
+
+
+class Sampler(base.Sampler):
+    """Create a data sampler for each client to use a divided partition of the
+    dataset, biased across labels according to the Dirichlet distribution."""
+
+    def __init__(self, datasource, client_id, testing):
+        super().__init__()
+
+        # Different clients should have a different bias across the labels & partition size
+        np.random.seed(self.random_seed * int(client_id))
+
+        # Concentration parameter to be used in the Dirichlet distribution
+        concentration = (
+            Config().data.concentration
+            if hasattr(Config().data, "concentration")
+            else 1.0
+        )
+
+        if testing:
+            target_list = datasource.get_test_set().targets
+        else:
+            # The list of labels (targets) for all the examples
+            target_list = datasource.targets()
+
+        class_list = datasource.classes()
+
+        target_proportions = np.random.dirichlet(
+            np.repeat(concentration, len(class_list))
+        )
+
+        if np.isnan(np.sum(target_proportions)):
+            target_proportions = np.repeat(0, len(class_list))
+            target_proportions[np.random.randint(0, len(class_list))] = 1
+
+        self.sample_weights = target_proportions[target_list]
+
+    def num_samples(self) -> int:
+        """Returns the length of the dataset after sampling."""
+        sampled_size = Config().data.partition_size
+
+        # Variable partition size across clients
+        if hasattr(Config().data, "partition_distribution"):
+            dist = Config().data.partition_distribution
+
+            if dist.distribution.lower() == "uniform":
+                sampled_size *= np.random.uniform(dist.low, dist.high)
+
+            if dist.distribution.lower() == "normal":
+                sampled_size *= np.random.normal(dist.mean, dist.high)
+
+            sampled_size = int(sampled_size)
+
+        return sampled_size
+
+    def get(self):
+        """Obtains an instance of the sampler."""
+        gen = torch.Generator()
+        gen.manual_seed(self.random_seed)
+
+        # Samples without replacement using the sample weights
+        subset_indices = list(
+            WeightedRandomSampler(
+                weights=self.sample_weights,
+                num_samples=self.num_samples(),
+                replacement=False,
+                generator=gen,
+            )
+        )
+
+        return SubsetRandomSampler(subset_indices, generator=gen)