fusion-bench 0.2.20__py3-none-any.whl → 0.2.22__py3-none-any.whl

fusion_bench/scripts/cli.py

@@ -12,6 +12,7 @@ import os
 import hydra
 from omegaconf import DictConfig, OmegaConf
 
+from fusion_bench.constants import PROJECT_ROOT_PATH
 from fusion_bench.programs import BaseHydraProgram
 from fusion_bench.utils import instantiate
 
@@ -20,11 +21,10 @@ log = logging.getLogger(__name__)
 
 def _get_default_config_path():
     for config_dir in ["fusion_bench_config", "config"]:
-        config_path = os.path.join(
-            importlib.import_module("fusion_bench").__path__[0], "..", config_dir
-        )
-        if os.path.exists(config_path) and os.path.isdir(config_path):
-            return os.path.abspath(config_path)
+        for config_path_root in [os.getcwd(), PROJECT_ROOT_PATH]:
+            config_path = os.path.join(config_path_root, config_dir)
+            if os.path.exists(config_path) and os.path.isdir(config_path):
+                return os.path.abspath(config_path)
     return None
 
 
@@ -34,6 +34,39 @@ def _get_default_config_path():
     version_base=None,
 )
 def main(cfg: DictConfig) -> None:
+    """
+    Main entry point for the FusionBench command-line interface.
+
+    This function serves as the primary entry point for the `fusion_bench` CLI command.
+    It is decorated with Hydra's main decorator to handle configuration management,
+    command-line argument parsing, and configuration file loading.
+
+    The function performs the following operations:
+    1. Resolves any interpolations in the configuration using OmegaConf
+    2. Instantiates the appropriate program class based on the configuration
+    3. Executes the program's run method to perform the fusion task
+
+    Args:
+        cfg (DictConfig): The Hydra configuration object containing all settings
+            for the fusion task. This includes method configuration, model pool
+            configuration, task pool configuration, and other runtime parameters.
+            The configuration is automatically loaded by Hydra from the specified
+            config files and command-line overrides.
+
+    Returns:
+        None: This function doesn't return a value but executes the fusion
+            program which may save results, log outputs, or perform other
+            side effects as configured.
+
+    Example:
+        This function is typically called automatically when running:
+        ```bash
+        fusion_bench method=... modelpool=... taskpool=...
+        ```
+
+        The Hydra decorator handles parsing these command-line arguments and
+        loading the corresponding configuration files to populate the cfg parameter.
+    """
     OmegaConf.resolve(cfg)
     program: BaseHydraProgram = instantiate(cfg)
     program.run()

fusion_bench/taskpool/base_pool.py

@@ -1,14 +1,15 @@
 from abc import abstractmethod
+from typing import Any, Dict
 
-from fusion_bench.mixins import BaseYAMLSerializableModel
+from fusion_bench.mixins import BaseYAMLSerializable
 
 
-class BaseTaskPool(BaseYAMLSerializableModel):
+class BaseTaskPool(BaseYAMLSerializable):
     _program = None
     _config_key = "taskpool"
 
     @abstractmethod
-    def evaluate(self, model, *args, **kwargs):
+    def evaluate(self, model: Any, *args: Any, **kwargs: Any) -> Dict[str, Any]:
         """
         Evaluate the model on all tasks in the task pool, and return a report.
 

fusion_bench/taskpool/clip_vision/taskpool.py

@@ -27,8 +27,9 @@ from tqdm.autonotebook import tqdm
 from transformers import CLIPModel, CLIPProcessor, CLIPVisionModel
 from transformers.models.clip.modeling_clip import CLIPVisionTransformer
 
+from fusion_bench import RuntimeConstants
 from fusion_bench.dataset import CLIPDataset
-from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.mixins import HydraConfigMixin, LightningFabricMixin
 from fusion_bench.models.hf_clip import HFCLIPClassifier
 from fusion_bench.taskpool import BaseTaskPool
 from fusion_bench.tasks.clip_classification import get_classnames_and_templates
@@ -56,6 +57,8 @@ class LayerWiseFeatureSaver:
         first_token_only: bool = True,
         max_num: Optional[int] = None,
     ):
+        if isinstance(save_path, str):
+            save_path = Path(save_path)
         self.save_path = save_path
         self.first_token_only = first_token_only
         self.max_num = max_num
@@ -84,8 +87,9 @@ class LayerWiseFeatureSaver:
 
 
 class CLIPVisionModelTaskPool(
-    BaseTaskPool,
+    HydraConfigMixin,
     LightningFabricMixin,
+    BaseTaskPool,
 ):
     """
     This class is used to define the image classification task for CLIP models.
