plato-learn 1.1__py3-none-any.whl

plato/processors/feature_additive_noise.py

@@ -0,0 +1,48 @@
+"""
+Implements a Processor for applying local differential privacy using additive noise mechanism.
+"""
+
+import logging
+from typing import Any
+
+import numpy
+
+from plato.processors import feature
+
+
+class Processor(feature.Processor):
+    """
+    Implements a Processor for applying local differential privacy using additive noise mechanism.
+    """
+
+    methods = {
+        "gaussian": numpy.random.normal,
+        "laplace": numpy.random.laplace,
+    }
+
+    def __init__(self, method="", scale=None, **kwargs) -> None:
+        self._method = method
+
+        def func(logits, targets):
+            return (
+                Processor.methods[method](logits, scale),
+                targets,
+            )
+
+        super().__init__(method=func, **kwargs)
+
+    def process(self, data: Any) -> Any:
+        """
+        Implements a Processor for applying randomized response as the local differential privacy
+        mechanism.
+        """
+
+        output = super().process(data)
+
+        logging.info(
+            "[Client #%d] Local differential privacy (using the %s mechanism) applied.",
+            self.client_id,
+            self._method,
+        )
+
+        return output

plato/processors/feature_dequantize.py

@@ -0,0 +1,34 @@
+"""
+Implements a Processor for applying dequantization to MistNet PyTorch features.
+"""
+
+import logging
+from typing import Any
+
+import torch
+
+from plato.processors import base
+
+
+class Processor(base.Processor):
+    """
+    Implements a Processor for applying dequantization to MistNet PyTorch features.
+    """
+
+    def __init__(self, server_id=None, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+        self.server_id = server_id
+
+    def process(self, data: Any) -> Any:
+        """
+        Implements a Processor for applying dequantization to MistNet PyTorch features.
+        """
+        feature_dataset = []
+
+        for logit, target in data:
+            feature_dataset.append((torch.dequantize(logit), target))
+
+        logging.info("[Server #%d] Dequantized features.", self.server_id)
+
+        return feature_dataset

plato/processors/feature_gaussian.py

@@ -0,0 +1,17 @@
+"""
+Implements a Processor for applying local differential privacy using gaussian mechanism.
+"""
+
+import math
+
+from plato.processors import feature_additive_noise
+
+
+class Processor(feature_additive_noise.Processor):
+    """
+    Implements a Processor for applying local differential privacy using gaussian mechanism.
+    """
+
+    def __init__(self, epsilon=None, delta=None, sensitivity=None, **kwargs) -> None:
+        scale = 2 * math.log(1.25 / delta) * sensitivity**2 / epsilon**2
+        super().__init__(method="gaussian", scale=scale, **kwargs)

plato/processors/feature_laplace.py

@@ -0,0 +1,15 @@
+"""
+Implements a Processor for applying local differential privacy using laplace mechanism.
+"""
+
+from plato.processors import feature_additive_noise
+
+
+class Processor(feature_additive_noise.Processor):
+    """
+    Implements a Processor for applying local differential privacy using the Laplace mechanism.
+    """
+
+    def __init__(self, epsilon=None, sensitivity=None, **kwargs) -> None:
+        scale = sensitivity / epsilon
+        super().__init__(method="laplace", scale=scale, **kwargs)

plato/processors/feature_quantize.py

@@ -0,0 +1,34 @@
+"""
+Implements a Processor for applying quantization to MistNet PyTorch features.
+"""
+
+import logging
+from typing import Any
+
+import torch
+
+from plato.processors import feature
+
+
+class Processor(feature.Processor):
+    """
+    Implements a Processor for applying quantization to MistNet PyTorch features.
+    """
+
+    def __init__(self, scale=0.1, zero_point=10, dtype=torch.quint8, **kwargs) -> None:
+        def func(logits, targets):
+            logits = torch.quantize_per_tensor(logits, scale, zero_point, dtype)
+            return logits, targets
+
+        super().__init__(method=func, use_numpy=False, **kwargs)
+
+    def process(self, data: Any) -> Any:
+        """
+        Implements a Processor for applying quantization to MistNet PyTorch features.
+        """
+
+        output = super().process(data)
+
+        logging.info("[Client #%d] Quantization applied.", self.client_id)
+
+        return output

plato/processors/feature_randomized_response.py

@@ -0,0 +1,50 @@
+"""
+Implements a Processor for applying local differential privacy using randomized response.
+"""
+
+import logging
+from typing import Any
+
+from plato.config import Config
+from plato.processors import feature
+from plato.utils import unary_encoding
+
+
+class Processor(feature.Processor):
+    """
+    Implements a Processor for applying local differential privacy using randomized response.
+    """
+
+    def __init__(self, **kwargs) -> None:
+        def func(logits, targets):
+            logits = unary_encoding.encode(logits)
+
+            if Config().algorithm.epsilon is None:
+                return logits, targets
+
+            _randomize = getattr(self.trainer, "randomize", None)
+            epsilon = Config().algorithm.epsilon
+
+            if callable(_randomize):
+                logits = self.trainer.randomize(logits, targets, epsilon)
+            else:
+                logits = unary_encoding.randomize(logits, epsilon)
+
+            return logits, targets
+
+        super().__init__(method=func, **kwargs)
+
+    def process(self, data: Any) -> Any:
+        """
+        Implements a Processor for applying randomized response as the
+        local differential privacy mechanism.
+        """
+
+        output = super().process(data)
+
+        logging.info(
+            "[Client #%d] Local differential privacy (using randomized response) applied.",
+            self.client_id,
+        )
+
+        return output

plato/processors/feature_unbatch.py

@@ -0,0 +1,39 @@
+"""
+Implements a Processor for unbatching MistNet PyTorch features into the dataset form.
+"""
+
+import logging
+from typing import Any
+
+import numpy as np
+
+from plato.processors import base
+
+
+class Processor(base.Processor):
+    """
+    Implements a Processor for unbatching MistNet PyTorch features into the dataset form.
+    """
+
+    def __init__(self, client_id=None, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+        self.client_id = client_id
+
+    def process(self, data: Any) -> Any:
+        """
+        Implements a Processor for unbatching MistNet PyTorch features into the dataset form.
+        """
+        feature_dataset = []
+
+        for logits, targets in data:
+            for i in np.arange(logits.shape[0]):  # each sample in the batch
+                feature_dataset.append((logits[i].clone(), targets[i].clone()))
+
+        logging.info(
+            "[Client #%d] Features extracted from %s examples.",
+            self.client_id,
+            len(feature_dataset),
+        )
+
+        return feature_dataset

plato/processors/inbound_feature_tensors.py

@@ -0,0 +1,39 @@
+"""
+Implements a Processor for converting MistNet features from numpy ndarrays to PyTorch tensors.
+"""
+
+import logging
+from typing import Any
+
+import torch
+
+from plato.processors import base
+
+
+class Processor(base.Processor):
+    """
+    Implements a Processor for converting MistNet features from numpy ndarrays to PyTorch tensors.
+    """
+
+    def __init__(self, server_id=None, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+        self.server_id = server_id
+
+    def process(self, data: Any) -> Any:
+        """
+        Converts MistNet features from numpy ndarrays to PyTorch tensors.
+        """
+        feature_dataset = []
+
+        for logit, target in data:
+            # Uses torch.as_tensor() as opposed to torch.tensor() to avoid data copying
+            # according to https://pytorch.org/docs/stable/generated/torch.tensor.html
+            feature_dataset.append((torch.as_tensor(logit), torch.as_tensor(target)))
+
+        logging.info(
+            "[Server #%d] Features converted from ndarrays to PyTorch tensors.",
+            self.server_id,
+        )
+
+        return feature_dataset

plato/processors/model.py

@@ -0,0 +1,55 @@
+"""
+Base processor for processing PyTorch models.
+"""
+
+import logging
+import pickle
+import sys
+from typing import OrderedDict
+
+import torch
+
+from plato.processors import base
+
+
+class Processor(base.Processor):
+    """Base processor for processing PyTorch models."""
+
+    def __init__(self, client_id=None, server_id=None, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.client_id = client_id
+        self.server_id = server_id
+
+    def process(self, data: OrderedDict) -> OrderedDict:
+        """
+        Processes PyTorch model parameter.
+        The data is a state_dict of a PyTorch model.
+        """
+        old_data_size = sys.getsizeof(pickle.dumps(data))
+
+        new_data = OrderedDict()
+        for layer_name, layer_params in data.items():
+            new_data[layer_name] = self._process_layer(layer_params)
+
+        new_data_size = sys.getsizeof(pickle.dumps(new_data))
+
+        if self.client_id is None:
+            logging.info(
+                "[Server #%d] Processed the model and changed its size from %.2f MB to %.2f MB.",
+                self.server_id,
+                old_data_size / 1024**2,
+                new_data_size / 1024**2,
+            )
+        else:
+            logging.info(
+                "[Client #%d] Processed the model and changed its size from %.2f MB to %.2f MB.",
+                self.client_id,
+                old_data_size / 1024**2,
+                new_data_size / 1024**2,
+            )
+
+        return new_data
+
+    def _process_layer(self, layer: torch.Tensor) -> torch.Tensor:
+        """Processes an individual layer of the model."""
+        return layer

plato/processors/model_compress.py

@@ -0,0 +1,34 @@
+"""
+Implements a Processor for compressing model weights.
+"""
+
+import logging
+import pickle
+from typing import Any
+
+import zstd
+
+from plato.processors import model
+
+
+class Processor(model.Processor):
+    """
+    Implements a Processor for compressing model parameters.
+    """
+
+    def __init__(self, compression_level=1, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+        self.compression_level = compression_level
+
+    def process(self, data: Any) -> Any:
+        """Implements a Processor for compressing model parameters."""
+
+        output = zstd.compress(pickle.dumps(data), self.compression_level)
+
+        if self.client_id is None:
+            logging.info("[Server #%d] Compressed model parameters.", self.server_id)
+        else:
+            logging.info("[Client #%d] Compressed model parameters.", self.client_id)
+
+        return output

plato/processors/model_decompress.py

@@ -0,0 +1,37 @@
+"""
+Implements a Processor for decompressing model weights.
+"""
+
+import logging
+import pickle
+from typing import Any
+
+import zstd
+
+from plato.processors import model
+
+
+class Processor(model.Processor):
+    """
+    Implements a Processor for decompressing model parameters.
+    """
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+    def process(self, data: Any) -> Any:
+        """Implements a Processor for decompressing model parameters."""
+
+        output = pickle.loads(zstd.decompress(data))
+
+        if self.client_id is None:
+            logging.info(
+                "[Server #%d] Decompressed received model parameters.",
+                self.server_id,
+            )
+        else:
+            logging.info(
+                "[Client #%d] Decompressed received model parameters.",
+                self.client_id,
+            )
+        return output

plato/processors/model_decrypt.py

@@ -0,0 +1,41 @@
+"""
+A processor that decrypts model weights of MaskCrypt.
+"""
+
+from typing import Any
+
+import torch
+
+from plato.processors import model
+from plato.utils import homo_enc
+
+
+class Processor(model.Processor):
+    """
+    A processor that decrypts model tensors
+    """
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+        self.context = homo_enc.get_ckks_context()
+        weight_shapes = {}
+        para_nums = {}
+        extract_model = self.trainer.model.cpu().state_dict()
+
+        for key in extract_model.keys():
+            weight_shapes[key] = extract_model[key].size()
+            para_nums[key] = torch.numel(extract_model[key])
+
+        self.weight_shapes = weight_shapes
+        self.para_nums = para_nums
+
+    def process(self, data: Any) -> Any:
+        """Deserialize and decrypt the model weights."""
+        deserialized_weights = homo_enc.deserialize_weights(data, self.context)
+
+        output = homo_enc.decrypt_weights(
+            deserialized_weights, self.weight_shapes, self.para_nums
+        )
+
+        return output

plato/processors/model_deepcopy.py

@@ -0,0 +1,21 @@
+"""
+Processor for creating a deep copy of the PyTorch model state_dict.
+"""
+
+import copy
+
+import torch
+
+from plato.processors import model
+
+
+class Processor(model.Processor):
+    """
+    Processor for creating a deep copy of the PyTorch model state_dict.
+    """
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+    def _process_layer(self, layer: torch.Tensor):
+        return copy.deepcopy(layer)

plato/processors/model_dequantize.py

@@ -0,0 +1,18 @@
+"""
+Implements a Processor for dequantizing model parameters.
+"""
+
+import torch
+
+from plato.processors import model
+
+
+class Processor(model.Processor):
+    """
+    Implements a Processor for dequantizing model parameters.
+    """
+
+    def _process_layer(self, layer: torch.Tensor) -> torch.Tensor:
+        """Quantizes each individual layer of the model."""
+
+        return layer.to(torch.float32)

plato/processors/model_dequantize_qsgd.py

@@ -0,0 +1,61 @@
+"""
+Implements a Processor to decompress and dequantize upload models.
+
+In more detail, this processor first decompresses each received parameter.
+Next, dequantize each upload parameter under the given quantization level.
+Hence, 8-bit received parameters can be converted into 32-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
+"""
+
+from struct import unpack
+from typing import Any
+
+import torch
+
+from plato.processors import model
+
+
+class Processor(model.Processor):
+    """
+    Implements a Processor to dequantize model parameters quantized 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:
+        """Dequantizes each individual layer of the model."""
+
+        # Step 1: decompress the header
+        tuning_param = self.quantization_level - 1
+        max_v = unpack("!f", layer[0:4])[0]
+        numel = unpack("!I", layer[4:8])[0]
+        dimensions = unpack("!h", layer[8:10])[0]
+        size = []
+        for i in range(dimensions):
+            size.append(unpack("!h", layer[10 + 2 * i : 12 + 2 * i])[0])
+
+        # Step 2: decompress the content
+        layer = layer[10 + 2 * dimensions :]
+        zeta = []
+        prefix = b"\x00\x00\x00"
+        for i in range(numel):
+            tmp = unpack("!I", prefix + layer[i : i + 1])[0]
+            if tmp >= 128:
+                tmp = -1 * (tmp - 128)
+            zeta.append(tmp)
+        zeta = torch.tensor(zeta).reshape(size)
+
+        # Step 3: dequantize the content
+        zeta = zeta * max_v / tuning_param
+
+        return zeta

plato/processors/model_encrypt.py

@@ -0,0 +1,43 @@
+"""
+A processor that encrypts model weights in MaskCrypt.
+"""
+
+import logging
+from typing import Any
+
+import torch
+
+from plato.processors import model
+from plato.utils import homo_enc
+
+
+class Processor(model.Processor):
+    """
+    A processor that encrypts model weights with given encryption mask.
+    """
+
+    def __init__(self, mask=None, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.context = homo_enc.get_ckks_context()
+        self.mask = mask
+
+        para_nums = {}
+        extract_model = self.trainer.model.cpu().state_dict()
+        for key in extract_model.keys():
+            para_nums[key] = torch.numel(extract_model[key])
+        self.para_nums = para_nums
+
+    def process(self, data: Any) -> Any:
+        logging.info(
+            "[Client #%d] Encrypt the model weights with given encryption mask.",
+            self.client_id,
+        )
+
+        encrypted_weights = homo_enc.encrypt_weights(
+            data,
+            serialize=True,
+            context=self.context,
+            indices=self.mask,
+        )
+
+        return encrypted_weights

plato/processors/model_quantize.py

@@ -0,0 +1,18 @@
+"""
+Implements a Processor for quantizing model parameters.
+"""
+
+import torch
+
+from plato.processors import model
+
+
+class Processor(model.Processor):
+    """
+    Implements a Processor to quantize model parameters to 16-bit floating points.
+    """
+
+    def _process_layer(self, layer: torch.Tensor) -> torch.Tensor:
+        """Dequantizes each individual layer of the model."""
+
+        return layer.to(torch.bfloat16)