plato-learn 1.1__py3-none-any.whl

plato/models/multimodal/blending.py

@@ -0,0 +1,142 @@
+"""
+
+the overfitting value is the gap between the train loss L_i^T and
+the groundtruth L_i^* w.r.t the hypothetical target distribution
+Note: the  L^* is approximated by the validation loss L^V
+
+"""
+
+import numpy as np
+
+
+def compute_overfitting_o(eval_avg_loss, train_avg_loss):
+    """We define overfitting at epoch N as the gap between LTN and
+    L∗N (approximated by ON in fig. 3)."""
+    return eval_avg_loss - train_avg_loss
+
+
+def compute_delta_overfitting_o(
+    n_eval_avg_loss, n_train_avg_loss, N_eval_avg_loss, N_train_avg_loss
+):
+    """Compute the overfitting O based on losses between step n and N, (n < N)"""
+    delta_O = compute_overfitting_o(
+        n_eval_avg_loss, n_train_avg_loss
+    ) - compute_overfitting_o(N_eval_avg_loss, N_train_avg_loss)
+    return delta_O
+
+
+def compute_generalization_g(eval_avg_loss):
+    """Compute the generalization g which is actually the evluation loss"""
+    return eval_avg_loss
+
+
+def compute_delta_generalization(eval_avg_loss_n, eval_avg_loss_N):
+    """Compute the difference of the generalization"""
+    return compute_generalization_g(eval_avg_loss_n) - compute_generalization_g(
+        eval_avg_loss_N
+    )
+
+
+# n < N,
+def OGR_n2N(n_eval_avg_loss, n_train_avg_loss, N_eval_avg_loss, N_train_avg_loss):
+    """Compute the OGR = abs(delta_O/delta_G)"""
+    delta_O = compute_delta_overfitting_o(
+        n_eval_avg_loss, n_train_avg_loss, N_eval_avg_loss, N_train_avg_loss
+    )
+    delta_G = compute_delta_generalization(n_eval_avg_loss, N_eval_avg_loss)
+
+    ogr = abs(delta_O / delta_G)
+    return ogr
+
+
+# Optimal Gradient Blend
+#   x << N
+def get_optimal_gradient_blend_weights(modalities_losses_n, modalities_losses_N):
+    """Get the weights of modaliteis for optimal gradient blending
+
+    Args:
+        modalities_losses_n (dict): contains the train/eval losses for each modality in epoch n
+        modalities_losses_N (dict): contains the train/eval losses for each modality in epoch N
+
+        The structure of the above two dicts should be: (for example)
+            {"train": {"RGB": float, "Flow": float},
+            "eval": {"RGB": float, "Flow": float}}
+
+    The equation:
+        w^i = <∇L^*, v_i>/(σ_i)^2 * 1/Z
+            = <∇L^*, v_i>/(σ_i)^2 * 1/(sum_i <∇L^*, v_i>/2*(σ_i)^2)
+            = G^i / (O^i)^2 * 1 / (sum_i G^i / (2 * (O^i)^2))
+
+        where G^i = G_N,n =  L^*_n − L^*_N = compute_delta_generalization,
+                O^i = O_N,n =  O_N - O_n = compute_delta_overfitting_O
+    """
+    modality_names = list(modalities_losses_n["train"].keys())
+
+    Z = 0
+    modls_GO = dict()
+    for modality_nm in modality_names:
+        modl_eval_avg_loss_n = modalities_losses_n["eval"][modality_nm]
+        modl_subtrain_avg_loss_n = modalities_losses_n["train"][modality_nm]
+        modl_eval_avg_loss_N = modalities_losses_N["eval"][modality_nm]
+        modl_subtrain_avg_loss_N = modalities_losses_N["train"][modality_nm]
+        G_i = compute_delta_generalization(
+            eval_avg_loss_n=modl_eval_avg_loss_n, eval_avg_loss_N=modl_eval_avg_loss_N
+        )
+        O_i = compute_delta_overfitting_o(
+            n_eval_avg_loss=modl_eval_avg_loss_n,
+            n_train_avg_loss=modl_subtrain_avg_loss_n,
+            N_eval_avg_loss=modl_eval_avg_loss_N,
+            N_train_avg_loss=modl_subtrain_avg_loss_N,
+        )
+
+        modls_GO[modality_nm] = G_i / (O_i * O_i)
+        Gi_div_sqr_Oi = G_i / (2 * O_i * O_i)
+
+        Z += Gi_div_sqr_Oi
+
+    optimal_weights = dict()
+    for modality_nm in modality_names:
+        optimal_weights[modality_nm] = modls_GO[modality_nm] / Z
+
+    return optimal_weights
+
+
+# Optimal Gradient Blend
+#   x << N
+def get_optimal_gradient_blend_weights_og(delta_OGs):
+    """Get the weights of clients for optimal gradient blending
+
+    Args:
+        delta_OGs (list): each item is a tuple that contains (delta_O, delta_G)
+
+        The structure of the above two dicts should be: (for example)
+            [(0.2, 0.45), (0.3, 0.67)]
+
+    The equation that is the same as the weights computation for modalities:
+        w^i = <∇L^*, v_i>/(σ_i)^2 * 1/Z
+            = <∇L^*, v_i>/(σ_i)^2 * 1/(sum_i <∇L^*, v_i>/2*(σ_i)^2)
+            = G^i / (O^i)^2 * 1 / (sum_i G^i / (2 * (O^i)^2))
+
+        where G^i = G_N,n =  delta_G,
+                O^i = O_N,n =  O_N - O_n = delta_O
+    """
+    num_of_clients = len(delta_OGs)
+
+    Z = 0
+    clients_ratios = list()
+    for cli_i in range(num_of_clients):
+        cli_delta_O, cli_delta_G = delta_OGs[cli_i]
+
+        G_i = cli_delta_G
+        O_i = cli_delta_O
+
+        # models_GO[modality_nm] = G_i / (O_i * O_i)
+        Gi_div_sqr_Oi = G_i / (2 * O_i * O_i)
+
+        clients_ratios.append(Gi_div_sqr_Oi)
+
+        Z += Gi_div_sqr_Oi
+
+    optimal_weights = np.array(clients_ratios) / Z
+
+    return optimal_weights

