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

fusion_bench/taskpool/clip_vision/taskpool.py

@@ -27,7 +27,7 @@ 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 import RuntimeConstants, auto_register_config
 from fusion_bench.dataset import CLIPDataset
 from fusion_bench.mixins import HydraConfigMixin, LightningFabricMixin
 from fusion_bench.models.hf_clip import HFCLIPClassifier
@@ -86,6 +86,7 @@ class LayerWiseFeatureSaver:
             torch.save(features, self.save_path)
 
 
+@auto_register_config
 class CLIPVisionModelTaskPool(
     HydraConfigMixin,
     LightningFabricMixin,
@@ -134,11 +135,13 @@ class CLIPVisionModelTaskPool(
         layer_wise_feature_first_token_only: bool = True,
         layer_wise_feature_max_num: Optional[int] = None,
         fast_dev_run: Optional[bool] = None,
+        move_to_device: bool = True,
         **kwargs,
     ):
         """
         Initialize the CLIPVisionModelTaskPool.
         """
+        super().__init__(**kwargs)
         self._test_datasets = test_datasets
         self._processor = processor
         self._data_processor = data_processor
@@ -159,7 +162,6 @@ class CLIPVisionModelTaskPool(
             self.fast_dev_run = RuntimeConstants().debug
         else:
             self.fast_dev_run = fast_dev_run
-        super().__init__(**kwargs)
 
     def setup(self):
         """
@@ -220,7 +222,9 @@ class CLIPVisionModelTaskPool(
             for name, dataset in self.test_datasets.items()
         }
         self.test_dataloaders = {
-            name: self.fabric.setup_dataloaders(dataloader)
+            name: self.fabric.setup_dataloaders(
+                dataloader, move_to_device=self.move_to_device
+            )
             for name, dataloader in self.test_dataloaders.items()
         }
 
@@ -273,6 +277,8 @@ class CLIPVisionModelTaskPool(
                 task_name=task_name,
             )
             logits: Tensor = outputs["logits"]
+            if logits.device != targets.device:
+                targets = targets.to(logits.device)
 
             loss = F.cross_entropy(logits, targets)
             loss_metric.update(loss.detach().cpu())
@@ -321,7 +327,8 @@ class CLIPVisionModelTaskPool(
             self.clip_model,
             processor=self.processor,
         )
-        classifier = cast(HFCLIPClassifier, self.fabric.to_device(classifier))
+        if self.move_to_device:
+            classifier = cast(HFCLIPClassifier, self.fabric.to_device(classifier))
         # collect basic model information
         training_params, all_params = count_parameters(model)
         report["model_info"] = {

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/devices.py

@@ -39,7 +39,12 @@ def clear_cuda_cache():
         log.warning("CUDA is not available. No cache to clear.")
 
 
-def to_device(obj: T, device: Optional[torch.device], **kwargs: Any) -> T:
+def to_device(
+    obj: T,
+    device: Optional[torch.device],
+    copy_on_move: bool = False,
+    **kwargs: Any,
+) -> T:
     """
     Move a given object to the specified device.
 
@@ -49,12 +54,20 @@ def to_device(obj: T, device: Optional[torch.device], **kwargs: Any) -> T:
     Args:
         obj: The object to be moved to the device. This can be a torch.Tensor, torch.nn.Module, list, tuple, or dict.
         device (torch.device): The target device to move the object to. This can be `None`.
-        **kwargs: Additional keyword arguments to be passed to the `to` method of torch.Tensor or torch.nn.Module. For example, `non_blocking=True`, `dtype=torch.float16`.
+        copy_on_move (bool, optional): Whether to force a copy operation when moving tensors to a different device.
+            If True, tensors will be copied when moved to a different device (copy=True is passed to tensor.to()).
+            If False (default), tensors are moved without forcing a copy operation, allowing PyTorch to optimize
+            the operation. This parameter only affects torch.Tensor objects; modules and other types are unaffected.
+            Defaults to False.
+        **kwargs: Additional keyword arguments to be passed to the `to` method of torch.Tensor or torch.nn.Module.
+            For example, `non_blocking=True`, `dtype=torch.float16`. Note that if `copy_on_move=True`, the `copy`
+            keyword argument will be automatically set and should not be provided manually.
 
     Returns:
         The object moved to the specified device. The type of the returned object matches the type of the input object.
 
     Examples:
+        ```python
         >>> tensor = torch.tensor([1, 2, 3])
         >>> to_device(tensor, torch.device('cuda'))
         tensor([1, 2, 3], device='cuda:0')
@@ -66,17 +79,26 @@ def to_device(obj: T, device: Optional[torch.device], **kwargs: Any) -> T:
         >>> data = [torch.tensor([1, 2]), torch.tensor([3, 4])]
         >>> to_device(data, torch.device('cuda'))
         [tensor([1, 2], device='cuda:0'), tensor([3, 4], device='cuda:0')]
+
+        >>> # Force copy when moving to different device
+        >>> tensor = torch.tensor([1, 2, 3], device='cpu')
+        >>> copied_tensor = to_device(tensor, torch.device('cuda'), copy_on_move=True)
+        >>> # tensor and copied_tensor will have different memory locations
+        ```
     """
-    if isinstance(obj, (torch.Tensor, torch.nn.Module)):
+    if isinstance(obj, torch.Tensor):
+        if copy_on_move:
+            if obj.device != torch.device(device):
+                kwargs["copy"] = True
+        return obj.to(device, **kwargs)
+    elif isinstance(obj, torch.nn.Module):
         return obj.to(device, **kwargs)
     elif isinstance(obj, list):
-        return [to_device(o, device) for o in obj]
+        return [to_device(o, device, **kwargs) for o in obj]
     elif isinstance(obj, tuple):
-        return tuple(to_device(o, device) for o in obj)
+        return tuple(to_device(o, device, **kwargs) for o in obj)
     elif isinstance(obj, dict):
-        for key in obj:
-            obj[key] = to_device(obj[key], device)
-        return obj
+        return {key: to_device(value, device, **kwargs) for key, value in obj.items()}
     else:
         # the default behavior is to return the object as is
         return obj

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)