fusion-bench 0.2.30__tar.gz → 0.2.32__tar.gz

{fusion_bench-0.2.30 → fusion_bench-0.2.32}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fusion-bench
-Version: 0.2.30
+Version: 0.2.32
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
 Project-URL: Repository, https://github.com/tanganke/fusion_bench
@@ -61,6 +61,8 @@ Dynamic: license-file
 
 FusionBench is a benchmark suite designed to evaluate the performance of various deep model fusion techniques. It aims to provide a comprehensive comparison of different methods on a variety of datasets and tasks.
 
+## :newspaper: News and Related
+
 Projects based on FusionBench and news from the community (descending order of date. If you have any work based on FusionBench, please feel free to let us know, we are willing to add it to the list. :partying_face:):
 
 <details>

{fusion_bench-0.2.30 → fusion_bench-0.2.32}/README.md

@@ -21,6 +21,8 @@
 
 FusionBench is a benchmark suite designed to evaluate the performance of various deep model fusion techniques. It aims to provide a comprehensive comparison of different methods on a variety of datasets and tasks.
 
+## :newspaper: News and Related
+
 Projects based on FusionBench and news from the community (descending order of date. If you have any work based on FusionBench, please feel free to let us know, we are willing to add it to the list. :partying_face:):
 
 <details>

{fusion_bench-0.2.30 → fusion_bench-0.2.32}/fusion_bench/__init__.py

@@ -86,6 +86,9 @@ _import_structure = {
         "set_print_function_call",
         "set_print_function_call_permeanent",
         "timeit_context",
+        "initialize_hydra_config",
+        "get_default_config_path",
+        "get_hydra_output_dir",
     ],
 }
 
@@ -144,8 +147,11 @@ if TYPE_CHECKING:
         StateDictType,
         TorchModelType,
         cache_with_joblib,
+        get_default_config_path,
+        get_hydra_output_dir,
         get_rankzero_logger,
         import_object,
+        initialize_hydra_config,
         instantiate,
         parse_dtype,
         print_parameters,

fusion_bench-0.2.32/fusion_bench/__main__.py

@@ -0,0 +1,4 @@
+from fusion_bench.scripts.cli import _hydra_main
+
+if __name__ == "__main__":
+    _hydra_main()

{fusion_bench-0.2.30 → fusion_bench-0.2.32}/fusion_bench/constants/runtime.py

@@ -89,7 +89,10 @@ class RuntimeConstants:
             self._initialized = True
 
     debug = False
-    """Global debug flag for enabling verbose logging and debugging features."""
+    """
+    Global debug flag for enabling verbose logging and debugging features.
+    Use `RuntimeConstants().debug` instead of `RuntimeConstants.debug`
+    """
 
     @property
     def cache_dir(self) -> Path:

{fusion_bench-0.2.30 → fusion_bench-0.2.32}/fusion_bench/dataset/__init__.py

@@ -38,10 +38,12 @@ _extra_objects = {
 }
 _import_structure = {
     "clip_dataset": ["CLIPDataset"],
+    "image_dataset": ["ImageClassificationDataset"],
 }
 
 if TYPE_CHECKING:
     from .clip_dataset import CLIPDataset
