plato-learn 1.1__py3-none-any.whl

plato/models/general_multilayer.py

@@ -0,0 +1,254 @@
+"""
+A factory that generates a Multi-layer perceptron (MLP), with the ability to build a fully-
+connected network based on a specific configuration. This a very flexible MLP network generator to
+define any type of MLP network.
+
+Note: The general order of components in one MLP layer is:
+    Schema A: From the original paper of bn and dropout.
+    fc -> bn -> activation -> dropout -> ....
+
+    Schema B: From the researcher "https://math.stackexchange.com/users/167500/pseudomarvin".
+    fc -> activation -> dropout -> bn -> ....
+
+    See more discussion on the website:
+    https://stackoverflow.com/questions/39691902/ordering-of-batch-normalization-and-dropout
+
+Our work use the schema A.
+
+Trick: One may just drop the Dropout (when you have BN) as BN eliminates the need for Dropout in
+some cases, since intuitively BN provides similar regularization benefits as Dropout.
+
+"""
+
+from typing import Union, Dict, List
+from collections import OrderedDict
+
+from torch import nn
+
+activations_func = {
+    "relu": nn.ReLU,
+    "sigmoid": nn.Sigmoid,
+    "softmax": nn.Softmax,
+    "tanh": nn.Tanh,
+}
+
+
+# pylint: disable=too-many-locals
+def build_mlp_from_config(
+    mlp_configs: Dict[str, Union[int, List[Union[str, None, dict]]]],
+    layer_name_prefix: str = "layer",
+):
+    """
+    Build the fully-connected network (Multi-layer perceptron)
+    based on the input configuration.
+
+    :param  mlp_configs: A Dict type containing the hyper-parameters for definition.
+        It should contains:
+            input_dim: Integar
+            output_dim: Integar
+            hidden_layers_dim: List[int], with length N - 1
+            batch_norms: List[Union[None, dict]], with length N
+            activations: List[Union[None, str]], with length N
+            dropout_ratios: List[float], with length N
+    :param  layer_name_prefix: A string added to the layer's name.
+    """
+    input_dim = mlp_configs["input_dim"]
+    output_dim = mlp_configs["output_dim"]
+
+    hidden_layers_dim = mlp_configs["hidden_layers_dim"]
+    hidden_n = len(hidden_layers_dim)
+
+    batch_norms = mlp_configs["batch_norms"]
+    activations = mlp_configs["activations"]
+    dropout_porbs = (
+        mlp_configs["dropout_ratios"]
+        if isinstance(mlp_configs["dropout_ratios"], list)
+        else [mlp_configs["dropout_ratios"]]
+    )
+
+    assert len(batch_norms) == len(activations) == len(dropout_porbs)
+    assert hidden_n == len(batch_norms) - 1
+
+    def build_one_layer(
+        layer_ipt_dim, layer_opt_dim, batch_norm_param, activation, dropout_prob
+    ):
+        """Build one layer of MLP. Default no hidden layer.
+
+        For the structure of one MLP layer. Please access the description
+        in the NOTE part.
+
+        """
+        layer_structure = OrderedDict()
+        layer_structure["fc"] = nn.Linear(layer_ipt_dim, layer_opt_dim)
+
+        if batch_norm_param is not None:
+            layer_structure["bn"] = nn.BatchNorm1d(layer_opt_dim, **batch_norm_param)
+        if activation is not None:
+            layer_structure[activation] = activations_func[activation]()
+        if dropout_prob != 0.0:
+            layer_structure["drop"] = nn.Dropout(p=dropout_prob)
+
+        return nn.Sequential(layer_structure)
+
+    mlp_layers = OrderedDict()
+
+    # add the final output layer to the hidden layer for building layers
+    hidden_layers_dim.append(output_dim)
+    for hid_id, hid_dim in enumerate(hidden_layers_dim):
+        layer_input_dim = input_dim if hid_id == 0 else hidden_layers_dim[hid_id - 1]
+        desired_batch_norm = batch_norms[hid_id]
+        activation = activations[hid_id]
+        dropout_prob = dropout_porbs[hid_id]
+        built_layer = build_one_layer(
+            layer_input_dim, hid_dim, desired_batch_norm, activation, dropout_prob
+        )
+        mlp_layers[layer_name_prefix + str(hid_id + 1)] = built_layer
+
+    return nn.Sequential(mlp_layers)
+
+
+class Model:
+    """
+    The Multi-layer perceptron (MLP) model.
+
+    The implemented mlp networks are:
+    - linear_mlp, The mlp with one hidden layer.
+    - simclr_projection_mlp, The projection layer of SimCLR method.
+    - simsiam_projection_mlp, The projection layer of SimSiam method.
+    - simsiam_prediction_mlp, The prediction layer of SimSiam method.
+    - byol_projection_mlp, The projection layer of BYOL method.
+    - byol_prediction_mlp, The prediction layer of BYOL method.
+    - moco_final_mlp, The final layer of MoCo method.
+    - plato_multilayer, The Plato's multilayer.
+    - customized_mlp, The customized layer.
+    """
+
+    # pylint: disable=too-few-public-methods
+    @staticmethod
+    def get(
+        model_name: str,
+        input_dim: int,
+        output_dim: int,
+        **kwargs: Dict[str, Union[int, List[Union[str, None, dict]]]],
+    ):
+        # pylint:disable=too-many-return-statements
+        """Get the desired MLP model with required hyper-parameters (input_dim)."""
+
+        if model_name == "linear_mlp":
+            return build_mlp_from_config(
+                mlp_configs=dict(
+                    output_dim=output_dim,
+                    input_dim=input_dim,
+                    hidden_layers_dim=[],
+                    batch_norms=[None],
+                    activations=[None],
+                    dropout_ratios=[0.0],
+                )
+            )
+
+        if model_name == "simclr_projection_mlp":
+            projection_hidden_dim = kwargs["projection_hidden_dim"]
+            return build_mlp_from_config(
+                dict(
+                    output_dim=output_dim,
+                    input_dim=input_dim,
+                    hidden_layers_dim=[projection_hidden_dim],
+                    batch_norms=[None, None],
+                    activations=["relu", None],
+                    dropout_ratios=[0.0, 0.0],
+                )
+            )
+
+        if model_name == "simsiam_projection_mlp":
+            projection_hidden_dim = kwargs["projection_hidden_dim"]
+            return build_mlp_from_config(
+                dict(
+                    output_dim=output_dim,
+                    input_dim=input_dim,
+                    hidden_layers_dim=[
+                        projection_hidden_dim,
+                        projection_hidden_dim,
+                    ],
+                    batch_norms=[
+                        dict(momentum=0.1, eps=1e-5),
+                        dict(momentum=0.1, eps=1e-5),
+                        dict(momentum=0.1, eps=1e-5),
+                    ],
+                    activations=["relu", "relu", None],
+                    dropout_ratios=[0.0, 0.0, 0.0],
+                )
+            )
+
+        if model_name == "simsiam_prediction_mlp":
+            prediction_hidden_dim = kwargs["prediction_hidden_dim"]
+            return build_mlp_from_config(
+                dict(
+                    output_dim=output_dim,
+                    input_dim=input_dim,
+                    hidden_layers_dim=[prediction_hidden_dim],
+                    batch_norms=[dict(momentum=0.1, eps=1e-5), None],
+                    activations=["relu", None],
+                    dropout_ratios=[0.0, 0.0],
+                )
+            )
+
+        if model_name == "byol_projection_mlp":
+            projection_hidden_dim = kwargs["projection_hidden_dim"]
+            return build_mlp_from_config(
+                dict(
+                    output_dim=output_dim,
+                    input_dim=input_dim,
+                    hidden_layers_dim=[projection_hidden_dim],
+                    batch_norms=[dict(momentum=0.1, eps=1e-5), None],
+                    activations=["relu", None],
+                    dropout_ratios=[0.0, 0.0],
+                )
+            )
+
+        if model_name == "byol_prediction_mlp":
+            prediction_hidden_dim = kwargs["prediction_hidden_dim"]
+            return build_mlp_from_config(
+                dict(
+                    output_dim=output_dim,
+                    input_dim=input_dim,
+                    hidden_layers_dim=[prediction_hidden_dim],
+                    batch_norms=[dict(momentum=0.1, eps=1e-5), None],
+                    activations=["relu", None],
+                    dropout_ratios=[0.0, 0.0],
+                )
+            )
+
+        if model_name == "moco_final_mlp":
+            projection_hidden_dim = kwargs["projection_hidden_dim"]
+            return build_mlp_from_config(
+                dict(
+                    output_dim=output_dim,
+                    input_dim=input_dim,
+                    hidden_layers_dim=[projection_hidden_dim],
+                    batch_norms=[None, None],
+                    activations=["relu", None],
+                    dropout_ratios=[0.0, 0.0],
+                )
+            )
+
+        if model_name == "plato_multilayer":
+            return build_mlp_from_config(
+                dict(
+                    output_dim=output_dim,
+                    input_dim=input_dim,
+                    hidden_layers_dim=[1024, 512, 256, 128],
+                    batch_norms=[None, None, None, None, None],
+                    activations=["tanh", "tanh", "tanh", "tanh", None],
+                    dropout_ratios=[0.0, 0.0, 0.0, 0.0, 0.0],
+                )
+            )
+
+        # obtain the customized mlp laye
+        # the user needs to put the corresponding hyper-parameters
+        # in the 'kwargs'
+        if model_name == "customized_mlp":
+            return build_mlp_from_config(
+                dict(output_dim=output_dim, input_dim=input_dim, **kwargs)
+            )
+
+        raise ValueError(f"No such MLP model: {model_name}")