plato/models/multimodal/fc_net.py

@@ -0,0 +1,77 @@
+"""
+Build the fully-connected net based on the configs
+
+An example of the fc layer is:
+fuse_model = dict(
+    type='FullyConnectedHead',
+    num_classes=400,
+    in_channels=rgb_model['cls_head']['in_channels'] +
+    flow_model['cls_head']['in_channels'] +
+    audio_model['cls_head']['in_channels'],
+    hidden_layer_size=[2014, 512],
+    dropout_ratio=0.5,
+)
+"""
+
+from collections import OrderedDict
+
+from torch import nn
+
+
+# FullyConnectedHead
+def build_fc_from_config(fc_configs):
+    """Build one fully-connected network based our settings"""
+
+    out_classes_n = fc_configs["num_classes"]
+
+    hidden_layer_dims = fc_configs["hidden_layer_size"]
+    hidden_n = len(hidden_layer_dims)
+
+    drop_out_porbs = (
+        fc_configs["dropout_ratio"]
+        if isinstance(fc_configs["dropout_ratio"], list)
+        else [fc_configs["dropout_ratio"]]
+    )
+    drop_out_porbs = (
+        drop_out_porbs * hidden_n if len(drop_out_porbs) == 1 else drop_out_porbs
+    )
+
+    fc_strcuture = OrderedDict()
+
+    # the first layer
+    fc_strcuture["fc1"] = nn.Linear(fc_configs["in_channels"], hidden_layer_dims[0])
+    fc_strcuture["relu1"] = nn.ReLU()
+    fc_strcuture["drop1"] = nn.Dropout(p=drop_out_porbs[0])
+    # the hidden layer
+    for hidden_l_i, layer_dim in enumerate(hidden_layer_dims):
+        layer_dim = hidden_layer_dims[hidden_l_i]
+        if hidden_l_i == hidden_n - 1:  # the final prediction layer
+            layer_name = "fcf"  # the final layer
+
+            fc_strcuture[layer_name] = nn.Linear(layer_dim, out_classes_n)
+            fc_strcuture["sigmoid"] = nn.Sigmoid()
+
+        else:
+            next_layer_in_dim = hidden_layer_dims[hidden_l_i + 1]
+            layer_name = "fc" + str(hidden_l_i + 2)
+            relu_name = "relu" + str(hidden_l_i + 2)
+            dropout_name = "dropout" + str(hidden_l_i + 2)
+            fc_strcuture[layer_name] = nn.Linear(layer_dim, next_layer_in_dim)
+            fc_strcuture[relu_name] = nn.ReLU()
+            fc_strcuture[dropout_name] = nn.Dropout(p=drop_out_porbs[hidden_l_i + 1])
+
+    fc_net = nn.Sequential(fc_strcuture)
+
+    return fc_net
+
+
+if __name__ == "__main__":
+    fuse_model = dict(
+        type="FullyConnectedHead",
+        num_classes=400,
+        in_channels=512 * 3,
+        hidden_layer_size=[1024, 512],
+        dropout_ratio=0.5,
+    )
+    fc_model = build_fc_from_config(fuse_model)
+    print(fc_model)

