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

fusion_bench/models/utils.py

@@ -1,9 +1,37 @@
-from typing import List
+from typing import Iterable, List, Optional
 
 import torch
 from torch import nn
+from torch.nn.modules.module import _IncompatibleKeys
 
-from fusion_bench.utils.type import StateDictType
+from fusion_bench.utils.dict import dict_merge
+from fusion_bench.utils.type import StateDictType, TorchModelType
+
+
+def is_leaf_module(module: nn.Module) -> bool:
+    return len(list(module.children())) == 0
+
+
+def named_leaf_modules(
+    module: nn.Module,
+    prefix: str = "",
+    ignore_empty: bool = True,
+) -> Iterable[tuple[str, nn.Module]]:
+    """
+    Recursively find the leaf modules in a module.
+
+    Args:
+        module (nn.Module): PyTorch module.
+        prefix (str): A prefix to add to the layer names.
+
+    Returns:
+        Iterable[tuple[str, nn.Module]]: An iterable of (name, module) tuples for each leaf module.
+    """
+    for name, submodule in module.named_modules(prefix=prefix):
+        if is_leaf_module(submodule):
+            if ignore_empty and len(list(submodule.parameters())) == 0:
+                continue
+            yield name, submodule
 
 
 def del_attr(obj, names: List[str]):
@@ -104,3 +132,163 @@ def disable_dropout(model: torch.nn.Module):
     for module in model.modules():
         if isinstance(module, torch.nn.Dropout):
             module.p = 0
