fusion-bench 0.2.24__py3-none-any.whl → 0.2.25__py3-none-any.whl

fusion_bench/programs/fabric_fusion_program.py

@@ -5,7 +5,7 @@ from pathlib import Path
 from typing import Any, Callable, Dict, Iterable, List, Optional, Union  # noqa: F401
 
 import lightning as L
-from lightning.fabric.utilities.rank_zero import rank_zero_only
+from lightning_utilities.core.rank_zero import rank_zero_only
 from omegaconf import DictConfig, OmegaConf
 from torch import nn
 from tqdm.auto import tqdm
@@ -236,7 +236,7 @@ class FabricModelFusionProgram(
 
         # create symbol link to hydra output directory
         if (
-            self.fabric.is_global_zero
+            rank_zero_only.rank == 0
             and self.log_dir is not None
             and os.path.abspath(self.log_dir) != os.path.abspath(get_hydra_output_dir())
         ):

fusion_bench/programs/fusion_program.py

@@ -0,0 +1,271 @@
+import json
+import os
+from typing import Any, Dict, Iterable, List, Optional, Union
+
+import lightning as L
+from lightning_utilities.core.rank_zero import rank_zero_only
+from omegaconf import DictConfig, OmegaConf
+from torch import nn
+from tqdm.auto import tqdm
+
+from fusion_bench import (
+    BaseAlgorithm,
+    BaseHydraProgram,
+    BaseModelPool,
+    BaseTaskPool,
+    RuntimeConstants,
+    auto_register_config,
+    get_rankzero_logger,
+    import_object,
+    instantiate,
+    timeit_context,
+)
+from fusion_bench.utils.json import print_json
+from fusion_bench.utils.rich_utils import print_bordered, print_config_tree
+
+log = get_rankzero_logger(__name__)
+
+
+@auto_register_config
+class ModelFusionProgram(BaseHydraProgram):
+    method: BaseAlgorithm
+    modelpool: BaseModelPool
+    taskpool: Optional[BaseTaskPool] = None
+
+    _config_mapping = BaseHydraProgram._config_mapping | {
+        "_method": "method",
+        "_modelpool": "modelpool",
+        "_taskpool": "taskpool",
+        "fast_dev_run": "fast_dev_run",
+        "seed": "seed",
+        "path": "path",
+    }
+
+    def __init__(
+        self,
+        method: DictConfig,
+        modelpool: DictConfig,
+        taskpool: Optional[DictConfig] = None,
+        *,
+        print_config: bool = True,
+        dry_run: bool = False,
+        report_save_path: Optional[str] = None,
+        merged_model_save_path: Optional[str] = None,
+        merged_model_save_kwargs: Optional[DictConfig] = None,
+        fast_dev_run: bool = False,
+        seed: Optional[int] = None,
+        print_function_call: bool = True,
+        path: DictConfig = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self._method = method
+        self._modelpool = modelpool
+        self._taskpool = taskpool
+        self.report_save_path = report_save_path
+        self.merged_model_save_path = merged_model_save_path
+        self.merged_model_save_kwargs = merged_model_save_kwargs
+        self.fast_dev_run = fast_dev_run
+        self.seed = seed
+        self.path = path
+        RuntimeConstants.debug = fast_dev_run
+        RuntimeConstants.print_function_call = print_function_call
+        if path is not None:
+            RuntimeConstants.cache_dir = path.get("cache_dir", None)
+
+        if print_config:
+            print_config_tree(
+                self.config,
+                print_order=["method", "modelpool", "taskpool"],
+            )
+        if dry_run:
+            log.info("The program is running in dry-run mode. Exiting.")
+            exit(0)
+
+    def _instantiate_and_setup(
+        self, config: DictConfig, compat_load_fn: Optional[str] = None
+    ):
+        R"""
+        Instantiates and sets up an object based on the provided configuration.
+
+        This method performs the following steps:
+        1. Checks if the configuration dictionary contains the key "_target_".
+        2. If "_target_" is not found (for v0.1.x), attempts to instantiate the object using a compatible load function if provided.
+           - Logs a warning if "_target_" is missing.
+           - If `compat_load_fn` is provided, imports the function and uses it to instantiate the object.
+           - If `compat_load_fn` is not provided, raises a ValueError.
+        3. If "_target_" is found (for v.0.2.0 and above), attempts to import and instantiate the object using the `instantiate` function.
+           - Ensures the target can be imported.
+           - Uses the `instantiate` function with `_recursive_` set based on the configuration.
+        4. Sets the `_program` attribute of the instantiated object to `self` if the object has this attribute.
+        5. Returns the instantiated and set up object.
+        """
+        if "_target_" not in config:
+            log.warning(
+                "No '_target_' key found in config. Attempting to instantiate the object in a compatible way."
+            )
+            if compat_load_fn is not None:
+                compat_load_fn = import_object(compat_load_fn)
+                if rank_zero_only.rank == 0:
+                    print_bordered(
+                        OmegaConf.to_yaml(config),
+                        title="instantiate compat object",
+                        style="magenta",
+                        code_style="yaml",
+                    )
+                obj = compat_load_fn(config)
+            else:
+                raise ValueError(
+                    "No load function provided. Please provide a load function to instantiate the object."
+                )
+        else:
+            # try to import the object from the target
+            # this checks if the target is valid and can be imported
+            import_object(config._target_)
+            obj = instantiate(
+                config,
+                _recursive_=config.get("_recursive_", False),
+            )
+        if hasattr(obj, "_program"):
+            obj._program = self
+        return obj
+
+    def save_merged_model(self, merged_model):
+        """
+        Saves the merged model to the specified path.
+        """
+        if self.merged_model_save_path is not None:
+            # path to save the merged model, use "{log_dir}" to refer to the logger directory
+            save_path: str = self.merged_model_save_path
+            if "{log_dir}" in save_path and self.log_dir is not None:
+                save_path = save_path.format(log_dir=self.log_dir)
+
+            if os.path.dirname(save_path):
+                os.makedirs(os.path.dirname(save_path), exist_ok=True)
+
+            # save the merged model
+            if self.merged_model_save_kwargs is not None:
+                merged_model_save_kwargs = self.merged_model_save_kwargs
+            else:
+                merged_model_save_kwargs = {}
+            with timeit_context(f"Saving the merged model to {save_path}"):
+                self.modelpool.save_model(
+                    merged_model,
+                    save_path,
+                    **merged_model_save_kwargs,
+                )
+        else:
+            print("No save path specified for the merged model. Skipping saving.")
+
+    def evaluate_merged_model(
+        self,
+        taskpool: BaseTaskPool,
+        merged_model: Union[nn.Module, Dict, Iterable],
+        *args: Any,
+        **kwargs: Any,
+    ) -> Union[Dict, List, Any]:
+        """
+        Evaluates the merged model using the provided task pool.
+
+        Depending on the type of the merged model, this function handles the evaluation differently:
+        - If the merged model is an instance of `nn.Module`, it directly evaluates the model.
+        - If the merged model is a dictionary, it extracts the model from the dictionary and evaluates it.
+          The evaluation report is then updated with the remaining dictionary items.
+        - If the merged model is an iterable, it recursively evaluates each model in the iterable.
+        - Raises a `ValueError` if the merged model is of an invalid type.
+
+        Args:
+            taskpool: The task pool used for evaluating the merged model.
+            merged_model: The merged model to be evaluated. It can be an instance of `nn.Module`, a dictionary, or an iterable.
+            *args: Additional positional arguments to be passed to the `evaluate` method of the taskpool.
+            **kwargs: Additional keyword arguments to be passed to the `evaluate` method of the taskpool.
+
+        Returns:
+            The evaluation report. The type of the report depends on the type of the merged model:
+            - If the merged model is an instance of `nn.Module`, the report is a dictionary.
+            - If the merged model is a dictionary, the report is a dictionary updated with the remaining dictionary items.
+            - If the merged model is an iterable, the report is a list of evaluation reports.
+        """
+        if isinstance(merged_model, nn.Module):
+            report = taskpool.evaluate(merged_model, *args, **kwargs)
+            return report
+        elif isinstance(merged_model, Dict):
+            report = {}
+            for key, item in merged_model.items():
+                if isinstance(item, nn.Module):
+                    report[key] = taskpool.evaluate(item, *args, **kwargs)
+                elif key == "models":
+                    # for multi-model evaluation
+                    report[key] = self.evaluate_merged_model(
+                        taskpool, item, *args, **kwargs
+                    )
+                else:
+                    # metadata
+                    report[key] = item
+            return report
+        elif isinstance(merged_model, Iterable):
+            return [
+                self.evaluate_merged_model(taskpool, m, *args, **kwargs)
+                for m in tqdm(merged_model, desc="Evaluating models")
+            ]
+        else:
+            raise ValueError(f"Invalid type for merged model: {type(merged_model)}")
+
+    def run(self):
+        """
+        Executes the model fusion program.
+        """
+        if self.seed is not None:
+            L.seed_everything(self.seed)
+
+        log.info("Running the model fusion program.")
+        # setup the modelpool, method, and taskpool
+        log.info("loading model pool")
+        self.modelpool = self._instantiate_and_setup(
+            self._modelpool,
+            compat_load_fn="fusion_bench.compat.modelpool.load_modelpool_from_config",
+        )
+        log.info("loading method")
+        self.method = self._instantiate_and_setup(
+            self._method,
+            compat_load_fn="fusion_bench.compat.method.load_algorithm_from_config",
+        )
+        if self._taskpool is not None:
+            log.info("loading task pool")
+            self.taskpool = self._instantiate_and_setup(
+                self._taskpool,
+                compat_load_fn="fusion_bench.compat.taskpool.load_taskpool_from_config",
+            )
+
+        self.method.on_run_start()
+        merged_model = self.method.run(self.modelpool)
+        self.method.on_run_end()
+
+        if merged_model is None:
+            log.info(
+                "No merged model returned by the method. Skipping saving and evaluation."
+            )
+        else:
+            self.save_merged_model(merged_model)
+            if self.taskpool is not None:
+                report = self.evaluate_merged_model(self.taskpool, merged_model)
+                try:
+                    if rank_zero_only.rank == 0:
+                        print_json(report, print_type=False)
+                except Exception as e:
+                    log.warning(f"Failed to pretty print the report: {e}")
+                    log.info(report)
+                if self.report_save_path is not None:
+                    # save report (Dict) to a file
+                    # if the directory of `save_report` does not exists, create it
+                    if (
+                        "{log_dir}" in self.report_save_path
+                        and self.path.log_dir is not None
+                    ):
+                        self.report_save_path = self.report_save_path.format(
+                            log_dir=self.path.log_dir
+                        )
+                    os.makedirs(os.path.dirname(self.report_save_path), exist_ok=True)
+                    json.dump(report, open(self.report_save_path, "w"))
+            else:
+                log.info("No task pool specified. Skipping evaluation.")

fusion_bench/tasks/clip_classification/__init__.py

@@ -183,3 +183,18 @@ class CLIPTemplateFactory:
 
 def get_classnames_and_templates(dataset_name: str) -> Tuple[List[str], List[Callable]]:
     return CLIPTemplateFactory.get_classnames_and_templates(dataset_name)
+
+
+def get_num_classes(dataset_name: str) -> int:
+    classnames, _ = get_classnames_and_templates(dataset_name)
+    return len(classnames)
+
+
+def get_classnames(dataset_name: str) -> List[str]:
+    classnames, _ = get_classnames_and_templates(dataset_name)
+    return classnames
+
+
+def get_templates(dataset_name: str) -> List[Callable]:
+    _, templates = get_classnames_and_templates(dataset_name)
+    return templates

fusion_bench/utils/__init__.py

@@ -1,23 +1,169 @@
 # flake8: noqa: F401
-import importlib
-from typing import Iterable
+import sys
+from typing import TYPE_CHECKING
 
-from . import data, functools, path, pylogger
-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,
-    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
-from .type import BoolStateDictType, StateDictType, TorchModelType
+from . import functools
+from .lazy_imports import LazyImporter
+
+_extra_objects = {
+    "functools": functools,
+}
+_import_structure = {
+    "cache_utils": [
+        "cache_to_disk",
+        "cache_with_joblib",
+        "set_default_cache_dir",
+    ],
+    "data": [
+        "InfiniteDataLoader",
+        "load_tensor_from_file",
+        "train_validation_split",
+        "train_validation_test_split",
+    ],
+    "devices": [
+        "clear_cuda_cache",
+        "get_current_device",
+        "get_device",
+        "get_device_capabilities",
+        "get_device_memory_info",
+        "num_devices",
+        "to_device",
+    ],
+    "dtype": ["get_dtype", "parse_dtype"],
+    "fabric": ["seed_everything_by_time"],
+    "instantiate_utils": [
+        "instantiate",
+        "is_instantiable",
+        "set_print_function_call",
+        "set_print_function_call_permeanent",
+    ],
+    "json": ["load_from_json", "save_to_json", "print_json"],
+    "lazy_state_dict": ["LazyStateDict"],
+    "misc": [
+        "first",
+        "has_length",
+        "join_lists",
+        "validate_and_suggest_corrections",
+    ],
+    "packages": ["compare_versions", "import_object"],
+    "parameters": [
+        "check_parameters_all_equal",
+        "count_parameters",
+        "get_parameter_statistics",
+        "get_parameter_summary",
+        "human_readable",
+        "print_parameters",
+        "state_dict_to_vector",
+        "trainable_state_dict",
+        "vector_to_state_dict",
+    ],
+    "path": [
+        "create_symlink",
+        "listdir_fullpath",
+        "path_is_dir_and_not_empty",
+    ],
+    "pylogger": [
+        "RankedLogger",
+        "RankZeroLogger",
+        "get_rankzero_logger",
+    ],
+    "state_dict_arithmetic": [
+        "ArithmeticStateDict",
+        "state_dicts_check_keys",
+        "num_params_of_state_dict",
+        "state_dict_to_device",
+        "state_dict_flatten",
+        "state_dict_avg",
+        "state_dict_sub",
+        "state_dict_add",
+        "state_dict_add_scalar",
+        "state_dict_mul",
+        "state_dict_div",
+        "state_dict_power",
+        "state_dict_interpolation",
+        "state_dict_sum",
+        "state_dict_weighted_sum",
+        "state_dict_diff_abs",
+        "state_dict_binary_mask",
+        "state_dict_hadamard_product",
+    ],
+    "timer": ["timeit_context"],
+    "type": [
+        "BoolStateDictType",
+        "StateDictType",
+        "TorchModelType",
+    ],
+}
+
+if TYPE_CHECKING:
+    from .cache_utils import cache_to_disk, cache_with_joblib, set_default_cache_dir
+    from .data import (
+        InfiniteDataLoader,
+        load_tensor_from_file,
+        train_validation_split,
+        train_validation_test_split,
+    )
+    from .devices import (
+        clear_cuda_cache,
+        get_current_device,
+        get_device,
+        get_device_capabilities,
+        get_device_memory_info,
+        num_devices,
+        to_device,
+    )
+    from .dtype import get_dtype, parse_dtype
+    from .fabric import seed_everything_by_time
+    from .instantiate_utils import (
+        instantiate,
+        is_instantiable,
+        set_print_function_call,
+        set_print_function_call_permeanent,
+    )
+    from .json import load_from_json, print_json, save_to_json
+    from .lazy_state_dict import LazyStateDict
+    from .misc import first, has_length, join_lists, validate_and_suggest_corrections
+    from .packages import compare_versions, import_object
+    from .parameters import (
+        check_parameters_all_equal,
+        count_parameters,
+        get_parameter_statistics,
+        get_parameter_summary,
+        human_readable,
+        print_parameters,
+        state_dict_to_vector,
+        trainable_state_dict,
+        vector_to_state_dict,
+    )
+    from .path import create_symlink, listdir_fullpath, path_is_dir_and_not_empty
+    from .pylogger import RankedLogger, RankZeroLogger, get_rankzero_logger
+    from .state_dict_arithmetic import (
+        ArithmeticStateDict,
+        num_params_of_state_dict,
+        state_dict_add,
+        state_dict_add_scalar,
+        state_dict_avg,
+        state_dict_binary_mask,
+        state_dict_diff_abs,
+        state_dict_div,
+        state_dict_flatten,
+        state_dict_hadamard_product,
+        state_dict_interpolation,
+        state_dict_mul,
+        state_dict_power,
+        state_dict_sub,
+        state_dict_sum,
+        state_dict_to_device,
+        state_dict_weighted_sum,
+        state_dicts_check_keys,
+    )
+    from .timer import timeit_context
+    from .type import BoolStateDictType, StateDictType, TorchModelType
+
+else:
+    sys.modules[__name__] = LazyImporter(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        extra_objects=_extra_objects,
+    )

fusion_bench/utils/lazy_imports.py

@@ -24,36 +24,78 @@ to publish it as a standalone package.
 import importlib
 import os
 from types import ModuleType
-from typing import Any
+from typing import Any, Dict, List, Optional, Set, Union
 
 
 class LazyImporter(ModuleType):
-    """Do lazy imports."""
+    """Lazy importer for modules and their components.
+
+    This class allows for lazy importing of modules, meaning modules are only
+    imported when they are actually accessed. This can help reduce startup
+    time and memory usage for large packages with many optional dependencies.
+
+    Attributes:
+        _modules: Set of module names available for import.
+        _class_to_module: Mapping from class/function names to their module names.
+        _objects: Dictionary of extra objects to include in the module.
+        _name: Name of the module.
+        _import_structure: Dictionary mapping module names to lists of their exports.
+    """
 
     # Very heavily inspired by optuna.integration._IntegrationModule
     # https://github.com/optuna/optuna/blob/master/optuna/integration/__init__.py
-    def __init__(self, name, module_file, import_structure, extra_objects=None):
+    def __init__(
+        self,
+        name: str,
+        module_file: str,
+        import_structure: Dict[str, List[str]],
+        extra_objects: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Initialize the LazyImporter.
+
+        Args:
+            name: The name of the module.
+            module_file: Path to the module file.
+            import_structure: Dictionary mapping module names to lists of their exports.
+            extra_objects: Optional dictionary of extra objects to include.
+        """
         super().__init__(name)
-        self._modules = set(import_structure.keys())
-        self._class_to_module = {}
+        self._modules: Set[str] = set(import_structure.keys())
+        self._class_to_module: Dict[str, str] = {}
         for key, values in import_structure.items():
             for value in values:
                 self._class_to_module[value] = key
         # Needed for autocompletion in an IDE
-        self.__all__ = list(import_structure.keys()) + sum(
+        self.__all__: List[str] = list(import_structure.keys()) + sum(
             import_structure.values(), []
         )
         self.__file__ = module_file
         self.__path__ = [os.path.dirname(module_file)]
-        self._objects = {} if extra_objects is None else extra_objects
+        self._objects: Dict[str, Any] = {} if extra_objects is None else extra_objects
         self._name = name
         self._import_structure = import_structure
 
     # Needed for autocompletion in an IDE
-    def __dir__(self):
+    def __dir__(self) -> List[str]:
+        """Return list of available attributes for autocompletion.
+
+        Returns:
+            List of all available attribute names.
+        """
         return super().__dir__() + self.__all__
 
     def __getattr__(self, name: str) -> Any:
+        """Get attribute lazily, importing the module if necessary.
+
+        Args:
+            name: The name of the attribute to retrieve.
+
+        Returns:
+            The requested attribute.
+
+        Raises:
+            AttributeError: If the attribute is not found in any module.
+        """
         if name in self._objects:
             return self._objects[name]
         if name in self._modules:
@@ -67,31 +109,68 @@ class LazyImporter(ModuleType):
         setattr(self, name, value)
         return value
 
-    def _get_module(self, module_name: str):
+    def _get_module(self, module_name: str) -> ModuleType:
+        """Import and return the specified module.
+
+        Args:
+            module_name: Name of the module to import.
+
+        Returns:
+            The imported module.
+        """
         return importlib.import_module("." + module_name, self.__name__)
 
-    def __reduce__(self):
+    def __reduce__(self) -> tuple:
+        """Support for pickling the LazyImporter.
+
+        Returns:
+            Tuple containing the class and arguments needed to reconstruct the object.
+        """
         return (self.__class__, (self._name, self.__file__, self._import_structure))
 
 
-class LazyModule(ModuleType):
+class LazyPyModule(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.
+    This can help reduce startup time and memory usage by deferring module imports
+    until they are needed.
 
     Args:
         name: Name of module to apply lazy import.
+
+    Attributes:
+        _name: The name of the module to be lazily imported.
     """
 
     def __init__(self, name: str) -> None:
+        """Initialize the LazyPyModule.
+
+        Args:
+            name: The name of the module to be lazily imported.
+        """
         super().__init__(name)
-        self._name = name
+        self._name: str = name
 
     def _load(self) -> ModuleType:
+        """Load the actual module and update this object's dictionary.
+
+        Returns:
+            The loaded module.
+        """
         module = importlib.import_module(self._name)
         self.__dict__.update(module.__dict__)
         return module
 
     def __getattr__(self, item: str) -> Any:
+        """Get attribute from the lazily loaded module.
+
+        Args:
+            item: The name of the attribute to retrieve.
+
+        Returns:
+            The requested attribute from the loaded module.
+        """
         return getattr(self._load(), item)