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

fusion_bench/models/open_clip/modeling.py

@@ -0,0 +1,176 @@
+from typing import Callable, List
+
+import open_clip
+import torch
+from torch import Tensor
+
+from . import utils
+from .variables_and_paths import CACHEDIR, MODELS, OPENCLIP_CACHEDIR
+
+
+class ImageEncoder(torch.nn.Module):
+    R"""
+    Examples:
+
+        load the image encoder for a given model name
+
+        >>> from fusion_bench.models.open_clip import ImageEncoder
+        >>> image_encoder = ImageEncoder(model_name="ViT-B-32")
+    """
+
+    def __init__(self, model_name: str, keep_lang=False):
+        super().__init__()
+        assert (
+            model_name in MODELS
+        ), f"Invalid model name: {model_name}. Valid models are: {MODELS}"
+
+        if "__pretrained__" in model_name:
+            name, pretrained = model_name.split("__pretrained__")
+        elif "__init__" in model_name:
+            print("Using random initialization.")
+            name, pretrained = model_name.split("__init__")[0], None
+        else:
+            name = model_name
+            pretrained = "openai"
+        (
+            self.model,
+            self.train_preprocess,
+            self.val_preprocess,
+        ) = open_clip.create_model_and_transforms(
+            name, pretrained=pretrained, cache_dir=OPENCLIP_CACHEDIR
+        )
+
+        self.cache_dir = CACHEDIR
+
+        if not keep_lang and hasattr(self.model, "transformer"):
+            delattr(self.model, "transformer")
+
+    def forward(self, images):
+        assert self.model is not None
+        return self.model.encode_image(images)
+
+    def __call__(self, inputs):
+        return self.forward(inputs)
+
+    def save(self, filename):
+        print(f"Saving image encoder to {filename}")
+        utils.torch_save(self, filename)
+
+    @classmethod
+    def load(cls, model_name, filename):
+        print(f"Loading image encoder from {filename}")
+
+        state_dict = torch.load(filename, map_location="cpu")
+
+        model = cls(model_name)
+        model.load_state_dict(state_dict)
+        return model
+
+
+class ClassificationHead(torch.nn.Linear):
+    def __init__(
+        self,
+        normalize: bool,
+        weights: Tensor,
+        biases: Tensor = None,
+    ):
+        output_size, input_size = weights.shape
+        super().__init__(input_size, output_size)
+        self.normalize = normalize
+        if weights is not None:
+            self.weight = torch.nn.Parameter(weights.clone())
+        if biases is not None:
+            self.bias = torch.nn.Parameter(biases.clone())
+        else:
+            self.bias = torch.nn.Parameter(torch.zeros_like(self.bias))
+
+    def forward(self, inputs: Tensor):
+        if self.normalize:
+            inputs = inputs / inputs.norm(dim=-1, keepdim=True)
+        return super().forward(inputs)
+
+    def __call__(self, inputs: Tensor):
+        return self.forward(inputs)
+
+    def save(self, filename):
+        print(f"Saving classification head to {filename}")
+        utils.torch_save(self, filename, save_state_dict=False)
+
+    @classmethod
+    def load(cls, filename):
+        # print(f"Loading classification head from {filename}")
+        return utils.torch_load(filename)
+
+
+class ImageClassifier(torch.nn.Module):
+    train_preprocess: Callable
+    val_preprocess: Callable
+
+    def __init__(
+        self,
+        image_encoder: ImageEncoder,
+        classification_head: ClassificationHead,
+    ):
+        super().__init__()
+        self.image_encoder = image_encoder
+        self.classification_head = classification_head
+        if self.image_encoder is not None:
+            self.train_preprocess = self.image_encoder.train_preprocess
+            self.val_preprocess = self.image_encoder.val_preprocess
+
+    def freeze_head(self):
+        self.classification_head.weight.requires_grad_(False)
+        self.classification_head.bias.requires_grad_(False)
+
+    def forward(self, inputs: Tensor):
+        features = self.image_encoder(inputs)
+        outputs = self.classification_head(features)
+        return outputs
+
+    def __call__(self, inputs):
+        return self.forward(inputs)
+
+    def save(self, filename):
+        print(f"Saving image classifier to {filename}")
+        utils.torch_save(self, filename)
+
+    @classmethod
+    def load(cls, filename):
+        print(f"Loading image classifier from {filename}")
+        return utils.torch_load(filename)
+
+
+class MultiHeadImageClassifier(torch.nn.Module):
+    def __init__(
+        self,
+        image_encoder: ImageEncoder,
+        classification_heads: List[ClassificationHead],
+    ):
+        super().__init__()
+        self.image_encoder = image_encoder
+        self.classification_heads = torch.nn.ModuleList(classification_heads)
+        if self.image_encoder is not None:
+            self.train_preprocess = self.image_encoder.train_preprocess
+            self.val_preprocess = self.image_encoder.val_preprocess
+
+    def freeze_head(self):
+        for idx in range(len(self.classification_heads)):
+            self.classification_heads[idx].weight.requires_grad_(False)
+            self.classification_heads[idx].bias.requires_grad_(False)
+
+    def forward(self, inputs, head_idx):
+        features = self.image_encoder(inputs)
+        outputs = self.classification_heads[head_idx](features)
+        return outputs
+
+    def __call__(self, inputs, head_idx):
+        return self.forward(inputs, head_idx)
+
+    def save(self, filename):
+        print(f"Saving image classifier to {filename}")
+        utils.torch_save(self, filename)
+
+    @classmethod
+    def load(cls, filename):
+        print(f"Loading image classifier from {filename}")
+        return utils.torch_load(filename)

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/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

fusion_bench/taskpool/__init__.py

@@ -10,12 +10,13 @@ _import_structure = {
     "clip_vision": [
         "CLIPVisionModelTaskPool",
         "SparseWEMoECLIPVisionModelTaskPool",
-        "RankoneWEMoECLIPVisionModelTaskPool",
+        "RankoneMoECLIPVisionModelTaskPool",
     ],
     "dummy": ["DummyTaskPool"],
     "gpt2_text_classification": ["GPT2TextClassificationTaskPool"],
-    "nyuv2_taskpool": ["NYUv2TaskPool"],
     "llama": ["LlamaTestGenerationTaskPool"],
+    "nyuv2_taskpool": ["NYUv2TaskPool"],
+    "openclip_vision": ["OpenCLIPVisionModelTaskPool"],
 }
 
 
@@ -23,13 +24,14 @@ if TYPE_CHECKING:
     from .base_pool import BaseTaskPool
     from .clip_vision import (
         CLIPVisionModelTaskPool,
-        RankoneWEMoECLIPVisionModelTaskPool,
+        RankoneMoECLIPVisionModelTaskPool,
         SparseWEMoECLIPVisionModelTaskPool,
     )
     from .dummy import DummyTaskPool
     from .gpt2_text_classification import GPT2TextClassificationTaskPool
     from .llama import LlamaTestGenerationTaskPool
     from .nyuv2_taskpool import NYUv2TaskPool
+    from .openclip_vision import OpenCLIPVisionModelTaskPool
 
 else:
     sys.modules[__name__] = LazyImporter(

fusion_bench/taskpool/clip_vision/__init__.py

@@ -1,4 +1,5 @@
 # flake8: noqa F401
 from .clip_rankone_moe_taskpool import RankoneMoECLIPVisionModelTaskPool
+from .clip_smile_taskpool import SmileCLIPVisionModelTaskPool
 from .clip_sparse_wemoe_taskpool import SparseWEMoECLIPVisionModelTaskPool
 from .taskpool import CLIPVisionModelTaskPool