+
+
+def get_target_state_dict(
+    module: nn.Module,
+    target_modules: str | Iterable[str] | None = None,
+    prefix: str = "",
+    keep_vars: bool = False,
+) -> StateDictType:
+    """
+    This function retrieves the state dictionary of specified target submodules within a given module
+    of a PyTorch model or merged state dictionary from multiple submodules.
+
+    For example, if a model has submodules named "layer1", "layer2", and "layer3", and you want to get the state dictionary of "layer1" and "layer3",
+    you can call this function with `target_modules` set to `["layer1", "layer3"]`.
+    The function will return a state dictionary that includes only the parameters and buffers from those specified submodules.
+
+    Args:
+        module (nn.Module): The PyTorch module containing the target submodules.
+        target_modules (str | Iterable[str]): A single target module name or an iterable of target module names.
+            If None, the entire module's state dictionary is returned if no special attribute is set (look up the `_fusion_bench_target_modules` attribute).
+        keep_vars (bool): If True, keeps the variables in the state dictionary. Default is False.
+
+    Returns:
+        StateDictType: The state dictionary of the specified target submodules, merged if multiple are provided.
+    """
+    if target_modules is None:
+        if (
+            hasattr(module, "_fusion_bench_target_modules")
+            and module._fusion_bench_target_modules is not None
+        ):
+            return get_target_state_dict(
+                module,
+                target_modules=module._fusion_bench_target_modules,
+                prefix=prefix,
+                keep_vars=keep_vars,
+            )
+        else:
+            return module.state_dict(prefix=prefix, keep_vars=keep_vars)
+
+    if isinstance(target_modules, str):
+        target_modules = [target_modules]
+
+    state_dicts = []
+    for target_module in target_modules:
+        submodule_prefix = (
+            f"{prefix}{target_module}." if prefix else f"{target_module}."
+        )
+        submodule = module.get_submodule(target_module)
+        state_dict = submodule.state_dict(prefix=submodule_prefix, keep_vars=keep_vars)
+        state_dicts.append(state_dict)
+
+    merged_state_dict = dict_merge(state_dicts, disjoint=True)
+    return merged_state_dict
+
+
+def validate_target_modules_equal(modules: Iterable[nn.Module]) -> None:
+    """
+    Validates that the `_fusion_bench_target_modules` attribute is the same across all provided modules.
+
+    Args:
+        modules (Iterable[nn.Module]): An iterable of PyTorch modules to validate.
+
+    Raises:
+        ValueError: If the `_fusion_bench_target_modules` attribute differs among the modules.
+    """
+    model_iter = iter(modules)
+    first_module = next(model_iter)
+
+    if hasattr(first_module, "_fusion_bench_target_modules"):
+        target_modules = first_module._fusion_bench_target_modules
+    else:
+        # if the module does not have the attribute, set to None
+        target_modules = None
+
+    for module in model_iter:
+        if target_modules is None:
+            if (
+                hasattr(module, "_fusion_bench_target_modules")
+                and module._fusion_bench_target_modules != target_modules
+            ):
+                raise ValueError(
+                    "_fusion_bench_target_modules attribute differs among the provided modules."
+                )
+        else:
+            if (
+                not hasattr(module, "_fusion_bench_target_modules")
+                or module._fusion_bench_target_modules != target_modules
+            ):
+                raise ValueError(
+                    "_fusion_bench_target_modules attribute differs among the provided modules."
+                )
+
+
+def load_state_dict_into_target_modules(
+    module: TorchModelType,
+    state_dict: StateDictType,
+    target_modules: str | Iterable[str] | None = None,
+    strict: bool = True,
+    assign: bool = False,
+):
+    """
+    Load a state dictionary into specified target submodules within a given module of a PyTorch model.
+
+    This function allows you to load parameters and buffers from a state dictionary into specific submodules
+    of a PyTorch model. If the `target_modules` argument is provided, only the specified submodules will be updated
+    with the corresponding entries from the state dictionary.
+
+    Args:
+        module (nn.Module): The PyTorch module containing the target submodules.
+        state_dict (StateDictType): The state dictionary containing parameters and buffers to load.
+        target_modules (str | Iterable[str]): A single target module name or an iterable of target module names.
+            If None, the entire module's state dictionary is updated if no special attribute is set
+            (look up the `_fusion_bench_target_modules` attribute).
+        strict (bool): Whether to strictly enforce that the keys in `state_dict` match the keys returned by
+            the module's `state_dict()` function. Default is True.
+    """
+    if target_modules is None:
+        if (
+            hasattr(module, "_fusion_bench_target_modules")
+            and module._fusion_bench_target_modules is not None
+        ):
+            return load_state_dict_into_target_modules(
+                module,
+                state_dict,
+                target_modules=module._fusion_bench_target_modules,
+                strict=strict,
+                assign=assign,
+            )
+        else:
+            return module.load_state_dict(state_dict, strict=strict, assign=assign)
+
+    if isinstance(target_modules, str):
+        target_modules = [target_modules]
+
+    assert (
+        len(target_modules) > 0
+    ), "target_modules should contain at least one module name."
+    results: list[_IncompatibleKeys] = []
+    for target_module in target_modules:
+        submodule_prefix = f"{target_module}."
+        submodule_prefix_len = len(submodule_prefix)
+        submodule = module.get_submodule(target_module)
+
+        # Extract the relevant portion of the state dictionary for the submodule
+        submodule_state_dict = {
+            key[submodule_prefix_len:]: value for key, value in state_dict.items()
+        }
+
+        # Load the extracted state dictionary into the submodule
+        result = submodule.load_state_dict(
+            submodule_state_dict, strict=strict, assign=assign
+        )
+        results.append(result)
+
+    # Merge results from all submodules
+    merged_result = _IncompatibleKeys(
+        missing_keys=[key for res in results for key in res.missing_keys],
+        unexpected_keys=[key for res in results for key in res.unexpected_keys],
+    )
+    return merged_result

fusion_bench/models/wrappers/switch.py

