fusion-bench 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

fusion_bench/models/hf_clip.py

@@ -1,4 +1,5 @@
-from typing import Callable, Iterable, List  # noqa: F401
+import logging
+from typing import TYPE_CHECKING, Callable, Iterable, List  # noqa: F401
 
 import torch
 from torch import Tensor, nn
@@ -7,6 +8,11 @@ from transformers.models.clip.modeling_clip import BaseModelOutputWithPooling
 
 from fusion_bench.utils.devices import get_device
 
+if TYPE_CHECKING:
+    from fusion_bench.models.surgery.surgerymodelwrapper import SurgeryModelWrapper
+
+log = logging.getLogger(__name__)
+
 default_templates = [
     lambda c: f"a photo of a {c}",
 ]
@@ -33,6 +39,7 @@ class HFCLIPClassifier(nn.Module):
         self,
         clip_model: CLIPModel,
         processor: CLIPProcessor,
+        extra_module=None,
     ):
         """
         Initialize the HFCLIPClassifier.
@@ -56,6 +63,8 @@ class HFCLIPClassifier(nn.Module):
             persistent=False,
         )
 
+        self.extra_module = extra_module
+
     @property
     def text_model(self):
         """Get the text model component of CLIP."""
@@ -111,7 +120,13 @@ class HFCLIPClassifier(nn.Module):
 
         self.zeroshot_weights = zeroshot_weights
 
-    def forward(self, images, return_image_embeds=False, return_dict=False):
+    def forward(
+        self,
+        images: Tensor,
+        return_image_embeds=False,
+        return_dict=False,
+        task_name=None,
+    ):
         """
         Perform forward pass for zero-shot image classification.
 
@@ -120,6 +135,9 @@ class HFCLIPClassifier(nn.Module):
 
         Args:
             images (Tensor): Input images to classify.
+            return_image_embeds (bool): Whether to return the image embeddings.
+            return_dict (bool): Whether to return a dictionary with logits and image embeddings.
+            task_name (Optional[str]): The name of the task.
 
         Returns:
             Tensor: Classification logits for each input image.
@@ -131,16 +149,22 @@ class HFCLIPClassifier(nn.Module):
             raise ValueError("Must set classification task before forward pass")
         text_embeds = self.zeroshot_weights
 
-        image_embeds = self.vision_model(images)
-        if isinstance(image_embeds, Tensor):
-            pass
-        elif isinstance(image_embeds, BaseModelOutputWithPooling):
-            image_embeds = image_embeds[1]
-        image_embeds = self.clip_model.visual_projection(image_embeds)
-
+        image_embeds = self.get_image_features(images)
         # normalize embeddings
         image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
 
+        if (
+            hasattr(self.vision_model, "is_surgery_model")
+            and self.vision_model.is_surgery_model
+        ):
+            # Dealing with the surgery model, for more details, please refer to:
+            # (ICML 2024) Yang, et.al. Representation Surgery for Multi-Task Model Merging
+            # https://arxiv.org/abs/2402.02705
+            self.vision_model: "SurgeryModelWrapper" = self.vision_model
+            image_embeds, _, _ = self.vision_model.compute_surgery_features(
+                image_embeds, dataset_name=task_name
+            )
+
         # cosine similarity
         logit_scale = self.clip_model.logit_scale.exp()
         logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * logit_scale
@@ -156,3 +180,20 @@ class HFCLIPClassifier(nn.Module):
                 return logits_per_image, image_embeds
             else:
                 return logits_per_image
+
+    def get_image_features(self, images: Tensor) -> Tensor:
+        """
+        Compute the image embeddings.
+
+        Returns:
+            image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by
+            applying the projection layer to the pooled output of [`CLIPVisionModel`].
+        """
+
+        image_embeds = self.vision_model(images)
+        if isinstance(image_embeds, Tensor):
+            pass
+        elif isinstance(image_embeds, BaseModelOutputWithPooling):
+            image_embeds = image_embeds[1]
+        image_embeds = self.clip_model.visual_projection(image_embeds)
+        return image_embeds

fusion_bench/models/rankone_moe.py

@@ -0,0 +1,410 @@
+import functools
+import logging
+from typing import Dict, List, Tuple  # noqa: F401
+
+import torch
+import torch.func
+from torch import Tensor, nn
+from torch.func import functional_call
+from torch.nn import functional as F
+
+from fusion_bench.utils.type import StateDictType
+
+log = logging.getLogger(__name__)
+
+
+def join_list(list_of_list: List[List]):
+    ans = []
+    for l in list_of_list:
+        ans.extend(l)
+    return ans
+
+
+def del_attr(obj, names: List[str]):
+    """
+    Deletes an attribute from an object recursively.
+
+    Args:
+        obj (object): Object to delete attribute from.
+        names (list): List of attribute names to delete recursively.
+    """
+    if len(names) == 1:
+        delattr(obj, names[0])
+    else:
+        del_attr(getattr(obj, names[0]), names[1:])
+
+
+def set_attr(obj, names: List[str], val):
+    """
+    Sets an attribute of an object recursively.
+
+    Args:
+        obj (object): Object to set attribute of.
+        names (list): List of attribute names to set recursively.
+        val (object): Value to set the attribute to.
+    """
+    if len(names) == 1:
+        setattr(obj, names[0], val)
+    else:
+        set_attr(getattr(obj, names[0]), names[1:], val)
+
+
+def get_attr(obj, names: List[str]):
+    """
+    Gets an attribute of an object recursively.
+
+    Args:
+        obj (object): Object to get attribute of.
+        names (list): List of attribute names to get recursively.
+
+    Returns:
+        object: The attribute of the object.
+    """
+    if len(names) == 1:
+        return getattr(obj, names[0])
+    else:
+        return get_attr(getattr(obj, names[0]), names[1:])
+
+
+class Depth_0_Gate(nn.Module):
+    def __init__(self, num_experts: int):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(num_experts), requires_grad=True)
+
+    def init_weight(self, init_lambda: float):
+        nn.init.constant_(self.weight, init_lambda)
+
+    def forward(self, *args, **kwargs) -> Tensor:
+        return self.weight
+
+
+class Depth_1_Gate(nn.Module):
+    def __init__(self, hidden_size: int, num_experts: int):
+        super().__init__()
+        self.fc = nn.Linear(hidden_size, num_experts, bias=True)
+
+    def init_weight(self, init_lambda: float):
+        nn.init.normal_(self.fc.weight, std=0.01)
+        nn.init.constant_(self.fc.bias, init_lambda)
+
+    def forward(self, hidden_states: Tensor) -> Tensor:
+        return self.fc(hidden_states)
+
+
+class Depth_2_Gate(nn.Module):
+    def __init__(self, hidden_size: int, num_experts: int):
+        super().__init__()
+        self.fc1 = nn.Linear(hidden_size, num_experts * 2, bias=True)
+        self.fc2 = nn.Linear(num_experts * 2, num_experts, bias=True)
+
+    def init_weight(self, init_lambda: float):
+        nn.init.normal_(self.fc1.weight, std=0.01)
+        nn.init.zeros_(self.fc1.bias)
+        nn.init.normal_(self.fc2.weight, std=0.01)
+        nn.init.constant_(self.fc2.bias, init_lambda)
+
+    def forward(self, hidden_states: Tensor) -> Tensor:
+        hidden_states = F.relu(self.fc1(hidden_states))
+        return self.fc2(hidden_states)
+
+
+def construct_rankone_moe_gate(
+    hidden_size: int,
+    num_experts: int,
+    init_lambda: float,
+    num_hidden_layers: int = 2,
+):
+    if num_hidden_layers == 0:
+        gate = Depth_0_Gate(num_experts)
+    elif num_hidden_layers == 1:
+        gate = Depth_1_Gate(hidden_size, num_experts)
+    elif num_hidden_layers == 2:
+        gate = Depth_2_Gate(hidden_size, num_experts)
+    else:
+        raise ValueError(f"Unsupported number of hidden layers: {num_hidden_layers}")
+
+    gate.num_hidden_layers = num_hidden_layers
+    gate.init_weight(init_lambda)
+    return gate
+
+
+class ExpertNotTrainedError(Exception):
+    pass
+
+
+def _is_all_zeros(tensor: Tensor | List[Tensor]) -> bool:
+    """
+    Check if a tensor or a list of tensors are all zeros.
+    """
+    if isinstance(tensor, Tensor):
+        return torch.allclose(tensor, torch.zeros_like(tensor))
+    else:
+        return all(_is_all_zeros(t) for t in tensor)
+
+
+def _svd(w: Tensor, full_matrices=True) -> Tuple[Tensor, Tensor, Tensor]:
+    """
+    Perform Singular Value Decomposition (SVD) on a tensor.
+    """
+    u, s, vh = torch.linalg.svd(
+        w, full_matrices=full_matrices, driver="gesvd" if w.is_cuda else None
+    )
+    v = vh.T
+    return u, s, v
+
+
+def svd(
+    w: Tensor, full_matrices=True, accelerator=None
+) -> Tuple[Tensor, Tensor, Tensor]:
+    """
+    Perform SVD on a tensor, optionally using a specified accelerator.
+    """
+    if accelerator is None:
+        return _svd(w, full_matrices=full_matrices)
+    original_device = w.device
+    w = w.to(accelerator)
+    u, s, v = _svd(w)
+    return u.to(original_device), s.to(original_device), v.to(original_device)
+
+
+def fun_joint_svd(
+    w_list: List[Tensor], accelerator=None
+) -> Tuple[Tensor, Tensor, Tensor]:
+
+    w = torch.cat(w_list, dim=1)  # stacked_matrix
+    original_device = w.device
+    if accelerator is not None:
+        w = w.to(accelerator)
+    u_c, s_c, vh_c = torch.linalg.svd(
+        w, full_matrices=False, driver="gesvd" if w.is_cuda else None
+    )
+
+    svd_list = []
+    offset = 0
+    for matrix in w_list:
+        n_cols = matrix.size(1)
+        u = u_c
+        s = s_c
+        vh_ = vh_c[:, offset : offset + n_cols]
+        v = vh_.T
+        svd_list.append(
+            [u.to(original_device), s.to(original_device), v.to(original_device)]
+        )
+
+        offset += n_cols
+    return svd_list
+
+
+class RankOneMoE(nn.Module):
+    # variable to store the merged state dict temporarily
+    _merged_state_dict: StateDictType = None
+
+    def __init__(
+        self,
+        hidden_size: int,
+        base_model: nn.Module,
+        expert_models: List[nn.Module],
+        init_lambda: float = 0.2,
+        batch_first: bool = False,
+        router_hidden_layers: int = 2,
+        batch_reduce: bool = False,
+        svd_accelerator=False,
+        rank_k: int = -1,
+        select_k: int = -1,
+    ):
+        """
+        Initializes the RankOneMoE class.
+        https://github.com/EnnengYang/RankOne-MoE
+
+        Args:
+            hidden_size (int): The size of the hidden layer in the models.
+            base_model (nn.Module): The base model that will be used as a reference for the expert models.
+            expert_models (List[nn.Module]): A list of expert models that will be combined.
+            init_lambda (float, optional): The initial lambda value for the weight ensembling gate. Defaults to 0.2.
+            batch_first (bool, optional): If True, the input tensors are expected to have the batch size as the first dimension. Defaults to False.
+            router_hidden_layers (int, optional): The number of hidden layers in the router. Defaults to 2.
+            batch_reduce (bool): If True, the batch dimension of routing weights is reduced. Defaults to False.
+        """
+        super().__init__()
+        self.num_experts = len(expert_models)
+        self.hidden_size = hidden_size
+        self.batch_first = batch_first
+        self.batch_reduce = batch_reduce
+        self.svd_accelerator = svd_accelerator
+        self.rank_k = rank_k
+        self.select_k = select_k
+        self.init_lambda = init_lambda
+
+        self.gate = construct_rankone_moe_gate(
+            hidden_size=hidden_size,
+            num_experts=int(self.num_experts * self.rank_k),
+            init_lambda=init_lambda,
+            num_hidden_layers=router_hidden_layers,
+        )
+
+        # compute the task vectors
+        for name, param in base_model.named_parameters():
+            if not param.requires_grad:
+                for m in expert_models:
+                    del_attr(m, name.split("."))
+            else:
+                for m in expert_models:
+                    get_attr(m, name.split(".")).data = (
+                        get_attr(m, name.split(".")) - param
+                    )
+
+        # fix base model and expert models
+        self.base_model = base_model.requires_grad_(False)
+        for m in expert_models:
+            m.requires_grad_(False)
+
+        # task vecotr  (only bias term)
+        self.task_vectors_fc1_bias = nn.Parameter(
+            torch.stack([e.fc1.bias for e in expert_models], dim=0), requires_grad=False
+        )
+        self.task_vectors_fc2_bias = nn.Parameter(
+            torch.stack([e.fc2.bias for e in expert_models], dim=0), requires_grad=False
+        )
+
+        # SVD representation of task vector (only weight term)
+        self.task_vectors_fc1_u = nn.ParameterList()
+        self.task_vectors_fc1_svh = nn.ParameterList()
+        self.task_vectors_fc2_u = nn.ParameterList()
+        self.task_vectors_fc2_svh = nn.ParameterList()
+
+        for m in expert_models:
+            for name, param in m.named_parameters():
+                if ".weight" in name:
+
+                    if _is_all_zeros(param):
+                        # All fine-tuned models are identical to the pretrained model
+                        raise ExpertNotTrainedError()
+
+                    u, s, v = svd(param, accelerator=self.svd_accelerator)
+                    u = u[:, : self.rank_k]
+                    s = s[: self.rank_k]
+                    v = v[:, : self.rank_k]
+
+                    if "fc1.weight" == name:
+                        self.task_vectors_fc1_u.append(
+                            nn.Parameter(u.T, requires_grad=False)
+                        )
+                        self.task_vectors_fc1_svh.append(
+                            nn.Parameter((s * v).T, requires_grad=False)
+                        )
+                    elif "fc2.weight" == name:
+                        self.task_vectors_fc2_u.append(
+                            nn.Parameter(u.T, requires_grad=False)
+                        )
+                        self.task_vectors_fc2_svh.append(
+                            nn.Parameter((s * v).T, requires_grad=False)
+                        )
+
+        # remove the original module from fine-tuned models to save memory
+        for name, param in base_model.named_parameters():
+            name_list = name.split(".")
+            for m in expert_models:
+                set_attr(m, name_list, None)
+
+    @property
+    def forward_model(self):
+        return functools.partial(
+            functional_call,
+            self.base_model,
+            self._merged_state_dict,
+        )
+
+    def top_k_soft(self, s, k):
+        threshold, _ = torch.topk(s, k, largest=True, sorted=False)
+        min_threshold = threshold.min()
+        # sigmoid -> mask
+        mask = torch.sigmoid(100 * (s - min_threshold))
+        result = s * mask
+        return result
+
+    def merge_weights(self, expert_weights):
+        state_dict = self.base_model.state_dict(keep_vars=True)
+
+        # Select top-k experts from the expert pool for fusion
+        if self.select_k > 0:
+            expert_weights = self.top_k_soft(expert_weights, self.select_k)
+
+        for name in state_dict:
+            if name == "fc1.bias":
+                for param in self.task_vectors_fc1_bias:
+                    state_dict[name] = state_dict[name] + self.init_lambda * param
+            elif name == "fc2.bias":
+                for param in self.task_vectors_fc2_bias:
+                    state_dict[name] = state_dict[name] + self.init_lambda * param
+
+            elif name == "fc1.weight":
+                w_list = torch.split(
+                    expert_weights,
+                    int(expert_weights.size(-1) / self.num_experts),
+                    dim=-1,
+                )
+                for weight, u, svh in zip(
+                    w_list, self.task_vectors_fc1_u, self.task_vectors_fc1_svh
+                ):
+                    weight_diag = torch.diag(weight)
+                    weight_u = torch.mm(weight_diag, u)
+                    result = torch.matmul(weight_u.T, svh)
+                    state_dict[name] = state_dict[name] + result
+
+            elif name == "fc2.weight":
+                w_list = torch.split(
+                    expert_weights,
+                    int(expert_weights.size(-1) / self.num_experts),
+                    dim=-1,
+                )
+                for weight, u, svh in zip(
+                    w_list, self.task_vectors_fc2_u, self.task_vectors_fc2_svh
+                ):
+                    weight_diag = torch.diag(weight)
+                    weight_u = torch.mm(weight_diag, u)
+                    result = torch.matmul(weight_u.T, svh)
+                    state_dict[name] = state_dict[name] + result
+
+        self._merged_state_dict = state_dict
+        return state_dict
+
+    def forward(self, hidden_states: Tensor):
+        if self.gate.num_hidden_layers == 0:
+            gate_weights = self.gate()
+        else:
+            gate_weights = self.gate(hidden_states)
+            if self.batch_first:
+                # the input is in the shape of (batch_size, seq_len, hidden_size)
+                gate_weights = gate_weights.mean(dim=1)
+            else:
+                # the input is in the shape of (seq_len, batch_size, hidden_size)
+                gate_weights = gate_weights.mean(dim=0)
+
+        if self.gate.num_hidden_layers == 0:
+            self.merge_weights(gate_weights)
+            output_hidden_states = self.forward_model(hidden_states)
+        elif self.batch_reduce:
+            gate_weights = gate_weights.mean(dim=0)
+            self.merge_weights(gate_weights)
+            output_hidden_states = self.forward_model(hidden_states)
+        else:
+            output_hidden_states = []
+            for sample_idx, weights in enumerate(gate_weights):
+                self.merge_weights(weights)
+                if self.batch_first:
+                    output_hidden_states.append(
+                        self.forward_model(hidden_states[sample_idx : sample_idx + 1])
+                    )
+                else:
+                    output_hidden_states.append(
+                        self.forward_model(
+                            hidden_states[:, sample_idx : sample_idx + 1]
+                        )
+                    )
+            if self.batch_first:
+                output_hidden_states = torch.cat(output_hidden_states, dim=0)
+            else:
+                output_hidden_states = torch.cat(output_hidden_states, dim=1)
+
+        self._merged_state_dict = None
+        return output_hidden_states