@@ -122,14 +126,14 @@ class CLIPVisionModelTaskPool(
         self,
         test_datasets: Union[DictConfig, Dict[str, Dataset]],
         *,
-        processor: Union[DictConfig, CLIPProcessor],
-        data_processor: Union[DictConfig, CLIPProcessor],
-        clip_model: Union[DictConfig, CLIPModel],
+        processor: Union[str, DictConfig, CLIPProcessor],
+        clip_model: Union[str, DictConfig, CLIPModel],
+        data_processor: Union[DictConfig, CLIPProcessor] = None,
         dataloader_kwargs: DictConfig = None,
         layer_wise_feature_save_path: Optional[str] = None,
         layer_wise_feature_first_token_only: bool = True,
         layer_wise_feature_max_num: Optional[int] = None,
-        fast_dev_run: bool = False,
+        fast_dev_run: Optional[bool] = None,
         **kwargs,
     ):
         """
@@ -151,7 +155,10 @@ class CLIPVisionModelTaskPool(
         self.layer_wise_feature_first_token_only = layer_wise_feature_first_token_only
         self.layer_wise_feature_max_num = layer_wise_feature_max_num
 
-        self.fast_dev_run = fast_dev_run
+        if self.fast_dev_run is None:
+            self.fast_dev_run = RuntimeConstants().debug
+        else:
+            self.fast_dev_run = fast_dev_run
         super().__init__(**kwargs)
 
     def setup(self):
@@ -159,21 +166,35 @@ class CLIPVisionModelTaskPool(
         Set up the processor, data processor, CLIP model, test datasets, and data loaders.
         """
         # setup processor and clip model
-        self.processor = (
-            instantiate(self._processor)
-            if isinstance(self._processor, DictConfig)
-            else self._processor
-        )
-        self.data_processor = (
-            instantiate(self._data_processor)
-            if isinstance(self._data_processor, DictConfig)
-            else self._data_processor
-        )
-        self.clip_model = (
-            instantiate(self._clip_model)
-            if isinstance(self._clip_model, DictConfig)
-            else self._clip_model
-        )
+        if isinstance(self._processor, str):
+            self.processor = CLIPProcessor.from_pretrained(self._processor)
+        elif (
+            isinstance(self._processor, (dict, DictConfig))
+            and "_target_" in self._processor
+        ):
+            self.processor = instantiate(self._processor)
+        else:
+            self.processor = self._processor
+
+        if self._data_processor is None:
+            self.data_processor = self.processor
+        else:
+            self.data_processor = (
+                instantiate(self._data_processor)
+                if isinstance(self._data_processor, DictConfig)
+                else self._data_processor
+            )
+
+        if isinstance(self._clip_model, str):
+            self.clip_model = CLIPModel.from_pretrained(self._clip_model)
+        elif (
+            isinstance(self._clip_model, (dict, DictConfig))
+            and "_target_" in self._clip_model
+        ):
+            self.clip_model = instantiate(self._clip_model)
+        else:
+            self.clip_model = self._clip_model
+
         self.clip_model = self.fabric.to_device(self.clip_model)
         self.clip_model.requires_grad_(False)
         self.clip_model.eval()

fusion_bench/taskpool/dummy.py

@@ -4,13 +4,13 @@ This is the dummy task pool that is used for debugging purposes.
 
 from typing import Optional
 
+from lightning.pytorch.utilities import rank_zero_only
 from torch import nn
 
 from fusion_bench.models.separate_io import separate_save
 from fusion_bench.taskpool.base_pool import BaseTaskPool
 from fusion_bench.utils import timeit_context
 from fusion_bench.utils.parameters import count_parameters, print_parameters
-from lightning.pytorch.utilities import rank_zero_only
 
 
 def get_model_summary(model: nn.Module) -> dict:

fusion_bench/taskpool/lm_eval_harness/taskpool.py

@@ -1,6 +1,6 @@
 import logging
 import os
-from typing import List, Literal, Optional, Union, TYPE_CHECKING
+from typing import TYPE_CHECKING, List, Literal, Optional, Union
 
 import lightning.fabric
 import lm_eval
