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

fusion_bench/method/opcm/opcm.py

@@ -15,7 +15,7 @@ from tqdm.auto import tqdm
 from transformers import CLIPVisionModel
 
 from fusion_bench import BaseAlgorithm, BaseModelPool
-from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.mixins import LightningFabricMixin, SimpleProfilerMixin
 from fusion_bench.taskpool import CLIPVisionModelTaskPool
 from fusion_bench.utils import instantiate
 from fusion_bench.utils.json import load_from_json, save_to_json
@@ -31,6 +31,7 @@ if TYPE_CHECKING:
 class OPCMForCLIP(
     BaseAlgorithm,
     LightningFabricMixin,
+    SimpleProfilerMixin,
 ):
     def __init__(
         self,
@@ -64,7 +65,8 @@ class OPCMForCLIP(
             L.seed_everything(self.seed)
         accelerator = self.fabric.device
 
-        pretrained_model = modelpool.load_pretrained_model()
+        with self.profile("loading model"):
+            pretrained_model = modelpool.load_pretrained_model()
 
         model_names = modelpool.model_names
         if self.shuffle_order:
@@ -83,15 +85,17 @@ class OPCMForCLIP(
             )
 
         # get the average model
-        merged_model = modelpool.load_model(model_names[0])
+        with self.profile("loading model"):
+            merged_model = modelpool.load_model(model_names[0])
 
         if self.evaluate_on_every_step:
-            self.taskpool._is_setup = False
-            self.taskpool._test_datasets = DictConfig(
-                {model_names[0]: self._test_datasets[model_names[0]]}
-            )
-            report = self.taskpool.evaluate(deepcopy(merged_model))
-            save_to_json(report, Path(self.log_dir) / "report_0.json")
+            with self.profile("evaluating model"):
+                self.taskpool._is_setup = False
+                self.taskpool._test_datasets = DictConfig(
+                    {model_names[0]: self._test_datasets[model_names[0]]}
+                )
+                report = self.taskpool.evaluate(deepcopy(merged_model))
+                save_to_json(report, Path(self.log_dir) / "report_0.json")
 
         self.avg_task_vector_norm = get_task_vector_norm(merged_model, pretrained_model)
         self.all_task_vector_norm = [self.avg_task_vector_norm]
@@ -113,90 +117,95 @@ class OPCMForCLIP(
             enumerate(model_names[1:]), desc="Processing models"
         ):
             model_idx += 1
-            task_model = modelpool.load_model(model_name)
+            with self.profile("loading model"):
+                task_model = modelpool.load_model(model_name)
 
-            self.all_task_vector_norm.append(
-                get_task_vector_norm(task_model, pretrained_model)
-            )
-            self.avg_task_vector_norm = np.mean(self.all_task_vector_norm)
-            self.fabric.log(
-                "model/task_vector_norm", self.all_task_vector_norm[-1], step=model_idx
-            )
-            self.fabric.log(
-                "model/avg_task_vector_norm", self.avg_task_vector_norm, step=model_idx
-            )
+            with self.profile("merging model"):
+                self.all_task_vector_norm.append(
+                    get_task_vector_norm(task_model, pretrained_model)
+                )
+                self.avg_task_vector_norm = np.mean(self.all_task_vector_norm)
+                self.fabric.log(
+                    "model/task_vector_norm", self.all_task_vector_norm[-1], step=model_idx
+                )
+                self.fabric.log(
+                    "model/avg_task_vector_norm", self.avg_task_vector_norm, step=model_idx
+                )
 
-            self.lambda_t = 1  # temporary value
-
-            for module_name, module in tqdm(
-                list(merged_model.named_modules()),
-                desc=f"Processing {model_name}",
-                leave=False,
-            ):
-                if not is_leaf_module(module):
-                    continue
-
-                if isinstance(module, nn.Linear):
-                    module.weight.data = self.merge_linear_weights(
-                        module.weight,
-                        pretrained_model.get_submodule(module_name).weight,
-                        task_model.get_submodule(module_name).weight,
-                        param_name=".".join([module_name, "weight"]),
-                        alpha=self.alpha,
-                        accelerator=accelerator,
-                    )
-                    if module.bias is not None:
-                        module.bias.data = self.merge_other_parameters(
-                            module.bias,
-                            pretrained_model.get_submodule(module_name).bias,
-                            task_model.get_submodule(module_name).bias,
-                            param_name=".".join([module_name, "bias"]),
+                self.lambda_t = 1  # temporary value
+
+                for module_name, module in tqdm(
+                    list(merged_model.named_modules()),
+                    desc=f"Processing {model_name}",
+                    leave=False,
+                ):
+                    if not is_leaf_module(module):
+                        continue
+
+                    if isinstance(module, nn.Linear):
+                        module.weight.data = self.merge_linear_weights(
+                            module.weight,
+                            pretrained_model.get_submodule(module_name).weight,
+                            task_model.get_submodule(module_name).weight,
+                            param_name=".".join([module_name, "weight"]),
+                            alpha=self.alpha,
                             accelerator=accelerator,
                         )
-                else:
-                    for param_name, param in module.named_parameters():
-                        param.data = self.merge_other_parameters(
-                            merged_W=param,
-                            pretrained_W=pretrained_model.get_submodule(
-                                module_name
-                            ).get_parameter(param_name),
-                            task_W=task_model.get_submodule(module_name).get_parameter(
-                                param_name
-                            ),
-                            param_name=".".join([module_name, param_name]),
-                            accelerator=accelerator,
-                        )
-
-            task_vector_norm = get_task_vector_norm(merged_model, pretrained_model)
-            self.lambda_t *= task_vector_norm / self.avg_task_vector_norm
-            for param_name, param in merged_model.named_parameters():
-                param.data = pretrained_model.get_parameter(param_name) + (
-                    param - pretrained_model.get_parameter(param_name)
-                ) * (self.avg_task_vector_norm / task_vector_norm)
-            self.fabric.log("model/lambda_t", self.lambda_t, step=model_idx)
-            self.fabric.log(
-                "empirical/lambda_t", np.sqrt(model_idx + 1), step=model_idx
-            )
-            self.previous_lambda_t = self.lambda_t
-            self.lambda_t = None
+                        if module.bias is not None:
+                            module.bias.data = self.merge_other_parameters(
+                                module.bias,
+                                pretrained_model.get_submodule(module_name).bias,
+                                task_model.get_submodule(module_name).bias,
+                                param_name=".".join([module_name, "bias"]),
+                                accelerator=accelerator,
+                            )
+                    else:
+                        for param_name, param in module.named_parameters():
+                            param.data = self.merge_other_parameters(
+                                merged_W=param,
+                                pretrained_W=pretrained_model.get_submodule(
+                                    module_name
+                                ).get_parameter(param_name),
+                                task_W=task_model.get_submodule(module_name).get_parameter(
+                                    param_name
+                                ),
+                                param_name=".".join([module_name, param_name]),
+                                accelerator=accelerator,
+                            )
+
+                task_vector_norm = get_task_vector_norm(merged_model, pretrained_model)
+                self.lambda_t *= task_vector_norm / self.avg_task_vector_norm
+                for param_name, param in merged_model.named_parameters():
+                    param.data = pretrained_model.get_parameter(param_name) + (
+                        param - pretrained_model.get_parameter(param_name)
+                    ) * (self.avg_task_vector_norm / task_vector_norm)
+                self.fabric.log("model/lambda_t", self.lambda_t, step=model_idx)
+                self.fabric.log(
+                    "empirical/lambda_t", np.sqrt(model_idx + 1), step=model_idx
+                )
+                self.previous_lambda_t = self.lambda_t
+                self.lambda_t = None
 
-            self.fabric.log(
-                "model/merged_task_vector_norm",
-                get_task_vector_norm(merged_model, pretrained_model),
-                step=model_idx,
-            )
+                self.fabric.log(
+                    "model/merged_task_vector_norm",
+                    get_task_vector_norm(merged_model, pretrained_model),
+                    step=model_idx,
+                )
 
             if self.save_on_every_step:
-                self.save_merged_model(merged_model, model_idx)
+                with self.profile("saving model"):
+                    self.save_merged_model(merged_model, model_idx)
 
             if self.evaluate_on_every_step:
-                self.taskpool._is_setup = False
-                self.taskpool._test_datasets = DictConfig(
-                    {n: self._test_datasets[n] for n in model_names[: model_idx + 1]}
-                )
-                report = self.taskpool.evaluate(deepcopy(merged_model))
-                save_to_json(report, Path(self.log_dir) / f"report_{model_idx}.json")
+                with self.profile("evaluating model"):
+                    self.taskpool._is_setup = False
+                    self.taskpool._test_datasets = DictConfig(
+                        {n: self._test_datasets[n] for n in model_names[: model_idx + 1]}
+                    )
+                    report = self.taskpool.evaluate(deepcopy(merged_model))
+                    save_to_json(report, Path(self.log_dir) / f"report_{model_idx}.json")
 
+        self.print_profile_summary()
         return merged_model
 
     def save_merged_model(self, merged_model: CLIPVisionModel, step: int):
@@ -227,7 +236,7 @@ class OPCMForCLIP(
         split_rank = (s.cumsum(dim=0) / s.sum() > alpha).float().argmax().item()
 
         projected_task_tv = u.T @ task_tv @ v
-        projected_task_tv.diag().fill_(0)
+        projected_task_tv.diagonal().fill_(0)
 
         projected_task_tv[:split_rank, :split_rank] = 0
 

fusion_bench/method/opcm/task_arithmetic.py

@@ -15,7 +15,7 @@ from tqdm.auto import tqdm
 from transformers import CLIPVisionModel
 
 from fusion_bench import BaseAlgorithm, BaseModelPool
-from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.mixins import LightningFabricMixin, SimpleProfilerMixin
 from fusion_bench.taskpool import CLIPVisionModelTaskPool
 from fusion_bench.utils.json import load_from_json, save_to_json
 from fusion_bench.utils.state_dict_arithmetic import state_dict_add, state_dict_sub
@@ -24,7 +24,11 @@ if TYPE_CHECKING:
     from torch.utils.tensorboard import SummaryWriter
 
 
-class ContinualTaskArithmeticForCLIP(BaseAlgorithm, LightningFabricMixin):
+class ContinualTaskArithmeticForCLIP(
+    BaseAlgorithm,
+    LightningFabricMixin,
+    SimpleProfilerMixin,
+):
     def __init__(
         self,
         scaling_factor: float,
@@ -79,32 +83,42 @@ class ContinualTaskArithmeticForCLIP(BaseAlgorithm, LightningFabricMixin):
         for model_idx, model_name in tqdm(
             enumerate(model_names), desc="Processing models"
         ):
-            task_model = modelpool.load_model(model_name)
+            with self.profile("loading model"):
+                task_model = modelpool.load_model(model_name)
 
-            for param_name, param in task_model.named_parameters():
-                if not param.requires_grad:
-                    continue
+            with self.profile("merging model"):
+                for param_name, param in task_model.named_parameters():
+                    if not param.requires_grad:
+                        continue
 
-                task_param = param
-                merged_param = merged_model.get_parameter(param_name)
-                pretrained_param = pretrained_model.get_parameter(param_name)
+                    task_param = param
+                    merged_param = merged_model.get_parameter(param_name)
+                    pretrained_param = pretrained_model.get_parameter(param_name)
 
-                new_param = merged_param + self.scaling_factor * (
-                    task_param - pretrained_param
-                )
-                merged_model.get_parameter(param_name).data = new_param
+                    new_param = merged_param + self.scaling_factor * (
+                        task_param - pretrained_param
+                    )
+                    merged_model.get_parameter(param_name).data = new_param
 
             if self.save_on_every_step:
-                self.save_merged_model(merged_model, model_idx)
+                with self.profile("saving model"):
+                    self.save_merged_model(merged_model, model_idx)
 
             if self.evaluate_on_every_step:
-                self.taskpool._is_setup = False
-                self.taskpool._test_datasets = DictConfig(
-                    {n: self._test_datasets[n] for n in model_names[: model_idx + 1]}
-                )
-                report = self.taskpool.evaluate(deepcopy(merged_model))
-                save_to_json(report, Path(self.log_dir) / f"report_{model_idx}.json")
-
+                with self.profile("evaluating model"):
+                    self.taskpool._is_setup = False
+                    self.taskpool._test_datasets = DictConfig(
+                        {
+                            n: self._test_datasets[n]
+                            for n in model_names[: model_idx + 1]
+                        }
+                    )
+                    report = self.taskpool.evaluate(deepcopy(merged_model))
+                    save_to_json(
+                        report, Path(self.log_dir) / f"report_{model_idx}.json"
+                    )
+
+        self.print_profile_summary()
         return merged_model
 
     def save_merged_model(self, merged_model: CLIPVisionModel, step: int):

fusion_bench/method/opcm/ties_merging.py

@@ -20,7 +20,7 @@ from fusion_bench.method.ties_merging.ties_merging_utils import (
     ties_merging,
     vector_to_state_dict,
 )
-from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.mixins import LightningFabricMixin, SimpleProfilerMixin
 from fusion_bench.taskpool import CLIPVisionModelTaskPool
 from fusion_bench.utils.json import load_from_json, save_to_json
 from fusion_bench.utils.state_dict_arithmetic import state_dict_add, state_dict_sub
@@ -29,7 +29,11 @@ if TYPE_CHECKING:
     from torch.utils.tensorboard import SummaryWriter
 
 
-class ContinualTiesMergingForCLIP(BaseAlgorithm, LightningFabricMixin):
+class ContinualTiesMergingForCLIP(
+    BaseAlgorithm,
+    LightningFabricMixin,
+    SimpleProfilerMixin,
+):
     def __init__(
         self,
         scaling_factor: float,
@@ -84,68 +88,83 @@ class ContinualTiesMergingForCLIP(BaseAlgorithm, LightningFabricMixin):
             )
 
         # get the average model
-        pretrained_model = modelpool.load_pretrained_model()
+        with self.profile("loading model"):
+            pretrained_model = modelpool.load_pretrained_model()
         merged_model = deepcopy(pretrained_model)
 
         for model_idx, model_name in tqdm(
             enumerate(model_names), desc="Processing models"
         ):
-            task_model = modelpool.load_model(model_name)
+            with self.profile("loading model"):
+                task_model = modelpool.load_model(model_name)
 
-            task_vector = state_dict_sub(
-                task_model.state_dict(),
-                pretrained_model.state_dict(),
-            )
-            if model_idx == 0:
-                # if is the first model, the merged task vector is equal to the task vector
-                ties_merging_state_dict = task_vector
-            else:
-                # if is not the first model, we need to merge the task vector with the previous merged task vector
-                merged_tv = state_dict_sub(
-                    merged_model.state_dict(),
+            with self.profile("merging model"):
+                task_vector = state_dict_sub(
+                    task_model.state_dict(),
                     pretrained_model.state_dict(),
                 )
-                tv_flat_checks = torch.vstack(
-                    [
-                        state_dict_to_vector(merged_tv, remove_keys=self.remove_keys),
-                        state_dict_to_vector(task_vector, remove_keys=self.remove_keys),
-                    ]
-                )
-                # perform the TIES merging
-                ties_merging_tv = ties_merging(
-                    tv_flat_checks,
-                    reset_thresh=self.threshold,
-                    merge_func=self.merge_func,
-                )
-                # convert the merged task vector back to a state dict
-                ties_merging_state_dict = vector_to_state_dict(
-                    ties_merging_tv,
-                    merged_model.state_dict(),
-                    remove_keys=self.remove_keys,
-                )
-
-            for param_name, param in task_model.named_parameters():
-                if not param.requires_grad:
-                    continue
-
-                merged_param = merged_model.get_parameter(param_name)
-                new_param = (
-                    merged_param
-                    + self.scaling_factor * ties_merging_state_dict[param_name]
-                )
-                merged_model.get_parameter(param_name).data = new_param
+                if model_idx == 0:
+                    # if is the first model, the merged task vector is equal to the task vector
+                    ties_merging_state_dict = task_vector
+                else:
+                    # if is not the first model, we need to merge the task vector with the previous merged task vector
+                    merged_tv = state_dict_sub(
+                        merged_model.state_dict(),
+                        pretrained_model.state_dict(),
+                    )
+                    tv_flat_checks = torch.vstack(
+                        [
+                            state_dict_to_vector(
+                                merged_tv, remove_keys=self.remove_keys
+                            ),
+                            state_dict_to_vector(
+                                task_vector, remove_keys=self.remove_keys
+                            ),
+                        ]
+                    )
+                    # perform the TIES merging
+                    ties_merging_tv = ties_merging(
+                        tv_flat_checks,
+                        reset_thresh=self.threshold,
+                        merge_func=self.merge_func,
+                    )
+                    # convert the merged task vector back to a state dict
+                    ties_merging_state_dict = vector_to_state_dict(
+                        ties_merging_tv,
+                        merged_model.state_dict(),
+                        remove_keys=self.remove_keys,
+                    )
+
+                for param_name, param in task_model.named_parameters():
+                    if not param.requires_grad:
+                        continue
+
+                    merged_param = merged_model.get_parameter(param_name)
+                    new_param = (
+                        merged_param
+                        + self.scaling_factor * ties_merging_state_dict[param_name]
+                    )
+                    merged_model.get_parameter(param_name).data = new_param
 
             if self.save_on_every_step:
-                self.save_merged_model(merged_model, model_idx)
+                with self.profile("saving model"):
+                    self.save_merged_model(merged_model, model_idx)
 
             if self.evaluate_on_every_step:
-                self.taskpool._is_setup = False
-                self.taskpool._test_datasets = DictConfig(
-                    {n: self._test_datasets[n] for n in model_names[: model_idx + 1]}
-                )
-                report = self.taskpool.evaluate(deepcopy(merged_model))
-                save_to_json(report, Path(self.log_dir) / f"report_{model_idx}.json")
-
+                with self.profile("evaluating model"):
+                    self.taskpool._is_setup = False
+                    self.taskpool._test_datasets = DictConfig(
+                        {
+                            n: self._test_datasets[n]
+                            for n in model_names[: model_idx + 1]
+                        }
+                    )
+                    report = self.taskpool.evaluate(deepcopy(merged_model))
+                    save_to_json(
+                        report, Path(self.log_dir) / f"report_{model_idx}.json"
+                    )
+
+        self.print_profile_summary()
         return merged_model
 
     def save_merged_model(self, merged_model: CLIPVisionModel, step: int):

fusion_bench/method/task_singular_vector/TSVM.py

@@ -9,19 +9,19 @@ fusion_bench \
 ```
 """
 
-from typing import List, Optional, Union, Iterable
+from typing import Iterable, List, Optional, Union
 
 import torch
-from torch import Tensor, nn
 from omegaconf import ListConfig
+from torch import Tensor, nn
 
 from fusion_bench import BaseAlgorithm
 from fusion_bench.mixins import LightningFabricMixin
 from fusion_bench.utils import timeit_context
 from fusion_bench.utils.state_dict_arithmetic import (
     state_dict_add,
-    state_dict_sub,
     state_dict_mul,
+    state_dict_sub,
 )
 from fusion_bench.utils.type import StateDictType
 

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,