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

fusion_bench/scripts/cli.py

@@ -12,9 +12,9 @@ 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
-from fusion_bench.constants import PROJECT_ROOT_PATH
 
 log = logging.getLogger(__name__)
 
@@ -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/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
@@ -86,8 +87,9 @@ class LayerWiseFeatureSaver:
 
 
 class CLIPVisionModelTaskPool(
-    BaseTaskPool,
+    HydraConfigMixin,
     LightningFabricMixin,
+    BaseTaskPool,
 ):
     """
     This class is used to define the image classification task for CLIP models.
@@ -131,7 +133,7 @@ class CLIPVisionModelTaskPool(
         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,
     ):
         """
@@ -153,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):

fusion_bench/utils/__init__.py

@@ -18,4 +18,5 @@ 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/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/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}'"
+        )

fusion_bench/utils/modelscope.py

@@ -26,13 +26,13 @@ try:
     from huggingface_hub import snapshot_download as huggingface_snapshot_download
 except ImportError:
 
-    def _raise_hugggingface_not_installed_error(*args, **kwargs):
+    def _raise_huggingface_not_installed_error(*args, **kwargs):
         raise ImportError(
             "Hugging Face Hub is not installed. Please install it using `pip install huggingface_hub` to use Hugging Face models."
         )
 
-    huggingface_snapshot_download = _raise_hugggingface_not_installed_error
-    hf_hub_download = _raise_hugggingface_not_installed_error
+    huggingface_snapshot_download = _raise_huggingface_not_installed_error
+    hf_hub_download = _raise_huggingface_not_installed_error
 
 __all__ = [
     "load_dataset",

fusion_bench/utils/path.py

@@ -1,6 +1,9 @@
+import logging
 import os
 from typing import List
 
+log = logging.getLogger(__name__)
+
 
 def path_is_dir_and_not_empty(path: str):
     if path is None:
@@ -20,3 +23,56 @@ def listdir_fullpath(dir: str) -> List[str]:
     assert os.path.isdir(dir), "Argument 'dir' must be a Directory"
     names = os.listdir(dir)
     return [os.path.join(dir, name) for name in names]
+
+
+def create_symlink(src_dir: str, dst_dir: str, link_name: str = None):
+    """
+    Creates a symbolic link from src_dir to dst_dir.
+
+    Args:
+        src_dir (str): The source directory to link to.
+        dst_dir (str): The destination directory where the symlink will be created.
+        link_name (str, optional): The name of the symlink. If None, uses the basename of src_dir.
+
+    Raises:
+        OSError: If the symbolic link creation fails.
+        ValueError: If src_dir does not exist or is not a directory.
+    """
+    if not os.path.exists(src_dir):
+        raise ValueError(f"Source directory does not exist: {src_dir}")
+
+    if not os.path.isdir(src_dir):
+        raise ValueError(f"Source path is not a directory: {src_dir}")
+
+    # Avoid creating symlink if source and destination are the same
+    if os.path.abspath(src_dir) == os.path.abspath(dst_dir):
+        log.warning(
+            "Source and destination directories are the same, skipping symlink creation"
+        )
+        return
+
+    # Create destination directory if it doesn't exist
+    os.makedirs(dst_dir, exist_ok=True)
+
+    # Determine link name
+    if link_name is None:
+        link_name = os.path.basename(src_dir)
+
+    link_path = os.path.join(dst_dir, link_name)
+
+    try:
+        # if the system is windows, use the `mklink` command in "CMD" to create the symlink
+        if os.name == "nt":
+            os.system(
+                f"mklink /J {os.path.abspath(link_path)} {os.path.abspath(src_dir)}"
+            )
+        else:
+            os.symlink(
+                src_dir,
+                link_path,
+                target_is_directory=True,
+            )
+        log.info(f"Created symbolic link: {link_path} -> {src_dir}")
+    except OSError as e:
+        log.warning(f"Failed to create symbolic link: {e}")
+        raise

fusion_bench/utils/pylogger.py

@@ -74,7 +74,7 @@ RankZeroLogger.manager = logging.Manager(RankZeroLogger.root)
 RankZeroLogger.manager.setLoggerClass(RankZeroLogger)
 
 
-def getRankZeroLogger(name=None):
+def get_rankzero_logger(name=None):
     """
     Return a logger with the specified name, creating it if necessary.
 

{fusion_bench-0.2.21.dist-info → fusion_bench-0.2.22.dist-info}/METADATA

@@ -1,30 +1,8 @@
 Metadata-Version: 2.4
 Name: fusion_bench
-Version: 0.2.21
+Version: 0.2.22
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
-License: MIT License
-        
-        Copyright (c) 2024 Anke Tang
-        
-        Permission is hereby granted, free of charge, to any person obtaining a copy
-        of this software and associated documentation files (the "Software"), to deal
-        in the Software without restriction, including without limitation the rights
-        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-        copies of the Software, and to permit persons to whom the Software is
-        furnished to do so, subject to the following conditions:
-        
-        The above copyright notice and this permission notice shall be included in all
-        copies or substantial portions of the Software.
-        
-        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-        SOFTWARE.
-        
 Project-URL: Repository, https://github.com/tanganke/fusion_bench
 Project-URL: Homepage, https://github.com/tanganke/fusion_bench
 Project-URL: Issues, https://github.com/tanganke/fusion_bench/issues