@@ -12,7 +12,6 @@ from fusion_bench import BaseTaskPool
 from fusion_bench.mixins import LightningFabricMixin
 from fusion_bench.utils.strenum import _version
 
-
 log = logging.getLogger(__name__)
 
 

fusion_bench/tasks/clip_classification/__init__.py

@@ -1,6 +1,6 @@
 import importlib
 import warnings
-from typing import Any, Callable, Dict, List
+from typing import Any, Callable, Dict, List, Tuple
 
 from datasets import load_dataset
 
@@ -79,7 +79,9 @@ class CLIPTemplateFactory:
     }
 
     @staticmethod
-    def get_classnames_and_templates(dataset_name: str):
+    def get_classnames_and_templates(
+        dataset_name: str,
+    ) -> Tuple[List[str], List[Callable]]:
         """
         Retrieves class names and templates for the specified dataset.
 
@@ -169,7 +171,7 @@ class CLIPTemplateFactory:
         CLIPTemplateFactory._dataset_mapping[dataset_name] = dataset_info
 
     @staticmethod
-    def get_available_datasets():
+    def get_available_datasets() -> List[str]:
         """
         Get a list of all available dataset names.
 
@@ -179,5 +181,5 @@ class CLIPTemplateFactory:
         return list(CLIPTemplateFactory._dataset_mapping.keys())
 
 
-def get_classnames_and_templates(dataset_name: str):
+def get_classnames_and_templates(dataset_name: str) -> Tuple[List[str], List[Callable]]:
     return CLIPTemplateFactory.get_classnames_and_templates(dataset_name)

fusion_bench/utils/__init__.py

@@ -7,10 +7,16 @@ from .cache_utils import *
 from .devices import *
 from .dtype import parse_dtype
 from .fabric import seed_everything_by_time
-from .instantiate_utils import instantiate, is_instantiable
+from .instantiate_utils import (
+    instantiate,
+    is_instantiable,
+    set_print_function_call,
+    set_print_function_call_permeanent,
+)
 from .json import load_from_json, save_to_json
 from .lazy_state_dict import LazyStateDict
 from .misc import *
 from .packages import import_object
 from .parameters import *
+from .pylogger import get_rankzero_logger
 from .timer import timeit_context

fusion_bench/utils/cache_utils.py

@@ -1,15 +1,30 @@
 import logging
 import os
 import pickle
+import warnings
 from functools import wraps
 from pathlib import Path
 from typing import Any, Callable, Union
 
-__all__ = ["cache_to_disk"]
+from joblib import Memory
+
+__all__ = ["cache_to_disk", "cache_with_joblib", "set_default_cache_dir"]
 
 
 log = logging.getLogger(__name__)
 
+DEFAULT_CACHE_DIR = Path.cwd() / "outputs" / "cache"
+
+
+def set_default_cache_dir(path: str | Path):
+    global DEFAULT_CACHE_DIR
+    if path is None:
+        return
+
+    if isinstance(path, str):
+        path = Path(path)
+    DEFAULT_CACHE_DIR = path
+
 
 def cache_to_disk(file_path: Union[str, Path]) -> Callable:
     """
@@ -17,6 +32,11 @@ def cache_to_disk(file_path: Union[str, Path]) -> Callable:
     the result is loaded from the file. Otherwise, the function is executed and
     the result is saved to the file.
 
+    !!! warning "deprecated"
+        This function is deprecated. Use `cache_with_joblib` instead for better
+        caching capabilities including automatic cache invalidation, better object
+        handling, and memory efficiency.
+
     ## Example usage
 
     ```python
@@ -32,6 +52,13 @@ def cache_to_disk(file_path: Union[str, Path]) -> Callable:
     Returns:
         Callable: The decorated function.
     """
