fusion-bench 0.2.29__py3-none-any.whl → 0.2.31__py3-none-any.whl

fusion_bench/models/open_clip/utils.py

@@ -77,7 +77,16 @@ def torch_load_old(save_path: str, device=None):
     return classifier
 
 
-def torch_save(model, save_path, save_state_dict=True):
+def torch_save(model: torch.nn.Module, save_path: str, save_state_dict: bool = True):
+    """
+    Save a model to disk.
+
+    Args:
+        model: The model to save.
+        save_path (str): The path to save the model to.
+        save_state_dict (bool): Whether to save the state dict of the model (weights only).
+            If False, the entire model object is saved. Default is True.
+    """
     # TODO: hacky way to save state dict
     if save_state_dict and isinstance(model, torch.nn.Module):
         model = model.state_dict()
@@ -86,7 +95,9 @@ def torch_save(model, save_path, save_state_dict=True):
     torch.save(model, save_path)
 
 
-def torch_load(save_path, device=None):
+def torch_load(
+    save_path: str, device: Optional[torch.device] = None
+) -> torch.nn.Module:
     model = torch.load(save_path, map_location="cpu")
     if device is not None:
         model = model.to(device)

fusion_bench/models/wrappers/layer_wise_fusion.py

@@ -173,6 +173,24 @@ class LayerWiseMergedModel(nn.Module, Generic[TorchModelType]):
 
     @property
     def forward_model(self):
+        """
+        Get a functional model with merged parameters.
+
+        Returns a partial function that applies the pretrained model with the current
+        merged state dictionary. This allows for efficient forward passes without
+        modifying the original model's parameters.
+
+        Returns:
+            Callable: A partial function that can be called with (args, kwargs) to
+                perform forward pass with merged parameters.
+
+        Example:
+            ```python
+            # Internal usage during forward pass
+            forward_fn = merged_model.forward_model
+            output = forward_fn(args=(x,), kwargs={})
+            ```
+        """
         return functools.partial(
             functional_call,
             self.pretrained_model,
@@ -181,10 +199,30 @@ class LayerWiseMergedModel(nn.Module, Generic[TorchModelType]):
             strict=self.strict,
         )
 
-    def merge_and_unload(self, task_vector_mask: Optional[Dict[str, Tensor]] = None):
+    def merge_and_unload(
+        self,
+        task_vector_mask: Optional[Dict[str, Tensor]] = None,
+        copy: bool = False,
+    ) -> TorchModelType:
+        """
+        Merge models and return the final merged model.
+
+        Args:
+            task_vector_mask (Optional[Dict[str, Tensor]], optional): Optional masks
+                for selective parameter merging. Defaults to None.
+            copy (bool, optional): Whether to return a deep copy of the pretrained model.
+                Defaults to False. If True, the original pretrained model remains unchanged.
+
+        Returns:
+            TorchModelType: The pretrained model with merged parameters loaded.
+        """
         self.merge_weights(task_vector_mask=task_vector_mask)
-        self.pretrained_model.load_state_dict(self._merged_state_dict)
-        return self.pretrained_model
+        if copy:
+            model = deepcopy(self.pretrained_model)
+        else:
+            model = self.pretrained_model
+        model.load_state_dict(self._merged_state_dict)
+        return model
 
     def merge_weights(self, task_vector_mask: Optional[Dict[str, Tensor]] = None):
         """

fusion_bench/models/wrappers/task_wise_fusion.py

@@ -16,6 +16,7 @@ outputs = merged_model(inputs)
 
 import functools
 import logging
+from copy import deepcopy
 from typing import Any, Callable, Dict, Generic, Iterator, List, Optional  # noqa: F401
 
 import torch
@@ -327,7 +328,11 @@ class TaskWiseMergedModel(nn.Module, Generic[TorchModelType]):
         self._merged_state_dict = state_dict
         return state_dict
 
-    def merge_and_unload(self, task_vector_mask: Optional[Dict[str, Tensor]] = None):
+    def merge_and_unload(
+        self,
+        task_vector_mask: Optional[Dict[str, Tensor]] = None,
+        copy: bool = False,
+    ) -> TorchModelType:
         """
         Merge models and return the final merged model.
 
@@ -338,6 +343,8 @@ class TaskWiseMergedModel(nn.Module, Generic[TorchModelType]):
         Args:
             task_vector_mask (Optional[Dict[str, Tensor]], optional): Optional masks
                 for selective parameter merging. Defaults to None.