fusion_bench/models/surgery/surgerymodelwrapper.py

@@ -0,0 +1,157 @@
+import math
+from typing import TYPE_CHECKING, List, Union, Callable, Generic
+
+import torch
+from torch import nn
+from transformers.models.clip.modeling_clip import (
+    CLIPVisionModel,
+    CLIPVisionTransformer,
+)
+from fusion_bench.utils.type import TorchModelType
+
+
+def regularize_name(name: str):
+    name = name.replace("-", "_")
+    name = name.replace(".", "_")
+    return name
+
+
+class SurgeryModelWrapper(torch.nn.Module, Generic[TorchModelType]):
+
+    is_surgery_model = True
+    """A flag to indicate that this is a surgery model."""
+
+    def __init__(
+        self,
+        model: TorchModelType,
+        test_datasets: List[str],
+        projection_dim: int = 512,
+        hidden_dim: int = 16,
+    ):
+        super(SurgeryModelWrapper, self).__init__()
+        self.model = model
+        self.model.requires_grad_(False)
+
+        self.test_datasets = test_datasets
+        self.non_linear_func = torch.nn.ReLU()
+
+        self.projection_dim = projection_dim
+        self.hidden_dim = hidden_dim
+
+        for dataset_name in test_datasets:
+            self.add_surgery_module(dataset_name)
+
+    def add_surgery_module(self, dataset_name: str):
+        """
+        Add a surgery module for a given dataset.
+
+        Args:
+            dataset_name (str): The name of the dataset.
+        """
+        dataset_name = regularize_name(dataset_name)
+
+        down_proj = torch.nn.Linear(self.projection_dim, self.hidden_dim, bias=False)
+        up_proj = torch.nn.Linear(self.hidden_dim, self.projection_dim, bias=False)
+
+        torch.nn.init.kaiming_uniform_(down_proj.weight, a=math.sqrt(5))
+        torch.nn.init.zeros_(up_proj.weight)
+
+        self.add_module(
+            "feature_mapping_to_head_down_proj_{}".format(dataset_name), down_proj
+        )
+        self.add_module(
+            "feature_mapping_to_head_up_proj_{}".format(dataset_name), up_proj
+        )
+
+    def collect_trainable_params(self):
+        trainable_params = []
+
+        # surgery parameter
+        for dataset_name in self.test_datasets:
+            dataset_name = regularize_name(dataset_name)
+            down_proj = getattr(
+                self, "feature_mapping_to_head_down_proj_{}".format(dataset_name)
+            )
+            up_proj = getattr(
+                self, "feature_mapping_to_head_up_proj_{}".format(dataset_name)
+            )
+            trainable_params.append(down_proj.weight)
+            trainable_params.append(up_proj.weight)
+        return trainable_params
+
+    def collect_surgery_module(self):
+        surgery_module = {}
+
+        # surgery parameter
+        for dataset_name in self.test_datasets:
+            dataset_name = regularize_name(dataset_name)
+            down_proj = getattr(
+                self, "feature_mapping_to_head_down_proj_{}".format(dataset_name)
+            )
+            up_proj = getattr(
+                self, "feature_mapping_to_head_up_proj_{}".format(dataset_name)
+            )
+            surgery_module[
+                "feature_mapping_to_head_down_proj_{}".format(dataset_name)
+            ] = down_proj
+            surgery_module[
+                "feature_mapping_to_head_up_proj_{}".format(dataset_name)
+            ] = up_proj
+
+        surgery_module["non_linear_func"] = self.non_linear_func
+
+        return surgery_module
+
+    def compute_surgery_features(
+        self,
+        compute_features_fn: Union[
+            torch.Tensor, Callable[[TorchModelType], torch.Tensor]
+        ],
+        dataset_name: str,
+    ):
+        """
+        Compute the surgery features.
+
+        Args:
+            compute_features_fn (Union[torch.Tensor, Callable[[nn.Module], torch.Tensor]]): A function that computes the features or a tensor that represents the features.
+            dataset_name (str): The name of the dataset.
+
+        Returns:
+            feature (torch.Tensor): The surgery features.
+            feature0 (torch.Tensor): The original features.
+            feature_sub (torch.Tensor): feature0 - feature.
+        """
+        dataset_name = regularize_name(dataset_name)
+
+        if isinstance(compute_features_fn, torch.Tensor):
+            feature = compute_features_fn
+        elif callable(compute_features_fn):
+            feature = compute_features_fn(self.model)
+        else:
+            raise ValueError(
+                "compute_features_fn must be a tensor or a callable, but got {}".format(
+                    type(compute_features_fn)
+                )
+            )
+
+        feature0 = feature
+
+        # feature bias
+        down_proj = getattr(
+            self, "feature_mapping_to_head_down_proj_{}".format(dataset_name)
+        )
+        up_proj = getattr(
+            self, "feature_mapping_to_head_up_proj_{}".format(dataset_name)
+        )
+        feature_sub = down_proj(feature)
+        feature_sub = self.non_linear_func(feature_sub)
+        feature_sub = up_proj(feature_sub)
+
+        # surgery feature
+        feature = feature0 - feature_sub
+
+        return feature, feature0, feature_sub
+
+    def forward(self, *args, **kwargs):
+        """The wrappered model should just forward like normal."""
+        return self.model(*args, **kwargs)

fusion_bench/models/utils.py

@@ -1,5 +1,6 @@
 from typing import List
 
+import torch
 from torch import nn
 
 
@@ -70,3 +71,10 @@ def find_layers_with_type(
         if isinstance(submodule, tuple(layer_types)):
             res[name] = submodule
     return res
+
+
+def disable_dropout(model: torch.nn.Module):
+    """Disable dropout in a model."""
+    for module in model.modules():
+        if isinstance(module, torch.nn.Dropout):
+            module.p = 0