plato/models/huggingface.py

@@ -0,0 +1,27 @@
+"""
+Obtaining a model from the PyTorch Hub.
+"""
+
+from transformers import AutoModelForCausalLM, AutoConfig
+from plato.config import Config
+
+
+class Model:
+    """The CausalLM model loaded from HuggingFace."""
+
+    @staticmethod
+    def get(model_name=None, **kwargs):  # pylint: disable=unused-argument
+        """Returns a named model from HuggingFace."""
+        config_kwargs = {
+            "cache_dir": None,
+            "revision": "main",
+            "use_auth_token": None,
+        }
+
+        config = AutoConfig.from_pretrained(model_name, **config_kwargs)
+
+        return AutoModelForCausalLM.from_pretrained(
+            model_name,
+            config=config,
+            cache_dir=Config().params["model_path"] + "/huggingface",
+        )

plato/models/lenet5.py

@@ -0,0 +1,113 @@
+"""The LeNet-5 model for PyTorch.
+
+Reference:
+
+Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner. "Gradient-based learning applied to
+document recognition." Proceedings of the IEEE, November 1998.
+"""
+
+import collections
+
+from torch import nn
+import torch.nn.functional as F
+
+
+class Model(nn.Module):
+    """The LeNet-5 model.
+
+    :param num_classes: The number of classes. The default value is 10.
+    """
+
+    def __init__(self, num_classes: int = 10, cut_layer=None):
+        super().__init__()
+        self.cut_layer = cut_layer
+
+        # We pad the image to get an input size of 32x32 as for the
+        # original network in the LeCun paper
+        self.conv1 = nn.Conv2d(
+            in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=2, bias=True
+        )
+        self.relu1 = nn.ReLU()
+        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.conv2 = nn.Conv2d(
+            in_channels=6,
+            out_channels=16,
+            kernel_size=5,
+            stride=1,
+            padding=0,
+            bias=True,
+        )
+        self.relu2 = nn.ReLU()
+        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.conv3 = nn.Conv2d(
+            in_channels=16, out_channels=120, kernel_size=5, bias=True
+        )
+        self.relu3 = nn.ReLU()
+        self.fc4 = nn.Linear(120, 84)
+        self.relu4 = nn.ReLU()
+        self.fc5 = nn.Linear(84, 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["relu1"] = self.relu1
+        self.layerdict["pool1"] = self.pool1
+        self.layerdict["conv2"] = self.conv2
+        self.layerdict["relu2"] = self.relu2
+        self.layerdict["pool2"] = self.pool2
+        self.layerdict["conv3"] = self.conv3
+        self.layerdict["relu3"] = self.relu3
+        self.layerdict["flatten"] = self.flatten
+        self.layerdict["fc4"] = self.fc4
+        self.layerdict["relu4"] = self.relu4
+        self.layerdict["fc5"] = self.fc5
+        self.layers.append("conv1")
+        self.layers.append("relu1")
+        self.layers.append("pool1")
+        self.layers.append("conv2")
+        self.layers.append("relu2")
+        self.layers.append("pool2")
+        self.layers.append("conv3")
+        self.layers.append("relu3")
+        self.layers.append("flatten")
+        self.layers.append("fc4")
+        self.layers.append("relu4")
+        self.layers.append("fc5")
+
+    def flatten(self, x):
+        """Flatten the tensor."""
+        return x.view(x.size(0), -1)
+
+    def forward(self, x):
+        """Forward pass."""
+        if self.cut_layer is not None and self.training:
+            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)
+        else:
+            x = self.conv1(x)
+            x = self.relu1(x)
+            x = self.pool1(x)
+            x = self.conv2(x)
+            x = self.relu2(x)
+            x = self.pool2(x)
+            x = self.conv3(x)
+            x = self.relu3(x)
+            x = self.flatten(x)
+            x = self.fc4(x)
+            x = self.relu4(x)
+            x = self.fc5(x)
+
+        return F.log_softmax(x, dim=1)
+
+    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

