fusion-bench 0.2.11__py3-none-any.whl → 0.2.13__py3-none-any.whl

fusion_bench/models/open_clip/utils.py

@@ -0,0 +1,311 @@
+import copy
+import os
+import pickle
+from collections import OrderedDict
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+def compute_l1_norm(
+    model1: nn.Module, model2: nn.Module
+) -> Tuple[torch.Tensor, Dict[str, float]]:
+    """
+    Computes the L1 norm between the parameters of two models.
+
+    Args:
+        model1 (nn.Module): The first model.
+        model2 (nn.Module): The second model.
+
+    Returns:
+        Tuple[torch.Tensor, Dict[str, float]]: A tuple containing the total L1 norm and a dictionary
+        with the L1 norm for each layer.
+
+    """
+    norms = dict()
+    l1_norm = 0.0
+    for (n, p1), p2 in zip(model1.named_parameters(), model2.parameters()):
+        layer_l1_norm = torch.norm(p1 - p2, 1)
+        l1_norm += layer_l1_norm
+        norms[n] = layer_l1_norm.item()
+
+    return l1_norm, norms
+
+
+def assign_learning_rate(param_group, new_lr):
+    param_group["lr"] = new_lr
+
+
+def _warmup_lr(base_lr, warmup_length, step):
+    return base_lr * (step + 1) / warmup_length
+
+
+def cosine_lr(optimizer, base_lrs, warmup_length, steps):
+    if not isinstance(base_lrs, list):
+        base_lrs = [base_lrs for _ in optimizer.param_groups]
+    assert len(base_lrs) == len(optimizer.param_groups)
+
+    def _lr_adjuster(step):
+        for param_group, base_lr in zip(optimizer.param_groups, base_lrs):
+            if step < warmup_length:
+                lr = _warmup_lr(base_lr, warmup_length, step)
+            else:
+                e = step - warmup_length
+                es = steps - warmup_length
+                lr = 0.5 * (1 + np.cos(np.pi * e / es)) * base_lr
+            assign_learning_rate(param_group, lr)
+
+    return _lr_adjuster
+
+
+def accuracy(output: torch.Tensor, target: torch.Tensor, topk: List[int] = (1,)):
+    pred = output.topk(max(topk), 1, True, True)[1].t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+    return [
+        float(correct[:k].reshape(-1).float().sum(0, keepdim=True).cpu().numpy())
+        for k in topk
+    ]
+
+
+def torch_load_old(save_path: str, device=None):
+    with open(save_path, "rb") as f:
+        classifier = pickle.load(f)
+    if device is not None:
+        classifier = classifier.to(device)
+    return classifier
+
+
+def torch_save(model, save_path, save_state_dict=True):
+    # TODO: hacky way to save state dict
+    if save_state_dict and isinstance(model, torch.nn.Module):
+        model = model.state_dict()
+    if os.path.dirname(save_path) != "":
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    torch.save(model, save_path)
+
+
+def torch_load(save_path, device=None):
+    model = torch.load(save_path, map_location="cpu")
+    if device is not None:
+        model = model.to(device)
+    return model
+
+
+def get_logits(inputs, classifier):
+    assert callable(classifier)
+    if hasattr(classifier, "to"):
+        classifier = classifier.to(inputs.device)
+    return classifier(inputs)
+
+
+def get_probs(inputs, classifier):
+    if hasattr(classifier, "predict_proba"):
+        probs = classifier.predict_proba(inputs.detach().cpu().numpy())
+        return torch.from_numpy(probs)
+    logits = get_logits(inputs, classifier)
+    return logits.softmax(dim=1)
+
+
+class LabelSmoothing(torch.nn.Module):
+    def __init__(self, smoothing=0.0):
+        super(LabelSmoothing, self).__init__()
+        self.confidence = 1.0 - smoothing
+        self.smoothing = smoothing
+
+    def forward(self, x, target):
+        logprobs = torch.nn.functional.log_softmax(x, dim=-1)
+
+        nll_loss = -logprobs.gather(dim=-1, index=target.unsqueeze(1))
+        nll_loss = nll_loss.squeeze(1)
+        smooth_loss = -logprobs.mean(dim=-1)
+        loss = self.confidence * nll_loss + self.smoothing * smooth_loss
+        return loss.mean()
+
+
+class DotDict(dict):
+    """dot.notation access to dictionary attributes"""
+
+    __getattr__ = dict.get
+    __setattr__ = dict.__setitem__
+    __delattr__ = dict.__delitem__
+
+
+def find_optimal_coef(
+    results: Dict[str, Any],
+    metric: str = "avg_normalized_top1",
+    minimize: bool = False,
+    control_metric: Optional[str] = None,
+    control_metric_threshold: float = 0.0,
+) -> float:
+    """
+    Finds the optimal coefficient based on the given results and metric.
+
+    Args:
+        results (Dict[str, Any]): A dictionary containing the results for different scaling coefficients.
+        metric (str, optional): The metric to optimize. Defaults to "avg_normalized_top1".
+        minimize (bool, optional): Whether to minimize the metric. Defaults to False.
+        control_metric (str, optional): The control metric to check against. Defaults to None.
+        control_metric_threshold (float, optional): The threshold value for the control metric. Defaults to 0.0.
+
+    Returns:
+        The optimal coefficient based on the given results and metric.
+    """
+    best_coef = None
+    if minimize:
+        best_metric = 1
+    else:
+        best_metric = 0
+    for scaling_coef in results.keys():
+        if control_metric is not None:
+            if results[scaling_coef][control_metric] < control_metric_threshold:
+                print(f"Control metric fell below {control_metric_threshold} threshold")
+                continue
+        if minimize:
+            if results[scaling_coef][metric] < best_metric:
+                best_metric = results[scaling_coef][metric]
+                best_coef = scaling_coef
+        else:
+            if results[scaling_coef][metric] > best_metric:
+                best_metric = results[scaling_coef][metric]
+                best_coef = scaling_coef
+    return best_coef
+
+
+def nonlinear_advantage(nonlinear_acc, linear_acc, num_classes):
+    """Computes the normalized non-linear advantage of a finetuned model.
+
+    The nonlinear_advantage is defined as:
+        error_rate(linear_model) - error_rate(nonlinear_model) / (1 - 1 / num_classes)
+    and takes values between [-1, 1]. A value of 0 indicates that the nonlinear
+    model is no better than the linear one. Meanwhile, a value of 1 indicates
+    that the nonlinear model is perfect and the linear trivial, and a value of
+    -1 indicates the opposite.
+    """
+    return (nonlinear_acc - linear_acc) / (1.0 - 1.0 / num_classes)
+
+
+def to_cuda(input_dict):
+    cuda_dict = {}
+    for key, value in input_dict.items():
+        cuda_dict[key] = value.to("cuda")
+    return cuda_dict
+
+
+def state_dict_to_vector(state_dict, remove_keys=[]):
+    shared_state_dict = copy.deepcopy(state_dict)
+    for key in remove_keys:
+        if key in shared_state_dict:
+            del shared_state_dict[key]
+    sorted_shared_state_dict = OrderedDict(sorted(shared_state_dict.items()))
+    return torch.nn.utils.parameters_to_vector(
+        [value.reshape(-1) for key, value in sorted_shared_state_dict.items()]
+    )
+
+
+def vector_to_state_dict(vector, state_dict, remove_keys=[]):
+    # create a reference dict to define the order of the vector
+    reference_dict = copy.deepcopy(state_dict)
+    for key in remove_keys:
+        if key in reference_dict:
+            del reference_dict[key]
+    sorted_reference_dict = OrderedDict(sorted(reference_dict.items()))
+
+    # create a shared state dict using the reference dict
+    torch.nn.utils.vector_to_parameters(vector, sorted_reference_dict.values())
+
+    # add back the encoder and decoder embedding weights.
+    if "transformer.shared.weight" in sorted_reference_dict:
+        for key in remove_keys:
+            sorted_reference_dict[key] = sorted_reference_dict[
+                "transformer.shared.weight"
+            ]
+    return sorted_reference_dict
+
+
+def add_ptm_to_tv(tv_dict, ptm_dict):
+    assert set(tv_dict.keys()) == set(
+        ptm_dict.keys()
+    ), "Differing parameter names in models."
+    final_dict = copy.deepcopy(tv_dict)
+    for k, v in ptm_dict.items():
+        final_dict[k] = tv_dict[k] + v
+    return final_dict
+
+
+def check_parameterNamesMatch(checkpoints):
+    parameter_names = set(checkpoints[0].keys())
+
+    if len(checkpoints) >= 2:
+        # raise ValueError("Number of models is less than 2.")
+        for checkpoint in checkpoints[1:]:
+            current_parameterNames = set(checkpoint.keys())
+            if current_parameterNames != parameter_names:
+                raise ValueError(
+                    "Differing parameter names in models. "
+                    f"The different parameters are {parameter_names.symmetric_difference(current_parameterNames)}"
+                )
+
+
+def check_state_dicts_equal(state_dict1, state_dict2):
+    if set(state_dict1.keys()) != set(state_dict2.keys()):
+        return False
+
+    for key in state_dict1.keys():
+        if not torch.equal(state_dict1[key], state_dict2[key]):
+            return False
+
+    return True
+
+
+def topk_values_mask(M, K=0.7, return_mask=False, reshape_mask=False):
+    if K == 100:
+        # print("Not applying mask")
+        if return_mask:
+            return M, torch.ones_like(M), None
+        else:
+            return M, torch.ones_like(M)
+
+    if K >= 1:
+        K /= 100
+
+    original_shape = M.shape
+    if M.dim() == 1:
+        M = M.unsqueeze(0)
+
+    n, d = M.shape
+    k = int(d * K)
+    k = d - k  # Keep top k elements instead of bottom k elements
+
+    # Find the k-th smallest element by magnitude for each row
+    kth_values, _ = M.abs().kthvalue(k, dim=1, keepdim=True)
+    # Create a mask tensor with True for the top k elements in each row
+    mask = M.abs() >= kth_values
+    final_mask = mask.squeeze() if original_shape == M.squeeze().shape else mask
+
+    if reshape_mask:
+        final_mask = final_mask.reshape(M.shape)
+
+    if return_mask:
+        return M * final_mask, final_mask.float().mean(dim=1), final_mask
+    else:
+        return M * final_mask, final_mask.float().mean(dim=1)
+
+
+def cleanup_linear(state_dict):
+    # The linear model also has keys for the reference point $\theta_0$ in the state dict with the prefix `params0`.
+    state_dict = {k: v for k, v in state_dict.items() if "params." in k}
+    return state_dict
+
+
+def get_ptm_linear(state_dict: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    # rename keys so that they match afterwards
+    state_dict_new = {
+        k.replace("params0", "params"): v
+        for k, v in state_dict.items()
+        if "params0." in k
+    }
+    state_dict_remaining = {k: v for k, v in state_dict.items() if "params." not in k}
+
+    return state_dict_new, state_dict_remaining

fusion_bench/models/open_clip/variables_and_paths.py

@@ -0,0 +1,56 @@
+from pathlib import Path
+from typing import Literal
+
+TQDM_BAR_FORMAT = "{l_bar}{bar:10}{r_bar}{bar:-10b}"
+MODELS = ["ViT-B-32", "ViT-B-16", "ViT-L-14"]
+OPENCLIP_CACHEDIR = Path(Path.home(), "openclip-cachedir", "open_clip").as_posix()
+CACHEDIR = None
+
+ALL_DATASETS = [
+    "Cars",
+    "DTD",
+    "EuroSAT",
+    "GTSRB",
+    "MNIST",
+    "RESISC45",
+    "SVHN",
+    "SUN397",
+    "STL10",
+    "OxfordIIITPet",
+    "Flowers102",
+    "CIFAR100",
+    "PCAM",
+    "FER2013",
+    "CIFAR10",
+    "Food101",
+    "FashionMNIST",
+    "RenderedSST2",
+    "EMNIST",
+    "KMNIST",
+]
+
+DATASETS_8 = ALL_DATASETS[:8]
+DATASETS_14 = ALL_DATASETS[:14]
+DATASETS_20 = ALL_DATASETS[:20]
+
+
+def cleanup_dataset_name(dataset_name: str):
+    return dataset_name.replace("Val", "") + "Val"
+
+
+def get_zeroshot_path(root, dataset, model):
+    return Path(
+        root, model, cleanup_dataset_name(dataset), f"nonlinear_zeroshot.pt"
+    ).as_posix()
+
+
+def get_finetuned_path(root, dataset, model):
+    return Path(
+        root, model, cleanup_dataset_name(dataset), f"nonlinear_finetuned.pt"
+    ).as_posix()
+
+
+def get_single_task_accuracies_path(model):
+    return Path(
+        "results/single_task", model, f"nonlinear_ft_accuracies.json"
+    ).as_posix()

fusion_bench/models/parameter_dict.py

@@ -1,4 +1,4 @@
-from typing import List, Mapping
+from typing import List, Mapping, Optional, Tuple
 
 import torch
 from torch import nn
@@ -6,7 +6,13 @@ from torch import nn
 __all__ = "ParamterDictModel"
 
 
-def set_attr(obj, names: List[str], val, check_parent: bool = False):
+def _set_attr(
+    obj,
+    names: List[str],
+    val,
+    check_parent: bool = False,
+    parent_builder=nn.Module,
+):
     """
     Sets an attribute of an object recursively.
 
@@ -20,8 +26,14 @@ def set_attr(obj, names: List[str], val, check_parent: bool = False):
         setattr(obj, names[0], val)
     else:
         if check_parent and not hasattr(obj, names[0]):
-            setattr(obj, names[0], nn.Module())
-        set_attr(getattr(obj, names[0]), names[1:], val, check_parent=check_parent)
+            setattr(obj, names[0], parent_builder())
+        _set_attr(
+            getattr(obj, names[0]),
+            names[1:],
+            val,
+            check_parent=check_parent,
+            parent_builder=parent_builder,
+        )
 
 
 def has_attr(obj, names: List[str]):
@@ -49,17 +61,19 @@ class ParameterDictModel(nn.Module):
 
     def __init__(
         self,
-        parameters: Mapping[str, nn.Parameter],
+        parameters: Optional[Mapping[str, nn.Parameter]] = None,
     ):
         super().__init__()
-        for name, param in parameters.items():
-            assert isinstance(param, nn.Parameter), f"{name} is not a nn.Parameter"
-            set_attr(
-                self,
-                name.split("."),
-                param,
-                check_parent=True,
-            )
+        if parameters is not None:
+            for name, param in parameters.items():
+                assert isinstance(param, nn.Parameter), f"{name} is not a nn.Parameter"
+                _set_attr(
+                    self,
+                    name.split("."),
+                    param,
+                    check_parent=True,
+                    parent_builder=self.__class__,
+                )
 
     def __repr__(self):
         """
@@ -73,3 +87,30 @@ class ParameterDictModel(nn.Module):
             param_repr = f"{name}: {param.size()}"
             param_reprs.append(param_repr)
         return f"{self.__class__.__name__}({', '.join(param_reprs)})"
+
+    def __getitem__(self, key: str):
+        if not has_attr(self, key.split(".")):
+            raise KeyError(f"Key {key} not found in {self}")
+        key = key.split(".")
+        obj = self
+        for k in key:
+            obj = getattr(obj, k)
+        return obj
+
+    def __setitem__(self, key: str, value: nn.Parameter):
+        if not has_attr(self, key.split(".")):
+            _set_attr(self, key.split("."), value, check_parent=True)
+        else:
+            _set_attr(self, key.split("."), value, check_parent=False)
+
+    def __contains__(self, key: str):
+        return has_attr(self, key.split("."))
+
+    def keys(self):
+        return [name for name, _ in self.named_parameters()]
+
+    def items(self) -> List[Tuple[str, nn.Parameter]]:
+        return [(name, self[name]) for name in self.keys()]
+
+    def values(self) -> List[nn.Parameter]:
+        return [self[name] for name in self.keys()]

fusion_bench/models/wrappers/layer_wise_fusion.py

@@ -16,6 +16,7 @@ import torch
 from torch import Tensor, nn
 from torch.func import functional_call
 
+from fusion_bench.models.utils import del_attr, get_attr, set_attr
 from fusion_bench.utils.type import StateDictType, TorchModelType
 
 __all__ = ["get_layer_wise_weights", "fuse_weights", "LayerWiseMergedModel"]
@@ -23,52 +24,6 @@ __all__ = ["get_layer_wise_weights", "fuse_weights", "LayerWiseMergedModel"]
 log = logging.getLogger(__name__)
 
 
-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:])
-
-
 def get_layer_wise_weights(
     num_models: int,
     num_layers: int,

fusion_bench/models/wrappers/layer_wise_fusion_doge_ta.py

@@ -10,132 +10,17 @@ import torch
 from torch import Tensor, nn
 from torch.func import functional_call
 
+from fusion_bench.models.utils import del_attr, get_attr, set_attr
 from fusion_bench.utils.state_dict_arithmetic import state_dict_add
 from fusion_bench.utils.type import StateDictType
 
+from .layer_wise_fusion import fuse_weights, get_layer_wise_weights
+
 __all__ = ["get_layer_wise_weights", "fuse_weights", "LayerWiseMergedModel"]
 
 log = logging.getLogger(__name__)
 
 
-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:])
-
-
-def get_layer_wise_weights(
-    num_models: int,
-    num_layers: int,
-    init_values: float = None,
-    dtype: torch.dtype = torch.float32,
-):
-    """
-    Return a tensor of layer-wise weights for the given number of models and layers.
-
-    Args:
-        num_models (int): The number of models to fuse.
-        num_layers (int): The number of layers in each model.
-        init_values (float, optional): The initial value for each weight. Defaults to 1.0 / num_models.
-        dtype (torch.dtype): dtype of weights. This should be the same with model dtype.
-
-    Returns:
-        Tensor: A tensor of shape (num_models, num_layers) containing the layer-wise weights.
-    """
-    assert num_models >= 1, f"num_models must be >= 1, got {num_models}"
-    assert num_layers >= 1, f"num_layers must be >= 1, got {num_layers}"
-    if init_values is None:
-        init_values = 1.0 / num_models
-    return torch.full((num_models, num_layers), init_values, dtype=dtype)
-
-
-def _fuse_weights(layer_wise_weight: Tensor, tensors: List[Tensor]):
-    """
-    Fuse the layer-wise weights with the given state dictionaries.
-
-    Args:
-        layer_wise_weight (Tensor): A tensor of shape (num_models,) containing the layer-wise weights.
-        state_dicts (List[Tensor]): A list of state dictionaries, each containing the weights for a single layer.
-
-    Returns:
-        Tensor: A tensor of shape (num_params,) containing the fused weights.
-    """
-    assert len(layer_wise_weight) == len(
-        tensors
-    ), f"layer_wise_weight.shape={layer_wise_weight.shape}, len(tensors)={len(tensors)}"
-    return sum(
-        layer_wise_weight[i] * w.to(layer_wise_weight.device)
-        for i, w in enumerate(tensors)
-    )
-
-
-def fuse_weights(
-    layer_wise_weight: Tensor, state_dicts: List[StateDictType]
-) -> StateDictType:
-    """
-    Fuse the weights of multiple models using layer-wise fusion.
-
-    Args:
-        layer_wise_weight (Tensor): A tensor of shape (num_models, num_layers) representing the weight of each layer for each model.
-        state_dicts (List[StateDict]): A list of state dictionaries, one for each model.
-
-    Returns:
-        A dictionary mapping each weight tensor key to the fused weight tensor.
-    """
-    num_models = len(state_dicts)
-    num_layers = len(state_dicts[0])
-    assert layer_wise_weight.shape == (
-        num_models,
-        num_layers,
-    ), f"layer_wise_weight.shape={layer_wise_weight.shape}, expected (num_models, num_layers): ({num_models}, {num_layers})"
-    return {
-        k: _fuse_weights(
-            layer_wise_weight[:, i], [state_dict[k] for state_dict in state_dicts]
-        )
-        for i, k in enumerate(state_dicts[0].keys())
-    }
-
-
 class LayerWiseMergedModel(nn.Module):
     _merged_state_dict: StateDictType = None
 
@@ -390,7 +275,7 @@ class LayerWiseMergedModel(nn.Module):
         layer_vectors_scale = layer_vectors * layer_lamdas.view(-1, 1, 1)
         sum_over_num_vectors = layer_vectors_scale.sum(dim=0)
 
-        layer_delta_scale = layer_delta.unsqueeze(0) * layer_lamdas.view(-1, 1, 1)
+        layer_delta_scale = layer_delta * layer_lamdas.view(-1, 1, 1)
         sum_over_delta = layer_delta_scale.sum(dim=0)
 
         # Iterate through each vector and calculate the loss one by one

fusion_bench/scripts/nyuv2_mtl_train.py

@@ -1,5 +1,5 @@
 R"""
-This script is used to train a multi-task learning (MTL) model on the NYUv2 dataset. 
+This script is used to train a multi-task learning (MTL) model on the NYUv2 dataset.
 """
 
 import importlib