fusion-bench 0.2.30__py3-none-any.whl → 0.2.32__py3-none-any.whl

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/__main__.py

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

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/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/dataset/clip_dataset.py

@@ -2,80 +2,12 @@
 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 fusion_bench.utils import DeprecationWarningMeta
 
-import torch
-from torch.utils.data import Dataset
-from transformers import BaseImageProcessor, CLIPProcessor, ProcessorMixin
+from .image_dataset import ImageClassificationDataset
 
 __all__ = ["CLIPDataset"]
 
 
-class CLIPDataset(torch.utils.data.Dataset):
-    """
-    A dataset class for CLIP models that converts a dataset of dictionaries or tuples
-    into a format suitable for CLIP processing.
-
-    This class wraps an existing dataset and applies CLIP 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):
-        self.dataset = dataset
-        self.processor = processor
-
-    def __len__(self):
-        """Returns the number of items in the dataset."""
-        return len(self.dataset)
-
-    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
-        """
-        Retrieves and processes an item from the dataset.
-
-        Args:
-            idx (int): The index of the item to retrieve.
-
-        Returns:
-            tuple: A tuple containing the processed image tensor and the label.
-
-        Raises:
-            ValueError: If the item is neither a dictionary nor a tuple/list of length 2.
-        """
-        item = self.dataset[idx]
-        if isinstance(item, dict):
-            item = item
-        elif isinstance(item, (tuple, list)):
-            assert len(item) == 2, "Each item should be a tuple or list of length 2"
-            item = {"image": item[0], "label": item[1]}
-        else:
-            raise ValueError("Each item should be a dictionary or a tuple of length 2")
-        image = item["image"]
-        if self.processor is not None:
-            if isinstance(self.processor, (ProcessorMixin, BaseImageProcessor)):
-                # Apply the processor to the image to get the input tensor
-                image = image.convert("RGB")  # ensure image is in RGB format
-                inputs = self.processor(images=[image], return_tensors="pt")[
-                    "pixel_values"
-                ][0]
-            elif callable(self.processor):
-                inputs = self.processor(image)
-            else:
-                raise ValueError(
-                    "The processor should be a CLIPProcessor or a callable function"
-                )
-        else:
-            # if processor is None, return the raw image directly
-            inputs = image
-        # convert boolean label to int, this is for the case when the label is a binary classification task
-        if isinstance(item["label"], bool):
-            item["label"] = 1 if item["label"] else 0
-        return inputs, item["label"]
+class CLIPDataset(ImageClassificationDataset, metaclass=DeprecationWarningMeta):
+    pass

fusion_bench/dataset/image_dataset.py

@@ -1,35 +1,39 @@
-from typing import Any, Callable, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, Union
 
+import torch
 from torch.utils.data import Dataset
+from transformers import BaseImageProcessor, ProcessorMixin
 
 
-class TransformedImageDataset(Dataset):
+class ImageClassificationDataset(Dataset):
     """
-    A dataset class for image classification tasks that applies a transform to images.
+    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 a specified transform 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.
-        transform (Callable): A function/transform to apply on the image.
-
-    Attributes:
-        dataset: The wrapped dataset.
-        transform (Callable): The transform to be applied to the images.
     """
 
-    def __init__(self, dataset: Dataset, transform: Callable):
-        super().__init__()
+    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.transform = transform
+        self.processor = processor
 
     def __len__(self):
         """Returns the number of items in the dataset."""
         return len(self.dataset)
 
-    def __getitem__(self, idx: int) -> Tuple[Any, Any]:
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
         """
         Retrieves and processes an item from the dataset.
 
@@ -37,11 +41,13 @@ class TransformedImageDataset(Dataset):
             idx (int): The index of the item to retrieve.
 
         Returns:
-            tuple: A tuple containing the processed image and the label.
+            tuple: A tuple containing the processed image tensor and the label.
 
         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
@@ -50,6 +56,26 @@ class TransformedImageDataset(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"]
-        inputs = self.transform(image)
+        if self.processor is not None:
+            if isinstance(self.processor, (ProcessorMixin, BaseImageProcessor)):
+                # Apply the processor to the image to get the input tensor
+                image = image.convert("RGB")  # ensure image is in RGB format
+                inputs = self.processor(images=[image], return_tensors="pt")[
+                    "pixel_values"
+                ][0]
+            elif callable(self.processor):
+                inputs = self.processor(image)
+            else:
+                raise ValueError(
+                    "The processor should be a transformers Processor or a callable function"
+                )
+        else:
+            # if processor is None, return the raw image directly
+            inputs = image
+        # convert boolean label to int, this is for the case when the label is a binary classification task
+        if isinstance(item["label"], bool):
+            item["label"] = 1 if item["label"] else 0
         return inputs, item["label"]

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