plato/models/multilayer.py

@@ -0,0 +1,90 @@
+"""
+The Multi-Layer Perception model for PyTorch.
+The model follows the previous work to use tanh as activation
+Reference: https://www.comp.nus.edu.sg/~reza/files/Shokri-SP2019.pdf
+"""
+
+import collections
+
+import torch.nn as nn
+
+from plato.config import Config
+
+
+class Model(nn.Module):
+    """The Multi-Layer Perception model.
+
+    Arguments:
+        num_classes (int): The number of classes. Default: 10.
+    """
+
+    def __init__(self, input_dim=600, num_classes=10):
+        super().__init__()
+        self.fc1 = nn.Sequential(nn.Linear(input_dim, 1024), nn.Tanh())
+
+        self.fc2 = nn.Sequential(nn.Linear(1024, 512), nn.Tanh())
+
+        self.fc3 = nn.Sequential(
+            nn.Linear(512, 256),
+            nn.Tanh(),
+        )
+
+        self.fc4 = nn.Sequential(
+            nn.Linear(256, 128),
+            nn.Tanh(),
+        )
+
+        self.fc5 = nn.Linear(128, 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["fc1"] = self.fc1
+        self.layerdict["fc2"] = self.fc2
+        self.layerdict["fc3"] = self.fc3
+        self.layerdict["fc4"] = self.fc4
+        self.layerdict["fc5"] = self.fc5
+
+        self.layers.append("fc1")
+        self.layers.append("fc2")
+        self.layers.append("fc3")
+        self.layers.append("fc4")
+        self.layers.append("fc5")
+
+    def forward(self, x):
+        """Forward pass."""
+        x = self.fc1(x)
+        x = self.fc2(x)
+        x = self.fc3(x)
+        x = self.fc4(x)
+        x = self.fc5(x)
+        return x
+
+    def forward_to(self, x, cut_layer):
+        """Forward pass, but only to the layer specified by cut_layer."""
+        layer_index = self.layers.index(cut_layer)
+
+        for i in range(0, layer_index + 1):
+            x = self.layerdict[self.layers[i]](x)
+
+        return x
+
+    def forward_from(self, x, cut_layer):
+        """Forward pass, starting from the layer specified by cut_layer."""
+        layer_index = self.layers.index(cut_layer)
+
+        for i in range(layer_index + 1, len(self.layers)):
+            x = self.layerdict[self.layers[i]](x)
+
+        return x
+
+    @staticmethod
+    def get_model(*args):
+        """Obtaining an instance of this model."""
+        if hasattr(Config().trainer, "num_classes"):
+            return Model(
+                input_dim=Config().trainer.input_dim,
+                num_classes=Config().trainer.num_classes,
+            )
+        return Model()

plato/models/multimodal/base_net.py

@@ -0,0 +1,91 @@
+"""
+The base class used for all following classes
+
+"""
+
+import torch
+import torch.nn as nn
+
+from mmaction.models import build_model
+
+
+class BaseClassificationNet(nn.Module):
+    """Base class for classification networks"""
+
+    def __init__(self, net_configs, is_head_included=True):
+        super(BaseClassificationNet, self).__init__()
+
+        self.net_configs = net_configs
+        # 1 build the model based on the configurations
+        self._net = build_model(net_configs)
+
+        self.is_head_included = is_head_included
+
+        # the features must be forwarded the avg pool
+        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
+
+    def get_net(self):
+        """Get the built network"""
+        return self._net
+
+    def forward_train(self, ipt_data, labels, **kwargs):
+        """Defines the computation performed at every call when training."""
+
+        ipt_data_sz = ipt_data.reshape((-1,) + ipt_data.shape[2:])
+
+        # 1. forward the backbone
+        data_feat = self._net.extract_feat(ipt_data_sz)
+        # from [N * num_segs, in_channels, h, w]
+        #   to [N, in_channels, 1, 1]
+        immediate_feat = self.avg_pool(data_feat)
+        # to [N, in_channels]
+        immediate_feat = torch.squeeze(immediate_feat)
+
+        # 2. forward the classification head if possible and obtain the losses
+        loss_cls = 0.0
+        if self.is_head_included:
+            cls_score = self._net.cls_head(data_feat)
+
+            gt_labels = labels.squeeze()
+            loss_cls = self._net.cls_head.loss(cls_score, gt_labels, **kwargs)
+
+            return [immediate_feat, cls_score, loss_cls]
+
+        return [immediate_feat]
+
+    def forward_test(self, ipt_data, **kwargs):
+        """Defines the computation performed at every call when training."""
+
+        ipt_data = ipt_data.reshape((-1,) + ipt_data.shape[2:])
+        # 1. forward the backbone
+        data_feat = self._net.extract_feat(ipt_data)
+        # 2. forward the classification head if possible and obtain the losses
+        cls_score = 0.0
+        if self.is_head_included:
+            cls_score = self._net.cls_head(data_feat)
+
+            return [data_feat, cls_score]
+
+        return [data_feat]
+
+    def forward(self, ipt_data, label=None, return_loss=True, **kwargs):
+        """Defines the computation performed at every call.
+
+        Args:
+            ipt_data (torch.Tensor): The input data.
+                the size of x is (num_batches, channel, num_slices, h, w).
+        Returns:
+            torch.Tensor: The feature of the input
+            samples extracted by the backbone.
+        """
+
+        if return_loss:
+            if label is None:
+                raise ValueError("Label should not be None.")
+            if self._net.blending is not None:
+                blended_ipt_data, label = self._net.blending(ipt_data, label)
+            else:
+                blended_ipt_data = ipt_data
+            return self.forward_train(blended_ipt_data, label, **kwargs)
+
+        return self.forward_test(ipt_data, **kwargs)