fusion-bench 0.2.31__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",
@@ -122,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,
@@ -130,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,

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

@@ -1,6 +1,16 @@
 from collections import OrderedDict
 from numbers import Number
-from typing import Callable, Dict, List, Literal, Optional, Union, cast
+from typing import (
+    Callable,
+    Dict,
+    Iterator,
+    List,
+    Literal,
+    Mapping,
+    Optional,
+    Union,
+    cast,
+)
 
 import torch
 from torch import Tensor
@@ -462,6 +472,118 @@ class ArithmeticStateDict(OrderedDict):
         return cls(result_dict)
 
 
+class LazyStateDictExpr(Mapping[str, torch.Tensor]):
+    """
+    A lazy, key-wise expression over state_dict-like objects.
+    """
+
+    # ---- core Mapping API ----
+    def __getitem__(self, key: str) -> torch.Tensor:
+        raise NotImplementedError
+
+    def __iter__(self) -> Iterator[str]:
+        raise NotImplementedError
+
+    def __len__(self) -> int:
+        raise NotImplementedError
+
+    # ---- arithmetic (build graph only) ----
+    def __add__(self, other):
+        return BinaryOp(torch.add, self, ensure_expr(other))
+
+    def __sub__(self, other):
+        return BinaryOp(torch.sub, self, ensure_expr(other))
+
+    def __mul__(self, scalar):
+        return UnaryOp(lambda x: x * scalar, self)
+
+    def __rmul__(self, scalar):
+        return self.__mul__(scalar)
+
+    def __truediv__(self, scalar):
+        return UnaryOp(lambda x: x / scalar, self)
+
+    # ---- eager escape hatch ----
+    def materialize(
+        self, device=None, dtype=None, non_blocking=False, copy=False
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Eagerly evaluate into an OrderedDict.
+        """
+        out = {}
+        for k in self:
+            v = self[k]
+            out[k] = v.to(
+                device=device,
+                dtype=dtype,
+                non_blocking=non_blocking,
+                copy=copy,
+            )
+        return out
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(lazy)"
+
+
+class StateDictLeaf(LazyStateDictExpr):
+    def __init__(self, state_dict: Mapping[str, torch.Tensor]):
+        self._sd = state_dict
+
+    def __getitem__(self, key: str) -> torch.Tensor:
+        return self._sd[key]
+
+    def __iter__(self):
+        return iter(self._sd)
+
+    def __len__(self):
+        return len(self._sd)
+
+
+class UnaryOp(LazyStateDictExpr):
+    def __init__(self, op: Callable[[torch.Tensor], torch.Tensor], child):
+        self.op = op
+        self.child = child
+
+    def __getitem__(self, key: str):
+        return self.op(self.child[key])
+
+    def __iter__(self):
+        return iter(self.child)
+
+    def __len__(self):
+        return len(self.child)
+
+
+class BinaryOp(LazyStateDictExpr):
+    def __init__(
+        self,
+        op: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+        left,
+        right,
+    ):
+        self.op = op
+        self.left = left
+        self.right = right
+
+    def __getitem__(self, key: str):
+        return self.op(self.left[key], self.right[key])
+
+    def __iter__(self):
+        # assume key sets are aligned
+        return iter(self.left)
+
+    def __len__(self):
+        return len(self.left)
+
+
+def ensure_expr(x):
+    if isinstance(x, LazyStateDictExpr):
+        return x
+    if isinstance(x, Mapping):
+        return StateDictLeaf(x)
+    raise TypeError(f"Unsupported operand type: {type(x)}")
+
+
 def _validate_state_dict_list_not_empty(state_dicts: List[StateDictType]) -> None:
     """
     Validate that the list of state dicts is not empty and contains valid state dicts.
@@ -1228,3 +1350,63 @@ def state_dict_hadamard_product(a: StateDictType, b: StateDictType) -> StateDict
     """
     _validate_state_dict_same_keys([a, b])
     return OrderedDict((key, a[key] * b[key]) for key in a)
+
+
+def state_dict_max(
+    state_dicts: List[StateDictType],
+) -> StateDictType:
+    """
+    Compute the element-wise maximum across multiple state dicts.
+
+    Args:
+        state_dicts: List of state dicts to compute the maximum from.
+
+    Returns:
+        A state dict containing the element-wise maximums.
+    """
+    _validate_state_dict_list_not_empty(state_dicts)
+    _validate_state_dict_same_keys(state_dicts)
+
+    max_state_dict = OrderedDict()
+
+    for key in state_dicts[0]:
+        # Initialize with the first tensor
+        max_tensor = state_dicts[0][key].clone()
+
+        # Compute element-wise maximum
+        for state_dict in state_dicts[1:]:
+            max_tensor = torch.max(max_tensor, state_dict[key])
+
+        max_state_dict[key] = max_tensor
+
+    return max_state_dict
+
+
+def state_dict_max_abs(
+    state_dicts: List[StateDictType],
+) -> StateDictType:
+    """
+    Compute the element-wise maximum absolute value across multiple state dicts.
+
+    Args:
+        state_dicts: List of state dicts to compute the maximum absolute values from.
+
+    Returns:
+        A state dict containing the element-wise maximum absolute values.
+    """
+    _validate_state_dict_list_not_empty(state_dicts)
+    _validate_state_dict_same_keys(state_dicts)
+
+    max_abs_state_dict = OrderedDict()
+
+    for key in state_dicts[0]:
+        # Initialize with the absolute values of the first tensor
+        max_abs_tensor = state_dicts[0][key].abs()
+
+        # Compute element-wise maximum absolute value
+        for state_dict in state_dicts[1:]:
+            max_abs_tensor = torch.max(max_abs_tensor, state_dict[key].abs())
+
+        max_abs_state_dict[key] = max_abs_tensor
+
+    return max_abs_state_dict

fusion_bench/utils/tensorboard.py

@@ -2,14 +2,18 @@
 functions deal with tensorboard logs.
 """
 
-from typing import Dict, Iterable, List
+from pathlib import Path
+from typing import Dict, Iterable, List, Union
 
 import numpy as np
 import pandas as pd
 from tensorboard.backend.event_processing import event_accumulator
 
 
-def parse_tensorboard_as_dict(path: str, scalars: Iterable[str]):
+def parse_tensorboard_as_dict(
+    path: Union[str, Path],
+    scalars: Iterable[str],
+) -> Dict[str, pd.DataFrame]:
     """
     returns a dictionary of pandas dataframes for each requested scalar.
 
@@ -20,7 +24,19 @@ def parse_tensorboard_as_dict(path: str, scalars: Iterable[str]):
 
     Returns:
         Dict[str, pandas.DataFrame]: a dictionary of pandas dataframes for each requested scalar
+
+    Example:
+
+        >>> from fusion_bench.utils.tensorboard import parse_tensorboard_as_dict
+        >>> path = "path/to/tensorboard/logs"
+        >>> scalars = ["train/loss", "val/accuracy"]
+        >>> data = parse_tensorboard_as_dict(path, scalars)
+        >>> train_loss_df = data["train/loss"]
+        >>> val_accuracy_df = data["val/accuracy"]
     """
+    if isinstance(path, Path):
+        path = str(path)
+    assert isinstance(path, str), "path must be a string"
     ea = event_accumulator.EventAccumulator(
         path,
         size_guidance={event_accumulator.SCALARS: 0},
@@ -33,7 +49,9 @@ def parse_tensorboard_as_dict(path: str, scalars: Iterable[str]):
     return {k: pd.DataFrame(ea.Scalars(k)) for k in scalars}
 
 
-def parse_tensorboard_as_list(path: str, scalars: Iterable[str]):
+def parse_tensorboard_as_list(
+    path: Union[str, Path], scalars: Iterable[str]
+) -> List[pd.DataFrame]:
     """
     returns a list of pandas dataframes for each requested scalar.
 

{fusion_bench-0.2.31.dist-info → fusion_bench-0.2.32.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fusion-bench
-Version: 0.2.31
+Version: 0.2.32
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
 Project-URL: Repository, https://github.com/tanganke/fusion_bench
@@ -61,6 +61,8 @@ Dynamic: license-file
 
 FusionBench is a benchmark suite designed to evaluate the performance of various deep model fusion techniques. It aims to provide a comprehensive comparison of different methods on a variety of datasets and tasks.
 
+## :newspaper: News and Related
+
 Projects based on FusionBench and news from the community (descending order of date. If you have any work based on FusionBench, please feel free to let us know, we are willing to add it to the list. :partying_face:):
 
 <details>