plato/models/multimodal/fusion_net.py

@@ -0,0 +1,78 @@
+# @Date    : 2021-06-27 13:23:05
+"""
+This multimodal network is the core network used in our paper.
+    It can receives three datasets from three modalities(RGB, optical flow, and audio)
+    and then process them with
+three different networks:
+    - RGB and flow:  ResNet3D from the paper 'A closer look at spatiotemporal
+                    convolutions for action recognition'.
+                    This is the r2plus1d in the mmaction packet
+    - audio: ResNet: Deep residual learning for image recognition. In CVPR, 2016.
+    both with 50 layers.
+
+    - For fusion, we use a two-FC-layer network on concatenated
+        features from visual and audio backbones,
+    followed by one prediction layer.
+"""
+
+import torch
+import torch.nn as nn
+
+from mmaction.models import build_loss
+
+from plato.models.multimodal import fc_net
+
+
+class ConcatFusionNet(nn.Module):
+    """This supports concat the features of different modalities to one vector"""
+
+    def __init__(self, support_modalities, modalities_fea_dim, net_configs):
+        super(ConcatFusionNet, self).__init__()
+
+        # the support modality name is the pre-defined order that must be
+        # followed in the forward process
+        #   especially in the fusion part
+        self.support_modality_names = support_modalities  # a list
+        self.modalities_fea_dim = modalities_fea_dim
+        self.net_configs = net_configs
+        # 1 build the model based on the configurations
+        self._fuse_net = fc_net.build_fc_from_config(net_configs)
+
+        self.loss_cls = build_loss(self.net_configs["loss_cls"])
+
+    def create_fusion_feature(self, batch_size, modalities_features_container):
+        """[summary]
+
+        Args:
+            modalities_features_container (dict): [key is the name of the modality
+                                                while the value is the corresponding features]
+            modalities_features_dims_container (dict): [key is the name of the modality
+                                                while the value is the defined dim of the feature]
+        """
+        # obtain the fused feats by concating the modalities features
+        #   The order should follow the that in the support_modality_names
+        modalities_feature = []
+        for modality_name in self.support_modality_names:
+            if modality_name not in modalities_features_container:
+                modality_dim = self.modalities_fea_dim[modality_name]
+                # insert all zeros features if that modality is missing
+                modality_feature = torch.zeros(size=(batch_size, modality_dim))
+            else:
+                modality_feature = modalities_features_container[modality_name]
+
+            modalities_feature.append(modality_feature)
+
+        fused_feat = torch.cat(modalities_feature, 1)
+
+        return fused_feat
+
+    def forward(self, fused_features, gt_labels, return_loss):
+        """Forward the network"""
+        fused_cls_score = self._fuse_net(fused_features)
+
+        if return_loss:
+            fused_loss = self.loss_cls(fused_cls_score, gt_labels)
+
+            return [fused_cls_score, fused_loss]
+        else:
+            return [fused_cls_score]

plato/models/multimodal/multimodal_module.py

