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

fusion_bench/taskpool/image_classification.py

@@ -0,0 +1,270 @@
+import itertools
+import json
+import os
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+
+import torch
+from omegaconf import DictConfig, OmegaConf
+from torch import Tensor, nn
+from torch.nn import functional as F
+from torch.utils.data import DataLoader, Dataset
+from torchmetrics import Accuracy, MeanMetric
+from tqdm.auto import tqdm
+
+from fusion_bench import (
+    BaseTaskPool,
+    LightningFabricMixin,
+    RuntimeConstants,
+    auto_register_config,
+    get_rankzero_logger,
+    instantiate,
+)
+from fusion_bench.dataset import ImageClassificationDataset
+from fusion_bench.models.wrappers.switch import set_active_option
+from fusion_bench.tasks.clip_classification import get_classnames, get_num_classes
+from fusion_bench.utils import count_parameters
+
+if TYPE_CHECKING:
+    from transformers import AutoModelForImageClassification
+
+log = get_rankzero_logger(__name__)
+
+
+def _get_logits_from_model_output(outputs) -> Tensor:
+    """Extract logits from model outputs."""
+    match outputs:
+        case Tensor():
+            logits = outputs
+        case dict() | DictConfig() if "logits" in outputs:
+            logits = outputs["logits"]
+            assert isinstance(
+                logits, Tensor
+            ), "The 'logits' in the model output dictionary is not a Tensor."
+        case _:
+            if hasattr(outputs, "logits"):
+                logits = outputs.logits
+                assert isinstance(
+                    logits, Tensor
+                ), "The 'logits' attribute of the model output is not a Tensor."
+            else:
+                raise ValueError(
+                    "Model output is not a Tensor and does not have 'logits' attribute."
+                )
+    return logits
+
+
+@auto_register_config
+class ImageClassificationTaskPool(LightningFabricMixin, BaseTaskPool):
+    _config_mapping = BaseTaskPool._config_mapping | {
+        "_test_datasets": "test_datasets",
+        "_processor": "processor",
+    }
+
+    _processor_instance = None
+    _is_setup: bool = False
+
+    def __init__(
+        self,
+        test_datasets: DictConfig | Dict[str, Any],
+        *,
+        processor: DictConfig | Any,
+        dataloader_kwargs: DictConfig,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        # if the processor is given as a transformers processor instance, store it directly
+        if callable(processor):
+            self._processor_instance = processor
+
+    @property
+    def processor(self) -> Any:
+        if self._processor is None:
+            return None
+
+        if self._processor_instance is not None:
+            return self._processor_instance
+
+        match self._processor:
+            case dict() | DictConfig() if "_target_" in self._processor:
+                self._processor_instance = instantiate(self._processor)
+                return self._processor_instance
+            case str():
+                from transformers import AutoProcessor
+
+                self._processor_instance = AutoProcessor.from_pretrained(
+                    self._processor
+                )
+                return self._processor_instance
+
+        raise ValueError("Processor is not properly configured.")
+
+    def setup(self):
+        # Load test datasets
+        test_datasets = {
+            ds_name: ImageClassificationDataset(
+                self.load_test_dataset(ds_name), processor=self.processor
+            )
+            for ds_name in self._test_datasets
+        }
+        self.test_datasets = test_datasets
+        self.test_dataloaders = {
+            ds_name: self.fabric.setup_dataloaders(
+                self.get_dataloader(ds, stage="test")
+            )
+            for ds_name, ds in test_datasets.items()
+        }
+
+    def load_test_dataset(self, dataset_name: str, *args, **kwargs) -> Dataset:
+        """
+        Load the testing dataset for the specified model.
+
+        Args:
+            dataset_name (str): The name of the model.
+
+        Returns:
+            Dataset: The instantiated testing dataset.
+        """
+        test_dataset = self._test_datasets[dataset_name]
+        if isinstance(test_dataset, (DictConfig, dict)):
+            return instantiate(test_dataset, *args, **kwargs)
+        else:
+            return test_dataset
+
+    def get_dataloader(self, dataset, stage: str):
+        """Create a DataLoader for the specified dataset and training stage.
+
+        Constructs a PyTorch DataLoader with stage-appropriate configurations:
+        - Training stage: shuffling enabled by default
+        - Validation/test stages: shuffling disabled by default
+
+        Args:
+            dataset: The dataset to wrap in a DataLoader.
+            stage (str): Training stage, must be one of "train", "val", or "test".
+                Determines default shuffling behavior.
+
+        Returns:
+            DataLoader: Configured DataLoader for the given dataset and stage.
+        """
+        assert stage in ["train", "val", "test"], f"Invalid stage: {stage}"
+        dataloader_kwargs = dict(self.dataloader_kwargs)
+        if "shuffle" not in dataloader_kwargs:
+            dataloader_kwargs["shuffle"] = stage == "train"
+        return DataLoader(dataset, **dataloader_kwargs)
+
+    @torch.no_grad()
+    def _evaluate(
+        self,
+        classifier,
+        test_loader,
+        num_classes: int,
+        task_name: str = None,
+    ):
+        classifier.eval()
+        accuracy = Accuracy(task="multiclass", num_classes=num_classes)
+        loss_metric = MeanMetric()
+        if RuntimeConstants.debug:
+            log.info("Running under fast_dev_run mode, evaluating on a single batch.")
+            test_loader = itertools.islice(test_loader, 1)
+        else:
+            test_loader = test_loader
+
+        pbar = tqdm(
+            test_loader,
+            desc=f"Evaluating {task_name}" if task_name is not None else "Evaluating",
+            leave=False,
+            dynamic_ncols=True,
+        )
+        for batch in pbar:
+            inputs, targets = batch
+            outputs = classifier(inputs)
+            logits = _get_logits_from_model_output(outputs)
+            if logits.device != targets.device:
+                targets = targets.to(logits.device)
+
+            loss = F.cross_entropy(logits, targets)
+            loss_metric.update(loss.detach().cpu())
+            acc = accuracy(logits.detach().cpu(), targets.detach().cpu())
+            pbar.set_postfix(
+                {
+                    "accuracy": accuracy.compute().item(),
+                    "loss": loss_metric.compute().item(),
+                }
+            )
+
+        acc = accuracy.compute().item()
+        loss = loss_metric.compute().item()
+        results = {"accuracy": acc, "loss": loss}
+        return results
+
+    def evaluate(
+        self,
+        model: Union["AutoModelForImageClassification", nn.Module],
+        name: str = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        assert isinstance(
+            model, nn.Module
+        ), f"Expected model to be an instance of nn.Module, but got {type(model)}"
+
+        if not self._is_setup:
+            self.setup()
+
+        classifier = self.fabric.to_device(model)
+        classifier.eval()
+        report = {}
+        # collect basic model information
+        training_params, all_params = count_parameters(model)
+        report["model_info"] = {
+            "trainable_params": training_params,
+            "all_params": all_params,
+            "trainable_percentage": training_params / all_params,
+        }
+        if name is not None:
+            report["model_info"]["name"] = name
+
+        # evaluate on each task
+        pbar = tqdm(
+            self.test_dataloaders.items(),
+            desc="Evaluating tasks",
+            total=len(self.test_dataloaders),
+        )
+        for task_name, test_dataloader in pbar:
+            set_active_option(classifier, task_name)
+            num_classes = get_num_classes(task_name)
+            result = self._evaluate(
+                classifier,
+                test_dataloader,
+                num_classes=num_classes,
+                task_name=task_name,
+            )
+            report[task_name] = result
+
+        # calculate the average accuracy and loss
+        if "average" not in report:
+            report["average"] = {}
+            accuracies = [
+                value["accuracy"]
+                for key, value in report.items()
+                if "accuracy" in value
+            ]
+            if len(accuracies) > 0:
+                average_accuracy = sum(accuracies) / len(accuracies)
+                report["average"]["accuracy"] = average_accuracy
+            losses = [value["loss"] for key, value in report.items() if "loss" in value]
+            if len(losses) > 0:
+                average_loss = sum(losses) / len(losses)
+                report["average"]["loss"] = average_loss
+
+        log.info(f"Evaluation Result: {report}")
+        if self.fabric.is_global_zero and len(self.fabric._loggers) > 0:
+            save_path = os.path.join(self.log_dir, "report.json")
+            for version in itertools.count(1):
+                if not os.path.exists(save_path):
+                    break
+                # if the file already exists, increment the version to avoid overwriting
+                save_path = os.path.join(self.log_dir, f"report_{version}.json")
+            with open(save_path, "w") as fp:
+                json.dump(report, fp)
+            log.info(f"Evaluation report saved to {save_path}")
+        return report

fusion_bench/utils/__init__.py

@@ -31,6 +31,11 @@ _import_structure = {
     ],
     "dtype": ["get_dtype", "parse_dtype"],
     "fabric": ["seed_everything_by_time"],
+    "hydra_utils": [
+        "initialize_hydra_config",
+        "get_default_config_path",
+        "get_hydra_output_dir",
+    ],
     "instantiate_utils": [
         "instantiate",
         "is_instantiable",
@@ -40,6 +45,7 @@ _import_structure = {
     "json": ["load_from_json", "save_to_json", "print_json"],
     "lazy_state_dict": ["LazyStateDict"],
     "misc": [
+        "DeprecationWarningMeta",
         "first",
         "has_length",
         "join_lists",
@@ -53,6 +59,7 @@ _import_structure = {
         "get_parameter_summary",
         "human_readable",
         "print_parameters",
+        "print_trainable_parameters",
         "state_dict_to_vector",
         "trainable_state_dict",
         "vector_to_state_dict",
@@ -121,6 +128,11 @@ if TYPE_CHECKING:
     )
     from .dtype import get_dtype, parse_dtype
     from .fabric import seed_everything_by_time
+    from .hydra_utils import (
+        get_default_config_path,
+        get_hydra_output_dir,
+        initialize_hydra_config,
+    )
     from .instantiate_utils import (
         instantiate,
         is_instantiable,
@@ -129,7 +141,13 @@ if TYPE_CHECKING:
     )
     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 .misc import (
+        DeprecationWarningMeta,
+        first,
+        has_length,
+        join_lists,
+        validate_and_suggest_corrections,
+    )
     from .packages import compare_versions, import_object
     from .parameters import (
         check_parameters_all_equal,
@@ -138,6 +156,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/data.py

@@ -95,7 +95,7 @@ class InfiniteDataLoader:
             f"Failed to retrieve data from data loader after {self.max_retries} attempts. "
             f"Last error: [{type(last_exception).__name__}]{last_exception}. "
             + (
-                f"The data loader appears to be empty."
+                f"The data loader may be empty."
                 if isinstance(last_exception, StopIteration)
                 else ""
             )

fusion_bench/utils/dict.py

@@ -41,3 +41,22 @@ def dict_map(f, d: dict, *, max_level: int = -1, skip_levels=0, inplace=False):
 
     dict_map_impl(d, ans, 0)
     return ans
+
+
+def dict_merge(dicts: Iterable[dict], disjoint: bool = True) -> dict:
+    """Merge multiple dictionaries into one.
+
+    Args:
+        dicts (Iterable[dict]): iterable of dictionaries to merge
+        disjoint (bool, optional): if True, raises an error on key conflicts. Defaults to True.
+
+    Returns:
+        dict: merged dictionary
+    """
+    merged_dict = type(dicts[0])()
+    for d in dicts:
+        for k, v in d.items():
+            if disjoint and k in merged_dict:
+                raise ValueError(f"Key conflict when merging dictionaries: {k}")
+            merged_dict[k] = v
+    return merged_dict

fusion_bench/utils/dtype.py

@@ -146,3 +146,22 @@ def validate_expected_param_dtype(
             raise ValueError(
                 f"Parameter {name} has dtype {param.dtype}, but expected {dtype}"
             )
+
+
+def dtype_support_svd(dtype: torch.dtype) -> bool:
+    """
+    Check if the given dtype is supported for SVD operation in PyTorch.
+
+    Args:
+        dtype (torch.dtype): The data type to check.
+
+    Returns:
+        bool: True if the dtype is supported for SVD, False otherwise.
+    """
+    supported_dtypes = {
+        torch.float32,
+        torch.float64,
+        torch.complex64,
+        torch.complex128,
+    }
+    return dtype in supported_dtypes

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

@@ -9,6 +9,7 @@ __all__ = [
     "join_lists",
     "attr_equal",
     "validate_and_suggest_corrections",
+    "DeprecationWarningMeta",
 ]
 
 

fusion_bench/utils/packages.py

@@ -20,6 +20,10 @@ def _get_package_version(name: str) -> "Version":
         return version.parse("0.0.0")
 
 
+def is_ray_available():
+    return _is_package_available("ray")
+
+
 def is_pyav_available():
     return _is_package_available("av")
 

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:

fusion_bench/utils/rich_utils.py

@@ -93,11 +93,11 @@ def print_bordered(
     Print a message with a colored border.
 
     Args:
-    message (str): The message to print.
-    title (str, optional): The title of the panel. Defaults to None.
-    style (str, optional): The color style for the border. Defaults to "cyan".
-    code_style (str, optional): The syntax highlighting style if the message is code.
-                                Set to None for plain text. Defaults to "python".
+        message (str): The message to print.
+        title (str, optional): The title of the panel. Defaults to None.
+        style (str, optional): The color style for the border. Defaults to "cyan".
+        code_style (str, optional): The syntax highlighting style if the message is code.
+                                    Set to None for plain text. Defaults to "python".
     """
     if code_style:
         if format_code:
@@ -168,7 +168,7 @@ def print_config_tree(
         "callbacks",
         "logger",
         "trainer",
-        "paths",
+        "path",
         "extras",
     ),
     resolve: bool = False,
@@ -179,11 +179,20 @@ def print_config_tree(
 ) -> None:
     """Prints the contents of a DictConfig as a tree structure using the Rich library.
 
-    :param cfg: A DictConfig composed by Hydra.
-    :param print_order: Determines in what order config components are printed. Default is ``("data", "model",
-    "callbacks", "logger", "trainer", "paths", "extras")``.
-    :param resolve: Whether to resolve reference fields of DictConfig. Default is ``False``.
-    :param save_to_file: Whether to export config to the hydra output folder. Default is ``False``.
+    Args:
+        cfg (DictConfig): A DictConfig composed by Hydra.
+        print_order (Sequence[str], optional): Determines in what order config components are printed.
+            Defaults to ``("data", "model", "callbacks", "logger", "trainer", "paths", "extras")``.
+        resolve (bool, optional): Whether to resolve reference fields of DictConfig.
+            Defaults to ``False``.
+        save_to_file (bool, optional): Whether to export config to the hydra output folder.
+            Defaults to ``False``.
+        theme (str, optional): The theme to use for syntax highlighting. Defaults to "monokai".
+        background_color (str, optional): The background color to use for syntax highlighting.
+            Defaults to "default".
+
+    Returns:
+        None
     """
     style = "tree"
     tree = rich.tree.Tree("CONFIG", style=style, guide_style=style)
@@ -200,18 +209,13 @@ def print_config_tree(
             )
         )
 
-    # add all the other fields to queue (not specified in `print_order`)
-    for field in cfg:
-        if field not in queue:
-            queue.append(field)
-
     # generate config tree from queue
     for field in queue:
         branch = tree.add(field, style=style, guide_style=style)
 
         config_group = cfg[field]
         if isinstance(config_group, DictConfig):
-            branch_content = OmegaConf.to_yaml(config_group, resolve=resolve)
+            branch_content = OmegaConf.to_yaml(config_group, resolve=resolve).strip()
         else:
             branch_content = str(config_group)
 
@@ -224,13 +228,32 @@ def print_config_tree(
             )
         )
 
+    # add all the other fields to queue (not specified in `print_order`)
+    other_fields = [field for field in cfg if field not in queue]
+    if other_fields:
+        others_branch = tree.add(Text("[others]"), style=style, guide_style=style)
+
+        other_cfg = OmegaConf.create({field: cfg[field] for field in other_fields})
+        branch_content = OmegaConf.to_yaml(other_cfg, resolve=resolve).strip()
+
+        others_branch.add(
+            rich.syntax.Syntax(
+                branch_content, "yaml", theme=theme, background_color=background_color
+            )
+        )
+
     # print config tree
     rich.print(tree)
 
     # save config tree to file
     if save_to_file:
-        with open(Path(cfg.paths.output_dir, "config_tree.log"), "w") as file:
-            rich.print(tree, file=file)
+        if not cfg.get("paths") or not cfg.paths.get("output_dir"):
+            log.error(
+                "Cannot save config tree to file. 'paths.output_dir' is not specified in the config."
+            )
+        else:
+            with open(Path(cfg.path.output_dir, "config_tree.log"), "w") as file:
+                rich.print(tree, file=file)
 
 
 @rank_zero_only