+    warnings.warn(
+        "cache_to_disk is deprecated. Use cache_with_joblib instead for better "
+        "caching capabilities including automatic cache invalidation, better object "
+        "handling, and memory efficiency.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
     if isinstance(file_path, str):
         file_path = Path(file_path)
     assert isinstance(file_path, Path)
@@ -56,3 +83,76 @@ def cache_to_disk(file_path: Union[str, Path]) -> Callable:
         return wrapper
 
     return decorator
+
+
+def cache_with_joblib(
+    cache_dir: Union[str, Path] = None,
+    verbose: int = 0,
+) -> Callable:
+    """
+    A decorator to cache the result of a function using joblib.Memory. This provides
+    more advanced caching capabilities compared to cache_to_disk, including:
+    - Automatic cache invalidation when function arguments change
+    - Better handling of numpy arrays and other complex objects
+    - Memory-efficient storage
+    - Optional verbose output for cache hits/misses
+
+    ## Example usage
+
+    ```python
+    @cache_with_joblib("./cache", verbose=1)
+    def expensive_computation(x: int, y: str) -> Any:
+        # Function implementation
+        return complex_result
+
+    # Or with default settings:
+    @cache_with_joblib()
+    def another_function(x: int) -> int:
+        return x * 2
+    ```
+
+    Args:
+        cache_dir (Union[str, Path]): The directory where cache files should be stored.
+            If `None`, a default directory `outputs/cache` will be used.
+        verbose (int): Verbosity level for joblib.Memory (0=silent, 1=basic, 2++=verbose).
+
+    Returns:
+        Callable: A decorator function that can be applied to functions.
+    """
+
+    if cache_dir is None:
+        cache_dir = DEFAULT_CACHE_DIR
+
+    if isinstance(cache_dir, str):
+        cache_dir = Path(cache_dir)
+    assert isinstance(cache_dir, Path)
+
+    # Create the cache directory if it doesn't exist
+    cache_dir.mkdir(parents=True, exist_ok=True)
+
+    # Create a Memory object for this function
+    memory = Memory(location=cache_dir, verbose=verbose)
+
+    def decorator(func: Callable) -> Callable:
+        nonlocal memory
+
+        # Create the cached version of the function
+        cached_func = memory.cache(func)
+
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            return cached_func(*args, **kwargs)
+
+        # Expose useful methods from joblib.Memory
+        if not (
+            hasattr(cached_func, "clear")
+            or hasattr(cached_func, "call")
+            or hasattr(cached_func, "check_call_in_cache")
+        ):
+            wrapper.clear = cached_func.clear
+            wrapper.call = cached_func.call
+            wrapper.check_call_in_cache = cached_func.check_call_in_cache
+
+        return wrapper
+
+    return decorator

fusion_bench/utils/devices.py

@@ -1,4 +1,5 @@
 import gc
+import logging
 import os
 from typing import List, Optional, Union
 
@@ -12,7 +13,7 @@ from transformers.utils import (
 )
 
 __all__ = [
-    "cuda_empty_cache",
+    "clear_cuda_cache",
     "to_device",
     "num_devices",
     "get_device",
@@ -21,10 +22,19 @@ __all__ = [
     "get_device_capabilities",
 ]
 
+log = logging.getLogger(__name__)
 
-def cuda_empty_cache():
+
+def clear_cuda_cache():
+    """
+    Clears the CUDA memory cache to free up GPU memory.
+    Works only if CUDA is available.
+    """
     gc.collect()
-    torch.cuda.empty_cache()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    else:
+        log.warning("CUDA is not available. No cache to clear.")
 
 
 def to_device(obj, device: Optional[torch.device], **kwargs):
@@ -75,7 +85,7 @@ def num_devices(devices: Union[int, List[int], str]) -> int:
     Return the number of devices.
 
     Args:
-        devices: `devices` can be a single int to specify the number of devices, or a list of device ids, e.g. [0, 1, 2, 3]， or a str of device ids, e.g. "0,1,2,3" and "[0, 1, 2]".
+        devices: `devices` can be a single int to specify the number of devices, or a list of device ids, e.g. [0, 1, 2, 3], or a str of device ids, e.g. "0,1,2,3" and "[0, 1, 2]".
 
     Returns:
         The number of devices.

fusion_bench/utils/fabric.py

@@ -3,9 +3,9 @@ from typing import Optional
 
 import lightning as L
 
-from fusion_bench.utils.pylogger import getRankZeroLogger
+from fusion_bench.utils.pylogger import get_rankzero_logger
 
-log = getRankZeroLogger(__name__)
+log = get_rankzero_logger(__name__)
 
 
 def seed_everything_by_time(fabric: Optional[L.Fabric] = None):

fusion_bench/utils/instantiate_utils.py

@@ -28,7 +28,7 @@ PRINT_FUNCTION_CALL_FUNC = print
 Function to be used for printing function calls.
 """
 
-CATCH_EXCEPTION = True
+CATCH_EXCEPTION = False
 
 
 @contextmanager
@@ -41,10 +41,12 @@ def set_print_function_call(value: bool):
     finally:
         PRINT_FUNCTION_CALL = old_value
 
+
 def set_print_function_call_permeanent(value: bool):
     global PRINT_FUNCTION_CALL
     PRINT_FUNCTION_CALL = value
 
+
 def is_instantiable(config: Union[DictConfig, Any]) -> bool:
     if OmegaConf.is_dict(config):
         return "_target_" in config

fusion_bench/utils/lazy_imports.py

@@ -72,3 +72,26 @@ class LazyImporter(ModuleType):
 
     def __reduce__(self):
         return (self.__class__, (self._name, self.__file__, self._import_structure))
+
+
+class LazyModule(ModuleType):
+    """Module wrapper for lazy import.
+    Adapted from Optuna: https://github.com/optuna/optuna/blob/1f92d496b0c4656645384e31539e4ee74992ff55/optuna/__init__.py
+
+    This class wraps specified module and lazily import it when they are actually accessed.
+
+    Args:
+        name: Name of module to apply lazy import.
+    """
+
+    def __init__(self, name: str) -> None:
+        super().__init__(name)
+        self._name = name
+
+    def _load(self) -> ModuleType:
+        module = importlib.import_module(self._name)
+        self.__dict__.update(module.__dict__)
+        return module
+
+    def __getattr__(self, item: str) -> Any:
+        return getattr(self._load(), item)

fusion_bench/utils/lazy_state_dict.py

@@ -2,7 +2,7 @@ import json
 import logging
 import os
 from copy import deepcopy
-from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple, Type
+from typing import TYPE_CHECKING, Dict, Iterator, List, Mapping, Optional, Tuple, Type
 
 import torch
 from accelerate import init_empty_weights
@@ -49,7 +49,7 @@ def resolve_checkpoint_path(
         )
 
 
-class LazyStateDict:
+class LazyStateDict(Mapping[str, torch.Tensor]):
     """
     Dictionary-like object that lazily loads a state dict from a checkpoint path.
     """
@@ -168,12 +168,21 @@ class LazyStateDict:
     def config(self) -> "PretrainedConfig":
         return AutoConfig.from_pretrained(self._checkpoint)
 
+    @property
+    def dtype(self) -> torch.dtype:
+        """
+        `torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype).
+        """
+        first_key = next(iter(self.keys()))
+        first_param = self[first_key]
+        return first_param.dtype
+
     def state_dict(self, keep_vars: bool = False) -> "LazyStateDict":
         """
         Args:
             keep_vars (bool): Ignored, as LazyStateDict does not support keep_vars. Just for compatibility.
         """
-        return self
+        return deepcopy(self)
 
     def _resolve_checkpoint_files(self, checkpoint: str):
         # reference: https://huggingface.co/docs/accelerate/v0.17.1/en/usage_guides/big_modeling
@@ -290,6 +299,18 @@ class LazyStateDict:
             )
             return tensor
 
+    def pop(self, key: str):
+        assert key in list(
+            self.keys()
+        ), "KeyError: Cannot pop a tensor for a key that does not exist in the LazyStateDict."
+        if self._state_dict_cache is not None and key in self._state_dict_cache:
+            if key in self._index:
+                self._index.pop(key)
+            return self._state_dict_cache.pop(key)
+        if key in self._index:
+            self._index.pop(key)
+        return None
+
     def __setitem__(self, key: str, value: torch.Tensor) -> None:
         """
         Set a tensor in the LazyStateDict. This will update the state dict cache if it is enabled.
@@ -408,3 +429,17 @@ class LazyStateDict:
                     raise KeyError(f"Key {key} not found in LazyStateDict.")
         for key, value in state_dict.items():
             self[key] = value
+
+    def __getattr__(self, name: str):
+        if "meta_module" in self.__dict__:
+            meta_module = self.__dict__["meta_module"]
+            if meta_module is not None:
+                if "_parameters" in meta_module.__dict__:
+                    if name in meta_module.__dict__["_parameters"]:
+                        return self.get_parameter(name)
+                if "_modules" in meta_module.__dict__:
+                    if name in meta_module.__dict__["_modules"]:
+                        return self.get_submodule(name)
+        raise AttributeError(
+            f"'{type(self).__name__}' object has no attribute '{name}'"
+        )