@@ -0,0 +1,90 @@
+"""
+This module contains a wrapper for switching between different models.
+
+For example, it can be used to switch between different classification heads for a shared backbone.
+"""
+
+import logging
+from typing import Dict, Optional
+
+from torch import nn
+
+from fusion_bench.utils.misc import first, validate_and_suggest_corrections
+
+__all__ = ["SwitchModule", "set_active_option"]
+
+log = logging.getLogger(__name__)
+
+
+def _standardize_option_name(name: str) -> str:
+    """
+    Standardizes the option name by:
+
+    - Stripping whitespace and converting to lowercase.
+    - Replacing `-` with `_` if needed.
+    - Replacing `/` with `_` if needed.
+
+    Args:
+        name (str): The option name to standardize.
+    """
+    name = name.strip().lower()
+    name = name.replace("-", "_")
+    name = name.replace("/", "_")
+    return name
+
+
+class SwitchModule(nn.Module):
+    """
+    A wrapper module that contains multiple sub-modules (options) and allows switching between them.
+
+    This is useful for multi-head models or models where different parts are activated based on the task.
+    """
+
+    def __init__(self, modules: Dict[str, nn.Module]):
+        """
+        Args:
+            modules (Dict[str, nn.Module]): A dictionary of modules to switch between.
+        """
+        super().__init__()
+        standardized_modules = {
+            _standardize_option_name(name): module for name, module in modules.items()
+        }
+        self._option_modules = nn.ModuleDict(standardized_modules)
+        self._active_option = first(self._option_modules.keys())
+
+    def set_active_option(self, option_name: str):
+        standardized_name = _standardize_option_name(option_name)
+        validate_and_suggest_corrections(standardized_name, self._option_modules.keys())
+        self._active_option = standardized_name
+
+    def forward(self, *args, **kwargs):
+        active_module = self._option_modules[self._active_option]
+        return active_module(*args, **kwargs)
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            active_module = self._option_modules[self._active_option]
+            if hasattr(active_module, name):
+                return getattr(active_module, name)
+            raise
+
+
+def set_active_option(module: nn.Module, option_name: str) -> list[str]:
+    """
+    Utility function to set the active option for all SwitchModule instances within a given module.
+
+    Args:
+        module (nn.Module): The module to set the active option for.
+        option_name (str): The name of the option to activate.
+
+    Returns:
+        list[str]: A list of names of submodules that were activated.
+    """
+    activated_submodules = []
+    for name, submodule in module.named_modules():
+        if isinstance(submodule, SwitchModule):
+            submodule.set_active_option(option_name)
+            activated_submodules.append(name)
+    return activated_submodules

fusion_bench/programs/base_program.py

@@ -75,6 +75,12 @@ class BaseHydraProgram(BaseYAMLSerializable):
     - FusionBench CLI documentation for program execution details
     """
 
+    _program = None
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._program = self
+
     @abstractmethod
     def run(self):
         """

fusion_bench/programs/fabric_fusion_program.py

@@ -267,6 +267,7 @@ class FabricModelFusionProgram(
         merged_model = self.method.run(self.modelpool)
         self.method.on_run_end()
 
+        report = None
         if merged_model is None:
             log.info(
                 "No merged model returned by the method. Skipping saving and evaluation."
@@ -293,5 +294,8 @@ class FabricModelFusionProgram(
                         )
                     os.makedirs(os.path.dirname(self.report_save_path), exist_ok=True)
                     json.dump(report, open(self.report_save_path, "w"))
+                    self.log_artifact(local_path=self.report_save_path)
             else:
                 log.info("No task pool specified. Skipping evaluation.")
+
+        return {"merged_model": merged_model, "report": report}

fusion_bench/py.typed

@@ -0,0 +1 @@
+

fusion_bench/scripts/cli.py

@@ -3,36 +3,21 @@
 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_name="fabric_model_fusion",
-    version_base=None,
-)
 def main(cfg: DictConfig) -> None:
     """
     Main entry point for the FusionBench command-line interface.
@@ -68,7 +53,7 @@ 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")):
@@ -83,8 +68,25 @@ def main(cfg: DictConfig) -> None:
         err_msg += f"\n\nConfiguration content:\n{cfg}"
         raise TypeError(err_msg)
 
-    program.run()
+    try:
+        program_result = program.run()
+        return program_result
+    except BaseException as e:
+        # Log the exception before exiting
+        if hasattr(program, "finalize") and callable(getattr(program, "finalize")):
+            program.finalize()
+        log.error(e, exc_info=True)
+        raise e
+
+
+@hydra.main(
+    config_path=get_default_config_path(),
+    config_name="fabric_model_fusion",
+    version_base=None,
+)
+def _hydra_main(cfg: DictConfig) -> None:
+    main(cfg)
 
 
 if __name__ == "__main__":
-    main()
+    _hydra_main()

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):
         """