@@ -0,0 +1,152 @@
+"""
+This multimodal network is the core network used in our paper.
+    It can receives three datasets from three modalities(RGB, optical flow, and audio)
+    and then process them with
+three different networks:
+    - RGB and flow:  ResNet3D from the paper 'A closer look at spatiotemporal
+                    convolutions for action recognition'.
+                    This is the r2plus1d in the mmaction packet
+    - audio: ResNet: Deep residual learning for image recognition. In CVPR, 2016.
+    both with 50 layers.
+
+    - For fusion, we use a two-FC-layer network on concatenated features from
+        visual and audio backbones,
+    followed by one prediction layer.
+"""
+
+import logging
+
+import torch.nn as nn
+
+from plato.models.multimodal import base_net
+from plato.models.multimodal import fusion_net
+
+
+class DynamicMultimodalModule(nn.Module):
+    """DynamicMultimodalModule network.
+        This network supports the learning of several modalities (the modalities can be dynamic)
+
+    Args:
+        multimodal_nets_configs (namedtuple): a namedtuple contains the configurations for
+                                            different modalities, 'rgb_model', 'audio_model',
+                                            'flow_model', 'text_model'
+    """
+
+    def __init__(
+        self,
+        support_modality_names,
+        multimodal_nets_configs,  # multimodal_data_model
+        is_fused_head=True,
+    ):  # a cls head makes prediction based on the fused multimodal feature
+        super().__init__()
+
+        # ['rgb', "flow", "audio"]
+        self.support_modality_names = support_modality_names
+        self.support_nets = [
+            mod_nm + "_model" for mod_nm in self.support_modality_names
+        ]
+
+        self.is_fused_head = is_fused_head
+
+        assert all(
+            [s_net in multimodal_nets_configs.keys() for s_net in self.support_nets]
+        )
+        self.name_net_mapper = {}
+        self.modality_fea_dims_mapper = {}
+        for idx, modality_net in enumerate(self.support_nets):
+            modality_name = self.support_modality_names[idx]
+            modality_net = self.support_nets[idx]
+            if modality_net in multimodal_nets_configs.keys():
+                logging.info("Building the %s......", modality_net)
+                net_config = multimodal_nets_configs[modality_net]
+                is_head_included = "cls_head" in net_config.keys()
+                logging.info("The head is defined")
+
+                if is_head_included:
+                    # the feature dimension is the input dimension of the cls head
+                    fea_dims = net_config["cls_head"]["in_channels"]
+                    self.modality_fea_dims_mapper[modality_name] = fea_dims
+
+                self.name_net_mapper[modality_name] = base_net.BaseClassificationNet(
+                    net_configs=net_config, is_head_included=is_head_included
+                )
+
+        if is_fused_head:
+            fuse_net_config = multimodal_nets_configs["fuse_model"]
+
+            if "modalities_feature_dim" in list(fuse_net_config.keys()):
+                self.modality_fea_dims_mapper.update(
+                    fuse_net_config["modalities_feature_dim"]
+                )
+
+            self.cat_fusion_net = fusion_net.ConcatFusionNet(
+                support_modalities=support_modality_names,
+                modalities_fea_dim=self.modality_fea_dims_mapper,
+                net_configs=fuse_net_config,
+            )
+
+    def assing_weights(self, net_name, weights):
+        """Assign the weights to the specific network"""
+        self.name_net_mapper[net_name].load_state_dict(weights, strict=True)
+
+    def _freeze_stages(self):
+        """Prevent all the parameters from being optimized before
+        ``self.frozen_stages``."""
+        if self.frozen_stages >= 0:
+            self.conv1.eval()
+            for param in self.conv1.parameters():
+                param.requires_grad = False
+
+        for i in range(1, self.frozen_stages + 1):
+            layer_module = getattr(self, f"layer{i}")
+            layer_module.eval()
+            for param in layer_module.parameters():
+                param.requires_grad = False
+
+    def forward(self, data_container, label=None, return_loss=True, **kwargs):
+        """[Forward the data to the whole net]
+
+        Args:
+            data_container (dict): [key is the name of the modality
+                                    while the value is batch of data]
+            label (torch.tensor, optional): [the lable of the sample]. Defaults to None.
+            return_loss (bool, optional): [whether return the loss ]. Defaults to True.
+        """
+        modalities_pred_scores_container = dict()
+        modalities_losses_container = dict()
+        modalities_features_container = dict()
+
+        for modality_name in data_container.keys():
+            modality_net = self.name_net_mapper[modality_name]
+            modality_ipt_data = data_container[modality_name]
+            batch_size = modality_ipt_data.shape[0]
+
+            logging.debug("modality_name: %s", modality_name)
+            logging.debug("modality_net: %s", modality_net)
+            logging.debug("modality_net inner net: %s", modality_net.get_net())
+            logging.debug("modality_ipt_data: %s", modality_ipt_data.shape)
+            logging.debug("batch_size: %s", batch_size)
+
+            # obtain the modality fea and the class opt
+            modality_opt = modality_net.forward(
+                ipt_data=modality_ipt_data, label=label, return_loss=return_loss
+            )
+
+            modalities_features_container[modality_name] = modality_opt[0]
+            modalities_pred_scores_container[modality_name] = modality_opt[1]
+            modalities_losses_container[modality_name] = modality_opt[2]
+
+        if self.is_fused_head:
+            # obtain the fused feats by concating the modalities features
+            #   The order should follow the that in the support_modality_names
+            fused_feat = self.cat_fusion_net.create_fusion_feature(
+                batch_size=batch_size,
+                modalities_features_container=modalities_features_container,
+            )
+            fused_cls_score, fused_loss = self.cat_fusion_net.forward(
+                fused_feat, label, return_loss=return_loss
+            )
+            modalities_pred_scores_container["fused"] = fused_cls_score
+            modalities_losses_container["fused"] = fused_loss
+
+        return modalities_pred_scores_container, modalities_losses_container

plato/models/registry.py

@@ -0,0 +1,99 @@
+"""
+The registry for machine learning models.
+
+Having a registry of all available classes is convenient for retrieving an instance
+based on a configuration at run-time.
+"""
+
+from typing import Any, Dict, TypedDict, cast
+
+from plato.config import Config
+from plato.models import (
+    cnn_encoder,
+    dcgan,
+    general_multilayer,
+    huggingface,
+    lenet5,
+    multilayer,
+    resnet,
+    torch_hub,
+    vgg,
+    vit,
+)
+
+registered_models = {
+    "lenet5": lenet5.Model,
+    "dcgan": dcgan.Model,
+    "multilayer": multilayer.Model,
+}
+
+registered_factories = {
+    "resnet": resnet.Model,
+    "vgg": vgg.Model,
+    "cnn_encoder": cnn_encoder.Model,
+    "general_multilayer": general_multilayer.Model,
+    "torch_hub": torch_hub.Model,
+    "huggingface": huggingface.Model,
+    "vit": vit.Model,
+}
+
+
+class ModelKwargs(TypedDict, total=False):
+    model_name: str
+    model_type: str
+    model_params: Dict[str, Any]
+
+
+def get(**kwargs: Any) -> Any:
+    """Get the model with the provided name."""
+    config = Config()
+
+    # Get model name
+    model_name: str = ""
+    if "model_name" in kwargs:
+        model_name = cast(str, kwargs["model_name"])
+    elif hasattr(config, "trainer"):
+        trainer = getattr(config, "trainer")
+        if hasattr(trainer, "model_name"):
+            model_name = getattr(trainer, "model_name")
+
+    # Get model type
+    model_type: str = ""
+    if "model_type" in kwargs:
+        model_type = cast(str, kwargs["model_type"])
+    elif hasattr(config, "trainer"):
+        trainer = getattr(config, "trainer")
+        if hasattr(trainer, "model_type"):
+            model_type = getattr(trainer, "model_type")
+
+    # If model_type is still empty, derive it from model_name
+    if not model_type and model_name:
+        model_type = model_name.split("_")[0]
+
+    # Get model parameters
+    model_params: Dict[str, Any] = {}
+    if "model_params" in kwargs:
+        model_params = cast(Dict[str, Any], kwargs["model_params"])
+    elif hasattr(config, "parameters"):
+        parameters = getattr(config, "parameters")
+        if hasattr(parameters, "model"):
+            model = getattr(parameters, "model")
+            if hasattr(model, "_asdict"):
+                model_params = model._asdict()
+
+    if model_type in registered_models:
+        registered_model = registered_models[model_type]
+        return registered_model(**model_params)
+
+    if model_type in registered_factories:
+        return registered_factories[model_type].get(
+            model_name=model_name, **model_params
+        )
+
+    # The YOLOv8 model needs special handling as it needs to import third-party packages
+    if model_name == "yolov8":
+        from plato.models import yolov8
+
+        return yolov8.Model.get()
+
+    raise ValueError(f"No such model: {model_name}")