+    from .image_dataset import ImageClassificationDataset
 
 else:
     sys.modules[__name__] = LazyImporter(

fusion_bench-0.2.32/fusion_bench/dataset/clip_dataset.py

@@ -0,0 +1,13 @@
+"""
+This module provides a class to convert a dataset whose object is a list of dictionaries with keys "image" and "label" to a dataset whose object is a tuple of tensors (inputs, label) for CLIP models.
+"""
+
+from fusion_bench.utils import DeprecationWarningMeta
+
+from .image_dataset import ImageClassificationDataset
+
+__all__ = ["CLIPDataset"]
+
+
+class CLIPDataset(ImageClassificationDataset, metaclass=DeprecationWarningMeta):
+    pass

fusion_bench-0.2.30/fusion_bench/dataset/clip_dataset.py → fusion_bench-0.2.32/fusion_bench/dataset/image_dataset.py

@@ -1,35 +1,31 @@
-"""
-This module provides a class to convert a dataset whose object is a list of dictionaries with keys "image" and "label" to a dataset whose object is a tuple of tensors (inputs, label) for CLIP models.
-"""
-
-from typing import Optional, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, Union
 
 import torch
 from torch.utils.data import Dataset
-from transformers import BaseImageProcessor, CLIPProcessor, ProcessorMixin
-
-__all__ = ["CLIPDataset"]
+from transformers import BaseImageProcessor, ProcessorMixin
 
 
-class CLIPDataset(torch.utils.data.Dataset):
+class ImageClassificationDataset(Dataset):
     """
-    A dataset class for CLIP models that converts a dataset of dictionaries or tuples
-    into a format suitable for CLIP processing.
+    A dataset class for image classification models that converts a dataset of dictionaries or tuples
+    into a format suitable for model processing.
 
-    This class wraps an existing dataset and applies CLIP preprocessing to the images.
+    This class wraps an existing dataset and applies preprocessing to the images.
     It expects each item in the dataset to be either a dictionary with 'image' and 'label' keys,
     or a tuple/list of (image, label).
-
-    Args:
-        dataset: The original dataset to wrap.
-        processor (CLIPProcessor): The CLIP processor for preparing inputs. If None, no preprocessing is applied and raw images are returned.
-
-    Attributes:
-        dataset: The wrapped dataset.
-        processor (CLIPProcessor): The CLIP processor used for image preprocessing.
     """
 
-    def __init__(self, dataset: Dataset, processor: Optional[CLIPProcessor] = None):
+    def __init__(
+        self,
+        dataset: Dataset,
+        processor: Optional[Union["ProcessorMixin", "BaseImageProcessor"]] = None,
+    ):
+        """
+        Args:
+            dataset (Dataset): The original dataset to wrap.
+            processor (Optional[Union[ProcessorMixin, BaseImageProcessor]]): The processor for preparing inputs.
+                If None, no preprocessing is applied and raw images are returned.
+        """
         self.dataset = dataset
         self.processor = processor
 
@@ -50,6 +46,8 @@ class CLIPDataset(torch.utils.data.Dataset):
         Raises:
             ValueError: If the item is neither a dictionary nor a tuple/list of length 2.
         """
+        # Standardize the item to a dictionary format
+        # {"image": ..., "label": ...}
         item = self.dataset[idx]
         if isinstance(item, dict):
             item = item
@@ -58,6 +56,8 @@ class CLIPDataset(torch.utils.data.Dataset):
             item = {"image": item[0], "label": item[1]}
         else:
             raise ValueError("Each item should be a dictionary or a tuple of length 2")
+
+        # Process the image using the provided processor, if any
         image = item["image"]
         if self.processor is not None:
             if isinstance(self.processor, (ProcessorMixin, BaseImageProcessor)):
@@ -70,7 +70,7 @@ class CLIPDataset(torch.utils.data.Dataset):
                 inputs = self.processor(image)
             else:
                 raise ValueError(
-                    "The processor should be a CLIPProcessor or a callable function"
+                    "The processor should be a transformers Processor or a callable function"
                 )
         else:
             # if processor is None, return the raw image directly

{fusion_bench-0.2.30 → fusion_bench-0.2.32}/fusion_bench/method/base_algorithm.py

@@ -59,6 +59,10 @@ class BaseAlgorithm(BaseYAMLSerializable):
     core fusion logic in the `run` method, while optional lifecycle hooks allow for
     setup and cleanup operations.
 
+    If model has `_fusion_bench_target_modules` attribute, the algorithm will only fuse
+    the specified target modules. This is useful for models where only certain layers
+    should be fused (e.g., classification heads on top of a shared backbone are not merged).
+
     Attributes:
         _program: Optional program reference for algorithm execution context.
         _config_key (str): Configuration key used for YAML serialization, defaults to "method".

{fusion_bench-0.2.30 → fusion_bench-0.2.32}/fusion_bench/method/classification/image_classification_finetune.py

@@ -173,6 +173,7 @@ class ImageClassificationFineTuning(BaseAlgorithm):
                 ),
             },
         )
+        lit_module.train()
 
         log_dir = (
             self._program.path.log_dir

fusion_bench-0.2.32/fusion_bench/method/concrete_subspace/clip_concrete_tsvm.py

@@ -0,0 +1,285 @@
+import logging
+import os
+from copy import deepcopy
+from typing import TYPE_CHECKING, Dict, Iterable, List, Literal, Optional
+
+import torch
+from omegaconf import DictConfig
+from tqdm import tqdm
+
+from fusion_bench import (
+    BaseAlgorithm,
+    OpenCLIPClassificationMixin,
+    OpenCLIPVisionModelPool,
+    SimpleProfilerMixin,
+    StateDictType,
+    auto_register_config,
+    get_rankzero_logger,
+    instantiate,
+)
+from fusion_bench.method.adamerging.entropy_loss import entropy_loss
+from fusion_bench.method.task_singular_vector import TaskSingularVectorMerging
+from fusion_bench.method.task_singular_vector.utils import (
+    TSVM_utils,
+    check_parameterNamesMatch,
+    check_state_dicts_equal,
+    state_dict_to_vector,
+    vector_to_state_dict,
+)
+from fusion_bench.models.masks import MaskModel, mask_sparsity
+from fusion_bench.models.open_clip import (
+    ClassificationHead,
+    ImageClassifier,
+    ImageEncoder,
+)
+from fusion_bench.models.wrappers.task_wise_fusion import (
+    TaskWiseMergedModel,
+    get_task_wise_weights,
+)
+from fusion_bench.utils.devices import clear_cuda_cache
+from fusion_bench.utils.dtype import parse_dtype
+from fusion_bench.utils.parameters import print_parameters, print_trainable_parameters
+from fusion_bench.utils.rich_utils import print_config_yaml
+from fusion_bench.utils.state_dict_arithmetic import (
+    _validate_state_dict_same_keys,
+    state_dict_add,
+    state_dict_hadamard_product,
+    state_dict_mul,
+    state_dict_sub,
+)
+
+log = get_rankzero_logger(__name__)
+
+
+@auto_register_config
+class ConcreteTSVMForOpenCLIP(
+    OpenCLIPClassificationMixin,
+    SimpleProfilerMixin,
+    BaseAlgorithm,
+):
+    def __init__(
+        self,
+        dataloader_kwargs: DictConfig,
+        optimizer: DictConfig,
+        lr_scheduler: DictConfig,
+        max_steps: int,
+        save_interval: int,
+        initial_logits: float,
+        temperature: float,
+        eval_mask_type: Literal["continuous", "discrete"],
+        mask_checkpoint: Optional[str],
+        merge_dtype: str,
+        clamp_weights: bool,
+        tie_weights: bool,
+        strict: bool,
+        skip_training: bool,
+        # === TSVM parameters ===
+        exclude_keys: Optional[List[str]],
+        alpha: float,
+        return_single_task_models: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        if not return_single_task_models:
+            log.warning("return_single_task_models is forced to be True here.")
+            self.return_single_task_models = True
+
+    @torch.no_grad()
+    def setup_models(self):
+        """
+        load the pre-trained model, task vectors, and construct the mask model.
+        """
+        merge_dtype = parse_dtype(self.merge_dtype)
+        modelpool = self.modelpool
+
+        # load the pre-trained model
+        pretrained_model = modelpool.load_pretrained_model()
+        self.set_clip_processor(stage="test", processor=pretrained_model.val_preprocess)
+
+        # constrcute mask model
+        mask_model = MaskModel(
+            pretrained_model, ignore_untrained_params=True, parameter_type="logits"
+        )
+        if merge_dtype is not None:
+            mask_model.to(merge_dtype)
+        mask_model.fill_(self.initial_logits)
+
+        if self.fabric.is_global_zero:
+            print("summary of mask model:")
+            print_parameters(mask_model)
+
+        if self.fabric.is_global_zero:
+            tsvm_algo = TaskSingularVectorMerging(
+                alpha=self.alpha,
+                exclude_keys=self.exclude_keys,
+                return_single_task_models=self.return_single_task_models,
+            )
+            tsvm_algo._fabric_instance = self.fabric
+            models = tsvm_algo.run(modelpool)
+
+            finetuned_models = [models[name] for name in modelpool.model_names]
+
+            task_wise_weight = get_task_wise_weights(
+                num_models=len(modelpool.model_names),
+                init_values=self.alpha,
+            )
+
+            # create a wrapped model
+            module = TaskWiseMergedModel(
+                task_wise_weight=task_wise_weight,
+                pretrained_model=pretrained_model,
+                finetuned_models=finetuned_models,
+                clamp_weights=self.clamp_weights,
+                tie_weights=self.tie_weights,
+                strict=self.strict,
+                task_vector_dtype=merge_dtype,
+            )
+            module = module.to(dtype=merge_dtype)
+
+            print("trainable parameter summary of merged model (TaskWiseMergedModel):")
+            print_trainable_parameters(module)
+        else:
+            module = None
+
+        with torch.no_grad():
+            self.fabric.barrier()
+            module = self.fabric.broadcast(module, src=0)
+
+        return module, mask_model
+
+    def train_mask(self, module: TaskWiseMergedModel, mask_model: MaskModel):
+        """
+        Train the mask model using the provided module.
+
+        This method configures the optimizer, sets up the mask model, and performs test-time adaptation to train the mask model.
+
+        Args:
+            module (TaskWiseMergedModel): The wrapped model with task-wise weights.
+            mask_model (MaskModel): The mask model to be trained.
+        """
+        config = self.config
+        merge_dtype = parse_dtype(self.merge_dtype)
+        log.info(f"Using merge dtype: {merge_dtype}")
+
+        optimizer: "torch.optim.Optimizer" = instantiate(
+            self.optimizer,
+            params=filter(lambda p: p.requires_grad, mask_model.parameters()),
+        )
+        print(f"{optimizer=}")
+        if self.lr_scheduler is not None:
+            lr_scheduler = instantiate(
+                self.lr_scheduler,
+                optimizer=optimizer,
+            )
+            print(f"{lr_scheduler=}")
+        else:
+            lr_scheduler = None
+
+        log.info("Setup models and optimizer with Fabric.")
+        mask_model, optimizer = self.fabric.setup(mask_model, optimizer)
+
+        log.info("Move the merged module to the correct device and disable gradients.")
+        module.requires_grad_(False)
+        module.to(mask_model.device)
+
+        mask_model.train()
+        optimizer.zero_grad()
+        for step_idx in (
+            pbar := tqdm(
+                range(self.config.max_steps if not self.is_debug_mode else 5),
+                ("[DEBUG MODE] " if self.is_debug_mode else "")
+                + "Concrete TSVM Test-Time Adaptation",
+                dynamic_ncols=True,
+                disable=not self.fabric.is_global_zero,
+            )
+        ):
+            metrics = {}
+            # sample a shared mask and merge weights
+            with self.profile("sample mask"):
+                mask = mask_model.sample_mask(
+                    mask_type="continuous", temperature=config.temperature
+                )
+                metrics["train/sparsity"] = mask_sparsity(mask)
+            with self.profile("merge weights"):
+                # rescale mask
+                for name, m in mask.items():
+                    mask[name] = m / torch.mean(m)
+                module.merge_weights(task_vector_mask=mask)
+
+            # ------ inner optimization goes here ------
+            # NOTE:
+            #   Because the algorithmic parameters of TSVM are assumed to be chosen on a validation test
+            #   set, we do not need to perform inner optimization here. So here we skip the inner optimization step.
+            # ------------------------------------------
+
+            total_loss = None
+            for task in self.modelpool.model_names:
+                with self.profile("data loading"):
+                    batch = next(self.get_shuffled_test_loader_iter(task))
+                    # NOTE: The labels are not allowed to be used during test-time adaptation
+                    images = batch[0].to(dtype=merge_dtype)
+                with self.profile("forward pass"):
+                    logits = self.compute_logits(module, images, task)
+                    loss = entropy_loss(logits)
+                    total_loss = loss if total_loss is None else total_loss + loss
+
+            with self.profile("compute grad"):
+                self.fabric.backward(total_loss)
+
+            with self.profile("optimizer step"):
+                optimizer.step()
+                optimizer.zero_grad()
+
+                if lr_scheduler is not None:
+                    lr_scheduler.step()
+
+            metrics.update({"train/loss": loss.item()})
+            self.fabric.log_dict(metrics, step=step_idx)
+            pbar.set_postfix(metrics)
+
+            if (step_idx + 1) % self.config.save_interval == 0:
+                with self.profiler.profile("save checkpoint"):
+                    save_dir = os.path.join(self.fabric.logger.log_dir, "checkpoints")
+                    if not os.path.exists(save_dir):
+                        os.makedirs(save_dir, exist_ok=True)
+                    save_path = os.path.join(save_dir, f"mask_steps_{step_idx}.pt")
+                    print(f"saving checkpoint to {save_path}")
+                    state = {"model": mask_model}
+                    self.fabric.save(save_path, state)
+
+                    # Create or update a symbolic link to the latest checkpoint
+                    if self.fabric.is_global_zero:
+                        symlink_path = os.path.join(save_dir, "latest_checkpoint.pt")
+                        if os.path.exists(symlink_path):
+                            os.remove(symlink_path)
+                        os.link(os.path.abspath(save_path), symlink_path)
+
+                self.print_profile_summary()
+
+    def run(self, modelpool: OpenCLIPVisionModelPool):
+        self.modelpool = modelpool
+        merge_dtype = parse_dtype(self.merge_dtype)
+
+        with self.profile("setup models"):
+            module, mask_model = self.setup_models()
+            self.setup_zero_shot_classification_head(freeze=True, dtype=merge_dtype)
+
+        if self.mask_checkpoint is None:
+            if not self.skip_training:
+                clear_cuda_cache()
+                self.train_mask(module, mask_model=mask_model)
+        else:
+            if self.fabric.is_global_zero:
+                print("loading mask from checkpoint", self.mask_checkpoint)
+            self.fabric.load(self.mask_checkpoint, {"model": mask_model})
+
+        with torch.no_grad():
+            clear_cuda_cache()
+            mask = mask_model.sample_mask(
+                mask_type=self.eval_mask_type, temperature=self.temperature
+            )
+            # rescale mask
+            for name, m in mask.items():
+                mask[name] = m / torch.mean(m)
+            model = module.merge_and_unload(mask)
+        return model.to(dtype=torch.float32)

{fusion_bench-0.2.30 → fusion_bench-0.2.32}/fusion_bench/method/dop/dop.py

@@ -79,28 +79,6 @@ class ContinualDOPForCLIP(BaseAlgorithm, LightningFabricMixin):
         ), "The alpha should be in the range of [0, 1]"
         super().__init__(**kwargs)
 
-    def print_params(self, pretrained_model):
-        total_params = 0
-        linear_params = 0
-        linear_weight_params = 0
-        for module_name, module in pretrained_model.named_modules():
-            if not is_leaf_module(module):
-                continue
-            if isinstance(module, nn.Linear):
-                linear_params += sum(p.numel() for n, p in module.named_parameters())
-                linear_weight_params += sum(
-                    p.numel() for n, p in module.named_parameters() if "weight" in n
-                )
-            total_params += sum(p.numel() for p in module.parameters())
-
-        linear_ratio = linear_params / total_params * 100
-        linear_weight_ratio = linear_weight_params / total_params * 100
-        print(f"Total Parameters: {total_params}")
-        print(f"Linear Parameters: {linear_params}")
-        print(f"Linear Weight Parameters: {linear_weight_params}")
-        print(f"Linear Ratio: {linear_ratio:.2f}%")
-        print(f"Linear Weight Ratio: {linear_weight_ratio:.2f}%")
-
     def run(self, modelpool: BaseModelPool):
         if self.seed is not None:
             L.seed_everything(self.seed)