plato-learn 1.1__py3-none-any.whl

plato/models/resnet.py

@@ -0,0 +1,190 @@
+"""
+The ResNet model (for the CIFAR-10 dataset only).
+
+Reference:
+
+https://github.com/kuangliu/pytorch-cifar/blob/master/models/resnet.py
+"""
+
+import collections
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1):
+        super().__init__()
+        self.conv1 = nn.Conv2d(
+            in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False
+        )
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(
+            planes, planes, kernel_size=3, stride=1, padding=1, bias=False
+        )
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(
+                    in_planes,
+                    self.expansion * planes,
+                    kernel_size=1,
+                    stride=stride,
+                    bias=False,
+                ),
+                nn.BatchNorm2d(self.expansion * planes),
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, in_planes, planes, stride=1):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(
+            planes, planes, kernel_size=3, stride=stride, padding=1, bias=False
+        )
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(
+            planes, self.expansion * planes, kernel_size=1, bias=False
+        )
+        self.bn3 = nn.BatchNorm2d(self.expansion * planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(
+                    in_planes,
+                    self.expansion * planes,
+                    kernel_size=1,
+                    stride=stride,
+                    bias=False,
+                ),
+                nn.BatchNorm2d(self.expansion * planes),
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = F.relu(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class Model(nn.Module):
+    def __init__(self, block, num_blocks, num_classes=10, cut_layer=None):
+        super().__init__()
+
+        self.in_planes = 64
+
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
+        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
+        self.linear = nn.Linear(512 * block.expansion, num_classes)
+
+        # Preparing named layers so that the model can be split and straddle
+        # across the client and the server
+        self.layers = []
+        self.layerdict = collections.OrderedDict()
+        self.layerdict["conv1"] = self.conv1
+        self.layerdict["bn1"] = self.bn1
+        self.layerdict["relu"] = F.relu
+        self.layerdict["layer1"] = self.layer1
+        self.layerdict["layer2"] = self.layer2
+        self.layerdict["layer3"] = self.layer3
+        self.layerdict["layer4"] = self.layer4
+        self.layers.append("conv1")
+        self.layers.append("bn1")
+        self.layers.append("relu")
+        self.layers.append("layer1")
+        self.layers.append("layer2")
+        self.layers.append("layer3")
+        self.layers.append("layer4")
+        self.cut_layer = cut_layer
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        strides = [stride] + [1] * (num_blocks - 1)
+        layers = []
+        for stride in strides:
+            layers.append(block(self.in_planes, planes, stride))
+            self.in_planes = planes * block.expansion
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = self.layer4(out)
+        out = F.avg_pool2d(out, 4)
+        out = out.view(out.size(0), -1)
+        out = self.linear(out)
+        return out
+
+    def forward_to(self, x):
+        """Forward pass, but only to the layer specified by cut_layer."""
+        layer_index = self.layers.index(self.cut_layer)
+
+        for i in range(0, layer_index + 1):
+            x = self.layerdict[self.layers[i]](x)
+        return x
+
+    def forward_from(self, x):
+        """Forward pass, starting from the layer specified by cut_layer."""
+        layer_index = self.layers.index(self.cut_layer)
+        for i in range(layer_index + 1, len(self.layers)):
+            x = self.layerdict[self.layers[i]](x)
+
+        out = F.avg_pool2d(x, 4)
+        out = out.view(out.size(0), -1)
+        out = self.linear(out)
+        return out
+
+    @staticmethod
+    def is_valid_model_type(model_type):
+        return (
+            model_type.startswith("resnet_")
+            and len(model_type.split("_")) == 2
+            and int(model_type.split("_")[1]) in [18, 34, 50, 101, 152]
+        )
+
+    @staticmethod
+    def get(model_name=None, num_classes=None, cut_layer=None, **kwargs):
+        """Returns a suitable ResNet model according to its type."""
+        if not Model.is_valid_model_type(model_name):
+            raise ValueError(f"Invalid Resnet model name: {model_name}")
+
+        resnet_type = int(model_name.split("_")[1])
+
+        if num_classes is None:
+            num_classes = 10
+
+        if resnet_type == 18:
+            return Model(BasicBlock, [2, 2, 2, 2], num_classes, cut_layer)
+        elif resnet_type == 34:
+            return Model(BasicBlock, [3, 4, 6, 3], num_classes, cut_layer)
+        elif resnet_type == 50:
+            return Model(Bottleneck, [3, 4, 6, 3], num_classes, cut_layer)
+        elif resnet_type == 101:
+            return Model(Bottleneck, [3, 4, 23, 3], num_classes, cut_layer)
+        elif resnet_type == 152:
+            return Model(Bottleneck, [3, 8, 36, 3], num_classes, cut_layer)
+
+        return None

plato/models/torch_hub.py

@@ -0,0 +1,19 @@
+"""
+Obtaining a model from the PyTorch Hub.
+"""
+
+import torch
+
+
+class Model:
+    """
+    The model loaded from PyTorch Hub.
+
+    We will soon be using the get_model() method for torchvision 0.14 when it is released.
+    """
+
+    @staticmethod
+    # pylint: disable=unused-argument
+    def get(model_name=None, **kwargs):
+        """Returns a named model from PyTorch Hub."""
+        return torch.hub.load("pytorch/vision", model_name, **kwargs)

plato/models/vgg.py

@@ -0,0 +1,113 @@
+"""
+A VGG-style neural network model for image classification.
+"""
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Model(nn.Module):
+    """A VGG-style neural network model for image classification."""
+
+    class ConvModule(nn.Module):
+        """A single convolutional module in a VGG network."""
+
+        def __init__(self, in_filters, out_filters):
+            super().__init__()
+            self.conv = nn.Conv2d(in_filters, out_filters, kernel_size=3, padding=1)
+            self.bn = nn.BatchNorm2d(out_filters)
+
+        def forward(self, x):
+            return F.relu(self.bn(self.conv(x)))
+
+    def __init__(self, plan, outputs=10):
+        super().__init__()
+
+        layers = []
+        filters = 3
+
+        for spec in plan:
+            if spec == "M":
+                layers.append(nn.MaxPool2d(kernel_size=2, stride=2))
+            else:
+                layers.append(Model.ConvModule(filters, spec))
+                filters = spec
+
+        self.layers = nn.Sequential(*layers)
+        self.fc = nn.Linear(512, outputs)
+
+    def forward(self, x):
+        x = self.layers(x)
+        x = nn.AvgPool2d(2)(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        return x
+
+    @staticmethod
+    def is_valid_model_name(model_name):
+        return (
+            model_name.startswith("vgg")
+            and len(model_name.split("_")) == 2
+            and model_name.split("_")[1].isdigit()
+            and int(model_name.split("_")[1]) in [11, 13, 16, 19]
+        )
+
+    @staticmethod
+    def get(model_name, num_classes=10):
+        """Returns a suitable VGG model corresponding to the provided name."""
+        if not Model.is_valid_model_name(model_name):
+            raise ValueError(f"Invalid VGG model name: {model_name}")
+
+        outputs = num_classes
+
+        num = int(model_name.split("_")[1])
+
+        if num == 11:
+            plan = [64, "M", 128, "M", 256, 256, "M", 512, 512, "M", 512, 512]
+        elif num == 13:
+            plan = [64, 64, "M", 128, 128, "M", 256, 256, "M", 512, 512, "M", 512, 512]
+        elif num == 16:
+            plan = [
+                64,
+                64,
+                "M",
+                128,
+                128,
+                "M",
+                256,
+                256,
+                256,
+                "M",
+                512,
+                512,
+                512,
+                "M",
+                512,
+                512,
+                512,
+            ]
+        elif num == 19:
+            plan = [
+                64,
+                64,
+                "M",
+                128,
+                128,
+                "M",
+                256,
+                256,
+                256,
+                256,
+                "M",
+                512,
+                512,
+                512,
+                512,
+                "M",
+                512,
+                512,
+                512,
+                512,
+            ]
+
+        return Model(plan, outputs)

plato/models/vit.py

@@ -0,0 +1,166 @@
+"""
+Obtaining a Vision Transformer (ViT) model for image classification from HuggingFace.
+
+Reference to the Tokens-to-Token ViT (T2T-ViT) model:
+https://github.com/yitu-opensource/T2T-ViT
+
+Reference to the Deep Vision Transformer (DeepViT) model:
+https://github.com/zhoudaquan/dvit_repo
+
+"""
+
+import torch
+from torch import nn
+from transformers import AutoConfig, AutoModelForImageClassification
+
+
+from plato.config import Config
+
+
+class ResolutionAdjustedModel(nn.Module):
+    """
+    Transforms the image resolution to the assigned resolution of a pretrained model.
+    """
+
+    def __init__(self, model_name, config) -> nn.Module:
+        super().__init__()
+
+        if (
+            hasattr(Config().parameters, "model")
+            and hasattr(Config().parameters.model, "pretrained")
+            and not Config().parameters.model.pretrained
+        ):
+            ignore_mismatched_sizes = True
+        else:
+            ignore_mismatched_sizes = False
+
+        self.model = AutoModelForImageClassification.from_pretrained(
+            model_name,
+            config=config,
+            ignore_mismatched_sizes=ignore_mismatched_sizes,
+            cache_dir=Config().params["model_path"] + "/huggingface",
+        )
+
+        if (
+            hasattr(Config().parameters, "model")
+            and hasattr(Config().parameters.model, "pretrained")
+            and not Config().parameters.model.pretrained
+        ):
+            self.model.init_weights()
+        self.resolution = config.image_size
+
+    def forward(self, image):
+        """
+        Adjusts the image resolution and outputs the logits.
+        """
+        if image.size(-1) != self.resolution:
+            image = nn.functional.interpolate(
+                image, size=self.resolution, mode="bicubic"
+            )
+        outputs = self.model(image)
+        return outputs.logits
+
+
+class T2TVIT(nn.Module):
+    """Wrapper for the T2T-ViT model."""
+
+    def __init__(self, name) -> nn.Module:
+        super().__init__()
+        # pylint:disable=import-outside-toplevel
+        from plato.models import t2tvit
+        from plato.models.t2tvit.models import t2t_vit
+
+        model_name = getattr(t2t_vit, name)
+        t2t = model_name(num_classes=Config().trainer.num_classes)
+
+        if (
+            hasattr(Config().parameters, "model")
+            and hasattr(Config().parameters.model, "pretrained")
+            and Config().parameters.model.pretrained
+        ):
+            t2tvit.utils.load_for_transfer_learning(
+                t2t,
+                Config().parameters.model.pretrain_path,
+                use_ema=True,
+                strict=False,
+                num_classes=Config().trainer.num_classes,
+            )
+        self.model = t2t
+        self.resolution = 224
+
+    def forward(self, feature):
+        """The forward pass."""
+        if feature.size(-1) != self.resolution:
+            feature = nn.functional.interpolate(
+                feature, size=self.resolution, mode="bicubic"
+            )
+        return self.model(feature)
+
+
+class DeepViT(nn.Module):
+    """Wrapper for the DeepViT model."""
+
+    def __init__(self, name) -> nn.Module:
+        super().__init__()
+        # pylint:disable=import-outside-toplevel
+        from plato.models.dvit.models import deep_vision_transformer
+
+        model_name = getattr(deep_vision_transformer, name)
+        deepvit = model_name(
+            pretrained=False, num_classes=Config.trainer.num_classes, in_chans=3
+        )
+        if (
+            hasattr(Config().parameters, "model")
+            and hasattr(Config().parameters.model, "pretrained")
+            and Config().parameters.model.pretrained
+        ):
+            state_dict = torch.load(
+                Config().parameters.model.pretrain_path, map_location="cpu"
+            )
+            del state_dict["head.weight"]
+            del state_dict["head.bias"]
+            deepvit.load_state_dict(state_dict)
+        self.model = deepvit
+        self.resolution = 224
+
+    def forward(self, feature):
+        """The forward pass."""
+        if feature.size(-1) != self.resolution:
+            feature = nn.functional.interpolate(
+                feature, size=self.resolution, mode="bicubic"
+            )
+
+        return self.model(feature)
+
+
+class Model:
+    """
+    The Transformer and other models loaded from HuggingFace.
+    Supported by HuggingFace
+    https://huggingface.co/docs/transformers/model_doc/auto#transformers.AutoModel
+    """
+
+    # pylint:disable=too-few-public-methods
+    @staticmethod
+    def get(model_name=None, **kwargs):  # pylint: disable=unused-argument
+        """Returns a named model from HuggingFace."""
+        if "t2t" in model_name:
+            return T2TVIT(model_name)
+
+        if "deepvit" in model_name:
+            return DeepViT(model_name)
+
+        config_kwargs = {
+            "cache_dir": None,
+            "revision": "main",
+            "use_auth_token": None,
+        }
+        config_kwargs.update(kwargs)
+
+        model_name = model_name.replace("@", "/")
+        # Only prepend "google/" if the model name doesn't already contain a namespace
+        if "/" not in model_name:
+            model_name = "google/" + model_name
+        config = AutoConfig.from_pretrained(model_name, **config_kwargs)
+
+        return ResolutionAdjustedModel(model_name, config)

plato/models/yolov8.py

@@ -0,0 +1,22 @@
+"""
+Obtaining a model from the Ultralytics.
+"""
+
+from ultralytics import YOLO
+
+from plato.config import Config
+
+
+class Model:
+    """
+    The model loaded from the YOLOv8.
+
+    """
+
+    @staticmethod
+    # pylint: disable=unused-argument
+    def get(model_name=None, **kwargs):
+        """Returns the YOLOV8 model loaded from the Ultralytics."""
+        model_type = Config().parameters.model.type
+
+        return YOLO(model_type)

plato/processors/base.py

@@ -0,0 +1,35 @@
+"""
+The Processor class is designed for pre-processing data payloads before or after they
+are transmitted over the network between the clients and the servers.
+"""
+
+from abc import abstractmethod
+from collections.abc import Iterable
+from typing import Any
+
+
+class Processor:
+    """
+    The base Processor class does nothing on the data payload.
+    """
+
+    def __init__(self, name=None, trainer=None, **kwargs) -> None:
+        """Constructor for Processor."""
+        self.name = name
+        self.trainer = trainer
+
+    @abstractmethod
+    def process(self, data: Any) -> Any:
+        """
+        Processing a data payload.
+        """
+        return data
+
+    def process_iterable(self, data: Iterable) -> Iterable:
+        """
+        Processing an Iterable of data payload.
+        """
+        return map(self.process, data)
+
+    def __repr__(self) -> str:
+        return self.name

plato/processors/compress.py

@@ -0,0 +1,46 @@
+"""
+Implements a Processor for compressing a numpy array.
+"""
+
+from typing import Any
+
+import zstd
+
+from plato.processors import base
+
+
+class Processor(base.Processor):
+    """Implements a Processor for compressing numpy array."""
+
+    def __init__(self, cr=1, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.compression_ratio = cr
+
+    def process(self, data: Any) -> Any:
+        """Implements a Processor for compressing numpy array."""
+        if isinstance(data, list):
+            ret = []
+            datashape_feature = data[0][0].shape
+            datatype_feature = data[0][0].dtype
+            ret.append((datashape_feature, datatype_feature))
+            for logits, targets in data:
+                datashape_target = targets.shape
+                datatype_target = targets.dtype
+                datacom_feature = zstd.compress(logits, self.compression_ratio)
+                datacom_target = zstd.compress(targets, self.compression_ratio)
+                ret.append(
+                    (
+                        datacom_feature,
+                        datacom_target,
+                        datashape_target,
+                        datatype_target,
+                    )
+                )
+        else:
+            ret = (
+                data.shape,
+                data.dtype,
+                zstd.compress(data, self.compression_ratio),
+            )
+
+        return ret

plato/processors/decompress.py

@@ -0,0 +1,48 @@
+"""
+Implements a Processor for decompressing a numpy array.
+"""
+
+from typing import Any
+
+import numpy as np
+import zstd
+
+from plato.processors import base
+
+
+class Processor(base.Processor):
+    """Implements a Processor for decompressing a numpy array."""
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+    def process(self, data: Any) -> Any:
+        """Implements a Processor for decompressing a numpy array."""
+        if isinstance(data, list):
+            ret = []
+            datashape_feature = data[0][0]
+            datatype_feature = data[0][1]
+            for (
+                datacom_feature,
+                datacom_target,
+                datashape_target,
+                datatype_target,
+            ) in data[1:]:
+                datacom_feature = zstd.decompress(datacom_feature)
+                datacom_feature = np.frombuffer(
+                    datacom_feature, datatype_feature
+                ).reshape(datashape_feature)
+                if len(datashape_target) > 0 and datashape_target[0] == 0:
+                    datacom_target = np.zeros(datashape_target)
+                else:
+                    datacom_target = zstd.decompress(datacom_target)
+                    datacom_target = np.frombuffer(
+                        datacom_target, datatype_target
+                    ).reshape(datashape_target)
+                ret.append((datacom_feature, datacom_target))
+        else:
+            shape, dtype, modelcom = data
+            modelcom = zstd.decompress(modelcom)
+            ret = np.frombuffer(modelcom, dtype).reshape(shape)
+
+        return ret

plato/processors/feature.py

@@ -0,0 +1,51 @@
+"""
+Implements a generalized Processor for applying operations onto MistNet PyTorch features.
+"""
+
+from typing import Any, Callable
+
+import torch
+
+from plato.processors import base
+
+
+class Processor(base.Processor):
+    """
+    Implements a generalized Processor for applying operations onto MistNet PyTorch features.
+    """
+
+    def __init__(
+        self,
+        method: Callable = lambda x, y: (x, y),
+        client_id=None,
+        use_numpy=True,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+
+        self.client_id = client_id
+        self.method = method
+        self.use_numpy = use_numpy
+
+    def process(self, data: Any) -> Any:
+        """
+        Implements a generalized Processor for applying operations onto MistNet PyTorch features.
+        """
+
+        output = []
+
+        for logits, targets in data:
+            if self.use_numpy:
+                logits = logits.detach().numpy()
+
+            logits, targets = self.method(logits, targets)
+
+            if self.use_numpy:
+                if self.trainer.device != "cpu":
+                    logits = torch.from_numpy(logits.astype("float16"))
+                else:
+                    logits = torch.from_numpy(logits.astype("float32"))
+
+            output.append((logits, targets))
+
+        return output