+            copy (bool, optional): Whether to return a deep copy of the pretrained model.
+                Defaults to False. If True, the original pretrained model remains unchanged.
 
         Returns:
             TorchModelType: The pretrained model with merged parameters loaded.
@@ -363,8 +370,12 @@ class TaskWiseMergedModel(nn.Module, Generic[TorchModelType]):
             The original pretrained model parameters will be lost.
         """
         self.merge_weights(task_vector_mask=task_vector_mask)
-        self.pretrained_model.load_state_dict(self._merged_state_dict)
-        return self.pretrained_model
+        if copy:
+            model = deepcopy(self.pretrained_model)
+        else:
+            model = self.pretrained_model
+        model.load_state_dict(self._merged_state_dict)
+        return model
 
     def forward(self, *args, **kwargs):
         """

fusion_bench/py.typed

@@ -0,0 +1 @@
+

fusion_bench/scripts/cli.py

@@ -3,33 +3,24 @@
 This is the CLI script that is executed when the user runs the `fusion_bench` command.
 The script is responsible for parsing the command-line arguments, loading the configuration file, and running the fusion algorithm.
 """
-
-import importlib
-import importlib.resources
 import logging
-import os
+from typing import TYPE_CHECKING
 
 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.utils.hydra_utils import get_default_config_path
 
-log = logging.getLogger(__name__)
-
+if TYPE_CHECKING:
+    from fusion_bench.programs import BaseHydraProgram
 
-def _get_default_config_path():
-    for config_path_root in [os.getcwd(), PROJECT_ROOT_PATH]:
-        for config_dir in ["config", "fusion_bench_config"]:
-            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
+log = logging.getLogger(__name__)
 
 
 @hydra.main(
-    config_path=_get_default_config_path(),
+    config_path=get_default_config_path(),
     config_name="fabric_model_fusion",
     version_base=None,
 )
@@ -68,7 +59,21 @@ def main(cfg: DictConfig) -> None:
         loading the corresponding configuration files to populate the cfg parameter.
     """
     OmegaConf.resolve(cfg)
-    program: BaseHydraProgram = instantiate(cfg)
+    program: "BaseHydraProgram" = instantiate(cfg)
+
+    # Validate that instantiation succeeded and returned an object with 'run' method
+    if not hasattr(program, "run") or not callable(getattr(program, "run")):
+        err_msg = (
+            f"Expected an object with a callable 'run' method, but got {type(program).__name__}. "
+            "Ensure that the configuration specifies a concrete program class with '_target_'."
+        )
+        if "_target_" not in cfg:
+            err_msg += "\nThe '_target_' field is missing from the root configuration."
+        else:
+            err_msg += f"\nFound '_target_': {cfg._target_}"
+        err_msg += f"\n\nConfiguration content:\n{cfg}"
+        raise TypeError(err_msg)
+
     program.run()
 
 

fusion_bench/scripts/imgui.py

@@ -9,7 +9,7 @@ import hydra
 from hydra import compose, initialize_config_dir
 from omegaconf import DictConfig, ListConfig, OmegaConf
 
-from fusion_bench.scripts.cli import _get_default_config_path
+from fusion_bench.scripts.cli import get_default_config_path
 
 # Keeping the ConfigGroupNode and AppState classes as they are
 from fusion_bench.scripts.webui import AppState, ConfigGroupNode, priority_iterable
@@ -40,7 +40,7 @@ class App:
         if self.args.config_path:
             return Path(self.args.config_path)
         else:
-            return _get_default_config_path()
+            return get_default_config_path()
 
     def generate_ui(self):
         dpg.create_context()

fusion_bench/scripts/webui.py

@@ -16,7 +16,7 @@ from hydra import compose, initialize_config_dir
 from hydra.core.hydra_config import HydraConfig
 from omegaconf import DictConfig, ListConfig, OmegaConf
 
-from fusion_bench.scripts.cli import _get_default_config_path
+from fusion_bench.scripts.cli import get_default_config_path
 
 
 def escape_overrides(value: str) -> str:
@@ -385,7 +385,7 @@ class App:
         if self.args.config_path:
             return Path(self.args.config_path)
         else:
-            return _get_default_config_path()
+            return get_default_config_path()
 
     def __getattr__(self, name):
         """

fusion_bench/utils/__init__.py

@@ -53,6 +53,7 @@ _import_structure = {
         "get_parameter_summary",
         "human_readable",
         "print_parameters",
+        "print_trainable_parameters",
         "state_dict_to_vector",
         "trainable_state_dict",
         "vector_to_state_dict",
@@ -138,6 +139,7 @@ if TYPE_CHECKING:
         get_parameter_summary,
         human_readable,
         print_parameters,
+        print_trainable_parameters,
         state_dict_to_vector,
         trainable_state_dict,
         vector_to_state_dict,

fusion_bench/utils/devices.py

@@ -32,11 +32,13 @@ def clear_cuda_cache():
     Clears the CUDA memory cache to free up GPU memory.
     Works only if CUDA is available.
     """
+
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
+        gc.collect()
     else:
-        log.warning("CUDA is not available. No cache to clear.")
+        log.debug("CUDA is not available. No cache to clear.")
 
 
 def to_device(

fusion_bench/utils/hydra_utils.py

@@ -1,4 +1,79 @@
+import logging
+import os
+
 import hydra.core.hydra_config
+from hydra import compose, initialize
+from omegaconf import DictConfig
+
+from fusion_bench.constants import PROJECT_ROOT_PATH
+
+log = logging.getLogger(__name__)
+
+
+def get_default_config_path():
+    """
+    Get the default configuration path by searching in common locations.
+    """
+    for config_path_root in [os.getcwd(), PROJECT_ROOT_PATH]:
+        for config_dir in ["config", "fusion_bench_config"]:
+            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
+
+
+def initialize_hydra_config(
+    config_name: str,
+    overrides: list[str] = None,
+    config_path: str = None,
+    return_hydra_config: bool = False,
+) -> DictConfig:
+    """
+    Load the Hydra configuration.
+
+    Args:
+        config_name (str): The name of the configuration file (without .yaml extension).
+        overrides (list[str]): A list of configuration overrides.
+        config_path (str): The path to the configuration directory. If None, it will be automatically detected.
+        return_hydra_config (bool): If True, return the Hydra configuration object.
+
+    Returns:
+        DictConfig: The loaded configuration.
+
+    Example:
+        >>> cfg = initialize_hydra_config(
+        ...     config_name="fabric_model_fusion",
+        ...     overrides=["method=dummy", "modelpool=dummy"],
+        ... )
+        >>> print(cfg.method)
+    """
+    if config_path is None:
+        config_path = get_default_config_path()
+
+    # check config_path validity
+    if config_path is None:
+        raise FileNotFoundError("Could not find configuration directory.")
+    if not os.path.isdir(config_path):
+        raise NotADirectoryError(
+            f"Configuration path {config_path} do not exists or is not a directory."
+        )
+
+    if overrides is None:
+        overrides = []
+
+    with initialize(
+        version_base=None,
+        config_path=os.path.relpath(
+            config_path,
+            start=os.path.dirname(__file__),
+        ),
+    ):
+        cfg = compose(
+            config_name=config_name,
+            overrides=overrides,
+            return_hydra_config=return_hydra_config,
+        )
+        return cfg
 
 
 def get_hydra_output_dir():

fusion_bench/utils/instantiate_utils.py

@@ -14,8 +14,8 @@ from lightning_utilities.core.rank_zero import rank_zero_only
 from omegaconf import DictConfig, OmegaConf, SCMode
 from omegaconf._utils import is_structured_config
 from rich import print
-from rich.panel import Panel
-from rich.syntax import Syntax
+
+from fusion_bench.utils.rich_utils import print_bordered
 
 PRINT_FUNCTION_CALL = True
 """
@@ -67,12 +67,22 @@ def _resolve_callable_name(f: Callable[..., Any]) -> str:
     return full_name
 
 
-def _format_args_kwargs(args, kwargs):
+def _get_obj_str(obj: Any) -> str:
+    if isinstance(obj, (str, int, float, bool, type(None))):
+        return repr(obj)
+    else:
+        return f"'<{type(obj).__name__} object>'"
+
+
+def _format_args_kwargs(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> str:
     result_strings = []
     if len(args) > 0:
-        result_strings.append(", ".join(repr(arg) for arg in args))
+        result_strings.append(", ".join(_get_obj_str(arg) for arg in args))
+
     if len(kwargs) > 0:
-        result_strings.append(", ".join(f"{k}={repr(v)}" for k, v in kwargs.items()))
+        result_strings.append(
+            ", ".join(f"{k}={_get_obj_str(v)}" for k, v in kwargs.items())
+        )
 
     if len(result_strings) == 0:
         return ""
@@ -145,14 +155,14 @@ def _call_target(
     if _partial_:
         if PRINT_FUNCTION_CALL and getattr(rank_zero_only, "rank", 0) == 0:
             call_str = f"functools.partial({_resolve_callable_name(_target_)}, {_format_args_kwargs(args, kwargs)})"
-            PRINT_FUNCTION_CALL_FUNC(
-                Panel(
-                    Syntax(call_str, "python", theme="monokai", word_wrap=True),
-                    title="Instantiate by calling partial",
-                    border_style="cyan",
-                )
+            print_bordered(
+                call_str,
+                code_style="python",
+                title=f"Instantiate by calling {'function' if not isinstance(_target_, type) else 'class'}",
+                style="cyan",
+                expand=False,
+                print_fn=PRINT_FUNCTION_CALL_FUNC,
             )
-
         if CATCH_EXCEPTION:
             try:
                 return functools.partial(_target_, *args, **kwargs)
@@ -169,12 +179,13 @@ def _call_target(
     else:
         if PRINT_FUNCTION_CALL and getattr(rank_zero_only, "rank", 0) == 0:
             call_str = f"{_resolve_callable_name(_target_)}({_format_args_kwargs(args, kwargs)})"
-            PRINT_FUNCTION_CALL_FUNC(
-                Panel(
-                    Syntax(call_str, "python", theme="monokai", word_wrap=True),
-                    title="Instantiate by calling function",
-                    border_style="green",
-                )
+            print_bordered(
+                call_str,
+                code_style="python",
+                title=f"Instantiate by calling {'function' if not isinstance(_target_, type) else 'class'}",
+                style="green",
+                expand=False,
+                print_fn=PRINT_FUNCTION_CALL_FUNC,
             )
         if CATCH_EXCEPTION:
             try:

fusion_bench/utils/misc.py

@@ -178,3 +178,19 @@ def validate_and_suggest_corrections(
     if matches:
         msg += f". Did you mean {', '.join(repr(m) for m in matches)}?"
     raise ValueError(msg)
+
+
+class DeprecationWarningMeta(type):
+    """
+    Metaclass that issues a deprecation warning whenever a class using it is instantiated.
+    """
+
+    def __call__(cls, *args, **kwargs):
+        import warnings
+
+        warnings.warn(
+            f"{cls.__name__} is deprecated and will be removed in a future version. ",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return super(DeprecationWarningMeta, cls).__call__(*args, **kwargs)

fusion_bench/utils/parameters.py

@@ -10,6 +10,7 @@ from .type import StateDictType
 __all__ = [
     "count_parameters",
     "print_parameters",
+    "print_trainable_parameters",
     "check_parameters_all_equal",
     "get_parameter_statistics",
     "state_dict_to_vector",
@@ -282,6 +283,38 @@ def print_parameters(
     )
 
 
+def print_trainable_parameters(
+    module: nn.Module,
+    is_human_readable: bool = True,
+    print_fn=print,
+    non_zero_only: bool = False,
+):
+    """
+    Print the names and number of trainable parameters in a PyTorch model.
+
+    Args:
+        module (nn.Module): The PyTorch model.
+        is_human_readable (bool, optional): Whether to print the number of parameters in a human-readable format. Defaults to True.
+        print_fn (callable, optional): The function to use for printing. Defaults to print.
+        non_zero_only (bool, optional): Whether to count only non-zero parameters. Defaults to False.
+
+    Prints:
+        The names and number of trainable parameters in the model.
+
+        ```python
+        print_trainable_parameters(model)
+        # weight: 1.50M parameters
+        # bias: 500.00K parameters
+        ```
+    """
+    for name, param in module.named_parameters():
+        if param.requires_grad:
+            num_params = _numel(param, non_zero_only=non_zero_only)
+            if is_human_readable:
+                num_params = human_readable(num_params)
+            print_fn(f"{name}: {num_params} parameters")
+
+
 def check_parameters_all_equal(
     list_of_param_names: List[Union[StateDictType, nn.Module, List[str]]],
 ) -> None: