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

fusion_bench/models/parameter_dict.py

@@ -1,12 +1,12 @@
-from typing import List, Mapping, Optional, Tuple
+from typing import Iterator, List, Mapping, Optional, Tuple, Union
 
 import torch
 from torch import nn
 
-__all__ = "ParamterDictModel"
+__all__ = ["ParameterDictModel"]
 
 
-def _set_attr(
+def set_nested_attr(
     obj,
     names: List[str],
     val,
@@ -27,7 +27,7 @@ def _set_attr(
     else:
         if check_parent and not hasattr(obj, names[0]):
             setattr(obj, names[0], parent_builder())
-        _set_attr(
+        set_nested_attr(
             getattr(obj, names[0]),
             names[1:],
             val,
@@ -36,7 +36,7 @@ def _set_attr(
         )
 
 
-def has_attr(obj, names: List[str]):
+def has_nested_attr(obj, names: List[str]):
     """
     Checks if an attribute exists in an object recursively.
 
@@ -50,26 +50,49 @@ def has_attr(obj, names: List[str]):
     if len(names) == 1:
         return hasattr(obj, names[0])
     else:
-        return has_attr(getattr(obj, names[0]), names[1:])
+        if not hasattr(obj, names[0]):
+            return False
+        return has_nested_attr(getattr(obj, names[0]), names[1:])
 
 
 class ParameterDictModel(nn.Module):
     """
-    This model is used to create a model with parameters from a dictionary.
-    It behaves like a normal `nn.ParameterDict`, but support keys with dots.
+    A module that stores parameters in a nested dictionary structure.
+
+    This model behaves similarly to `nn.ParameterDict`, but supports hierarchical keys
+    with dots (e.g., "layer1.weight"). Parameters are stored as nested attributes,
+    allowing for structured parameter access and manipulation.
+
+    Example:
+        >>> params = {
+        ...     "encoder.weight": nn.Parameter(torch.randn(10, 5)),
+        ...     "decoder.bias": nn.Parameter(torch.randn(5)),
+        ... }
+        >>> model = ParameterDictModel(params)
+        >>> model["encoder.weight"].shape
+        torch.Size([10, 5])
+        >>> "encoder.weight" in model
+        True
     """
 
     def __init__(
         self,
-        parameters: Optional[Mapping[str, nn.Parameter]] = None,
-    ):
+        parameters: Optional[Mapping[str, Union[nn.Parameter, torch.Tensor]]] = None,
+    ) -> None:
+        """
+        Args:
+            parameters: Optional mapping of parameter names to parameter tensors.
+                Keys can contain dots to create nested structures.
+                Values must be `nn.Parameter` or `nn.Buffer` instances.
+        """
+
         super().__init__()
         if parameters is not None:
             for name, param in parameters.items():
                 assert isinstance(
                     param, (nn.Parameter, nn.Buffer)
                 ), f"{name} is not a nn.Parameter or nn.Buffer"
-                _set_attr(
+                set_nested_attr(
                     self,
                     name.split("."),
                     param,
@@ -77,12 +100,13 @@ class ParameterDictModel(nn.Module):
                     parent_builder=__class__,
                 )
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """
         Generate a string representation of the model's parameters.
 
         Returns:
-            str: A string representation of the model's parameters.
+            A string representation of the model's parameters in the format:
+            "ParameterDictModel(name1: shape1, name2: shape2, ...)"
         """
         param_reprs = []
         for name, param in self.named_parameters():
@@ -90,32 +114,98 @@ class ParameterDictModel(nn.Module):
             param_reprs.append(param_repr)
         return f"{self.__class__.__name__}({', '.join(param_reprs)})"
 
-    def __getitem__(self, key: str):
-        if not has_attr(self, key.split(".")):
+    def __iter__(self) -> Iterator[str]:
+        """
+        Iterate over the model's parameters.
+
+        Yields:
+            Tuples of (parameter name, parameter tensor).
+        """
+        yield from self.keys()
+
+    def __getitem__(
+        self, key: str
+    ) -> Union[nn.Parameter, torch.Tensor, "ParameterDictModel"]:
+        """
+        Retrieve a parameter or nested submodule by key.
+
+        Args:
+            key: Parameter name, which can contain dots for nested access.
+
+        Returns:
+            The parameter, tensor, or nested ParameterDictModel at the specified key.
+
+        Raises:
+            KeyError: If the key is not found in the model.
+        """
+        assert isinstance(
+            key, str
+        ), f"Key must be a string, but got {type(key)}: {key}."
+        if not has_nested_attr(self, key.split(".")):
             raise KeyError(f"Key {key} not found in {self}")
-        key = key.split(".")
+        key_parts = key.split(".")
         obj = self
-        for k in key:
+        for k in key_parts:
             obj = getattr(obj, k)
         return obj
 
-    def __setitem__(self, key: str, value: nn.Parameter):
-        if not has_attr(self, key.split(".")):
-            _set_attr(self, key.split("."), value, check_parent=True)
+    def __setitem__(self, key: str, value: Union[nn.Parameter, torch.Tensor]) -> None:
+        """
+        Set a parameter at the specified key, creating nested structure if needed.
+
+        Args:
+            key: Parameter name, which can contain dots for nested assignment.
+            value: Parameter or tensor to assign.
+        """
+        if not has_nested_attr(self, key.split(".")):
+            set_nested_attr(self, key.split("."), value, check_parent=True)
         else:
-            _set_attr(self, key.split("."), value, check_parent=False)
+            set_nested_attr(self, key.split("."), value, check_parent=False)
+
+    def __contains__(self, key: str) -> bool:
+        """
+        Check if a parameter key exists in the model.
 
-    def __contains__(self, key: str):
-        return has_attr(self, key.split("."))
+        Args:
+            key: Parameter name, which can contain dots for nested checking.
+
+        Returns:
+            True if the key exists, False otherwise.
+        """
+        return has_nested_attr(self, key.split("."))
 
     def keys(self):
-        return [name for name, _ in self.named_parameters()]
+        """
+        Return a list of all parameter names in the model.
+
+        Returns:
+            List of parameter names (including nested names with dots).
+        """
+        return self.state_dict().keys()
+
+    def items(self):
+        """
+        Return a list of (name, parameter) tuples.
+
+        Returns:
+            List of tuples containing parameter names and their corresponding tensors.
+        """
+        yield from self.state_dict().items()
 
-    def items(self) -> List[Tuple[str, nn.Parameter]]:
-        return [(name, self[name]) for name in self.keys()]
+    def values(self):
+        """
+        Return a list of all parameter values in the model.
 
-    def values(self) -> List[nn.Parameter]:
-        return [self[name] for name in self.keys()]
+        Returns:
+            List of parameter tensors.
+        """
+        yield from self.state_dict().values()
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """
+        Return the number of parameters in the model.
+
+        Returns:
+            The total number of parameters.
+        """
         return len(self.keys())

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/scripts/cli.py

@@ -9,7 +9,6 @@ from typing import TYPE_CHECKING
 import hydra
 from omegaconf import DictConfig, OmegaConf
 
-from fusion_bench.constants import PROJECT_ROOT_PATH
 from fusion_bench.utils import instantiate
 from fusion_bench.utils.hydra_utils import get_default_config_path
 
@@ -19,11 +18,6 @@ if TYPE_CHECKING:
 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.
@@ -74,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()