fusion-bench 0.2.31__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/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/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)