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

fusion_bench/taskpool/base_pool.py

@@ -5,33 +5,115 @@ from fusion_bench.mixins import BaseYAMLSerializable
 
 
 class BaseTaskPool(BaseYAMLSerializable):
+    """Abstract base class for task pools in the FusionBench framework.
+
+    A task pool represents a collection of evaluation tasks that can be used to
+    assess model performance across multiple benchmarks or datasets. This base
+    class defines the common interface that all task pool implementations must
+    follow, ensuring consistency across different task types and evaluation
+    scenarios.
+
+    Task pools are designed to be configurable through YAML files and can be
+    used in various model fusion and evaluation workflows. They provide a
+    standardized way to evaluate models on multiple tasks and aggregate results.
+
+    The class inherits from BaseYAMLSerializable to support configuration
+    management and serialization capabilities.
+
+    Attributes:
+        _program: Optional program reference for execution context.
+        _config_key: Configuration key used for YAML configuration ("taskpool").
+
+    Abstract Methods:
+        evaluate: Must be implemented by subclasses to define task-specific
+            evaluation logic.
+
+    Example:
+        Implementing a custom task pool:
+
+        ```python
+        class MyTaskPool(BaseTaskPool):
+
+
+            def evaluate(self, model, **kwargs):
+                results = {}
+                for task_name in self.tasks:
+                    # Implement task-specific evaluation
+                    results[task_name] = self._evaluate_task(model, task_name)
+                return results
+        ```
+    """
+
     _program = None
     _config_key = "taskpool"
 
     @abstractmethod
     def evaluate(self, model: Any, *args: Any, **kwargs: Any) -> Dict[str, Any]:
-        """
-        Evaluate the model on all tasks in the task pool, and return a report.
+        """Evaluate a model on all tasks in the task pool and return aggregated results.
 
-        Take image classification as an example, the report will look like:
+        This abstract method defines the core evaluation interface that all task pool
+        implementations must provide. It should evaluate the given model on all tasks
+        managed by the pool and return a structured report of the results.
 
-        ```python
-        {
-            "mnist": {
-                "accuracy": 0.8,
-                "loss": 0.2,
-            },
-            <task_name>: {
-                <metric_name>: <metric_value>,
-                ...
-            },
-        }
-        ```
+        The evaluation process typically involves:
+        1. Iterating through all tasks in the pool
+        2. Running model inference on each task's dataset
+        3. Computing task-specific metrics
+        4. Aggregating results into a standardized report format
 
         Args:
-            model: The model to evaluate.
+            model: The model to evaluate. Can be any model type (PyTorch model,
+                Hugging Face model, etc.) that is compatible with the specific
+                task pool implementation.
+            *args: Additional positional arguments that may be needed for
+                task-specific evaluation procedures.
+            **kwargs: Additional keyword arguments for evaluation configuration,
+                such as batch_size, device, evaluation metrics, etc.
 
         Returns:
-            report (dict): A dictionary containing the results of the evaluation for each task.
+            Dict[str, Any]: A dictionary containing evaluation results for each task.
+                The structure follows the pattern:
+
+                ```python
+                {
+                    "task_name_1": {
+                        "metric_1": value,
+                        "metric_2": value,
+                        ...
+                    },
+                    "task_name_2": {
+                        "metric_1": value,
+                        "metric_2": value,
+                        ...
+                    },
+                    ...
+                }
+                ```
+
+        Example:
+            For an image classification task pool:
+
+            ```python
+            results = task_pool.evaluate(model)
+            # Returns:
+            # {
+            #     "mnist": {
+            #         "accuracy": 0.95,
+            #         "loss": 0.15,
+            #     },
+            #     "cifar10": {
+            #         "accuracy": 0.87,
+            #         "loss": 0.42,
+            #     }
+            # }
+            ```
+
+        Raises:
+            NotImplementedError: This method must be implemented by subclasses.
+
+        Note:
+            Implementations should ensure that the returned dictionary structure
+            is consistent and that metric names are standardized across similar
+            task types to enable meaningful comparison and aggregation.
         """
         pass

fusion_bench/taskpool/clip_vision/taskpool.py

@@ -27,8 +27,9 @@ 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.dataset import CLIPDataset
-from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.mixins import HydraConfigMixin, LightningFabricMixin
 from fusion_bench.models.hf_clip import HFCLIPClassifier
 from fusion_bench.taskpool import BaseTaskPool
 from fusion_bench.tasks.clip_classification import get_classnames_and_templates
@@ -86,8 +87,9 @@ class LayerWiseFeatureSaver:
 
 
 class CLIPVisionModelTaskPool(
-    BaseTaskPool,
+    HydraConfigMixin,
     LightningFabricMixin,
+    BaseTaskPool,
 ):
     """
     This class is used to define the image classification task for CLIP models.
@@ -131,7 +133,7 @@ class CLIPVisionModelTaskPool(
         layer_wise_feature_save_path: Optional[str] = None,
         layer_wise_feature_first_token_only: bool = True,
         layer_wise_feature_max_num: Optional[int] = None,
-        fast_dev_run: bool = False,
+        fast_dev_run: Optional[bool] = None,
         **kwargs,
     ):
         """
@@ -153,7 +155,10 @@ class CLIPVisionModelTaskPool(
         self.layer_wise_feature_first_token_only = layer_wise_feature_first_token_only
         self.layer_wise_feature_max_num = layer_wise_feature_max_num
 
-        self.fast_dev_run = fast_dev_run
+        if self.fast_dev_run is None:
+            self.fast_dev_run = RuntimeConstants().debug
+        else:
+            self.fast_dev_run = fast_dev_run
         super().__init__(**kwargs)
 
     def setup(self):
@@ -304,7 +309,7 @@ class CLIPVisionModelTaskPool(
             self.setup()
 
         report = {}
-        # CLIPVisionModel works the same with CLIPVisonTransformer, so we can use it directly
+        # CLIPVisionModel works the same with CLIPVisionTransformer, so we can use it directly
         if hasattr(model, "is_surgery_model") and model.is_surgery_model:
             log.info("running evaluation on a surgery model.")
             model: "SurgeryModelWrapper" = model

fusion_bench/taskpool/dummy.py

@@ -1,5 +1,10 @@
 """
-This is the dummy task pool that is used for debugging purposes.
+Dummy task pool implementation for debugging and testing purposes.
+
+This module provides a minimal task pool implementation that can be used for
+debugging model fusion workflows, testing infrastructure, and validating model
+architectures without running expensive evaluation procedures. It's particularly
+useful during development and prototyping phases.
 """
 
 from typing import Optional
@@ -14,14 +19,41 @@ from fusion_bench.utils.parameters import count_parameters, print_parameters
 
 
 def get_model_summary(model: nn.Module) -> dict:
-    """
-    Generate a report for the given model.
+    """Generate a comprehensive summary report for a PyTorch model.
+
+    Analyzes the given model to extract key information about its architecture,
+    parameter count, and training characteristics. This function is useful for
+    model introspection and comparative analysis during model fusion workflows.
+
+    The summary includes both trainable and total parameter counts, which helps
+    in understanding model complexity and memory requirements. The trainable
+    percentage is particularly useful for identifying models with frozen layers
+    or parameter-efficient fine-tuning setups.
 
     Args:
-        model: The model to generate the report for.
+        model: The PyTorch model to analyze. Can be any nn.Module instance
+            including complex models, fusion models, or pre-trained models.
 
     Returns:
-        dict: The generated report.
+        dict: A structured report containing model information:
+            - model_info: Dictionary with parameter statistics
+                - trainable_params: Number of trainable parameters
+                - all_params: Total number of parameters (trainable + frozen)
+                - trainable_percentage: Ratio of trainable to total parameters
+
+    Example:
+        ```python
+        >>> model = MyModel()
+        >>> summary = get_model_summary(model)
+        >>> print(summary)
+        {
+            "model_info": {
+                "trainable_params": 1234567,
+                "all_params": 1234567,
+                "trainable_percentage": 1.0
+            }
+        }
+        ```
     """
     report = {}
     training_params, all_params = count_parameters(model)
@@ -34,21 +66,77 @@ def get_model_summary(model: nn.Module) -> dict:
 
 
 class DummyTaskPool(BaseTaskPool):
+    """A lightweight task pool implementation for debugging and development workflows.
+
+    This dummy task pool provides a minimal evaluation interface that focuses on
+    model introspection rather than task-specific performance evaluation. It's
+    designed for development scenarios where you need to test model fusion
+    pipelines, validate architectures, or debug workflows without the overhead
+    of running actual evaluation tasks.
+
+    The task pool is particularly useful when:
+        - You want to verify model fusion works correctly
+        - You need to check parameter counts after fusion
+        - You're developing new fusion algorithms
+        - You want to test infrastructure without expensive evaluations
+
+    Example:
+        ```python
+        >>> taskpool = DummyTaskPool(model_save_path="/tmp/fused_model")
+        >>> results = taskpool.evaluate(fused_model)
+        >>> print(f"Model has {results['model_info']['trainable_params']} parameters")
+        ```
     """
-    This is a dummy task pool used for debugging purposes. It inherits from the base TaskPool class.
-    """
 
-    def __init__(self, model_save_path: Optional[str] = None):
-        super().__init__()
+    def __init__(self, model_save_path: Optional[str] = None, **kwargs):
+        """Initialize the dummy task pool with optional model saving capability.
+
+        Args:
+            model_save_path: Optional path where the evaluated model should be saved.
+                If provided, the model will be serialized and saved to this location
+                after evaluation using the separate_save utility. If None, no model
+                saving will be performed.
+
+        Example:
+            ```python
+            >>> # Create taskpool without saving
+            >>> taskpool = DummyTaskPool()
+
+            >>> # Create taskpool with model saving
+            >>> taskpool = DummyTaskPool(model_save_path="/path/to/save/model.pth")
+            ```
+        """
+        super().__init__(**kwargs)
         self.model_save_path = model_save_path
 
     def evaluate(self, model):
-        """
-        Evaluate the given model.
-        This method does nothing but print the parameters of the model in a human-readable format.
+        """Perform lightweight evaluation and analysis of the given model.
+
+        This method provides a minimal evaluation that focuses on model introspection
+        rather than task-specific performance metrics. It performs parameter analysis,
+        optionally saves the model, and returns a summary report.
+
+        The evaluation process includes:
+        1. Printing human-readable parameter information (rank-zero only)
+        2. Optionally saving the model if a save path was configured
+        3. Generating and returning a model summary report
 
         Args:
-            model: The model to evaluate.
+            model: The model to evaluate. Can be any PyTorch nn.Module including
+                fusion models, pre-trained models, or custom architectures.
+
+        Returns:
+            dict: A model summary report containing parameter statistics and
+                architecture information. See get_model_summary() for detailed
+                format specification.
+
+        Example:
+            ```python
+            >>> taskpool = DummyTaskPool(model_save_path="/tmp/model.pth")
+            >>> model = torch.nn.Linear(10, 5)
+            >>> results = taskpool.evaluate(model)
+            >>> print(f"Trainable params: {results['model_info']['trainable_params']}")
+            ```
         """
         if rank_zero_only.rank == 0:
             print_parameters(model, is_human_readable=True)

fusion_bench/taskpool/lm_eval_harness/taskpool.py

@@ -16,6 +16,47 @@ log = logging.getLogger(__name__)
 
 
 class LMEvalHarnessTaskPool(BaseTaskPool, LightningFabricMixin):
+    """A task pool implementation that interfaces with the LM Evaluation Harness framework.
+
+    This class provides a wrapper around the LM Evaluation Harness (lm-eval) library,
+    enabling evaluation of language models on various standardized benchmarks and tasks.
+    It inherits from BaseTaskPool and LightningFabricMixin to provide distributed
+    computing capabilities through PyTorch Lightning Fabric.
+
+    The task pool supports evaluation on multiple tasks simultaneously and provides
+    flexible configuration options for batch processing, output formatting, and
+    logging. It automatically handles model setup and wrapping for distributed
+    evaluation when using Lightning Fabric.
+
+    Args:
+        tasks: A single task name or list of task names to evaluate on.
+            Examples: "hellaswag", ["arc_easy", "arc_challenge", "hellaswag"]
+        apply_chat_template: Whether to apply chat template formatting to inputs.
+            Useful for instruction-tuned or chat models.
+        include_path: Path to additional task definitions or custom tasks.
+        batch_size: Number of samples to process in each batch. Larger values
+            may improve throughput but require more memory.
+        metadata: Additional metadata to include in evaluation results.
+        verbosity: Logging verbosity level for the evaluation process.
+        output_path: Custom path for saving evaluation results. If None,
+            results are saved to the default log directory.
+        log_samples: Whether to log individual sample predictions and targets.
+            Useful for debugging but increases output size significantly.
+        _usage_: Internal usage tracking string.
+        _version_: Internal version tracking string.
+        **kwargs: Additional arguments passed to the LM Evaluation Harness.
+
+    Example:
+        ```python
+        >>> taskpool = LMEvalHarnessTaskPool(
+        ...     tasks=["arc_easy", "hellaswag"],
+        ...     batch_size=8,
+        ...     verbosity="INFO"
+        ... )
+        >>> results = taskpool.evaluate(model)
+        ```
+    """
+
     def __init__(
         self,
         tasks: Union[str, List[str]],
@@ -44,6 +85,45 @@ class LMEvalHarnessTaskPool(BaseTaskPool, LightningFabricMixin):
         self.log_samples = log_samples
 
     def evaluate(self, model, *command_line_args, **kwargs):
+        """Evaluate a language model on the configured tasks using LM Evaluation Harness.
+
+        This method wraps the model with the LM Evaluation Harness framework and
+        executes evaluation on all configured tasks. It automatically handles
+        command-line argument construction, model wrapping with Lightning Fabric
+        for distributed evaluation, and result logging.
+
+        The evaluation process includes:
+        1. Building command-line arguments from instance configuration
+        2. Setting up the LM Evaluation Harness argument parser
+        3. Wrapping the model with Lightning Fabric if not already wrapped
+        4. Creating an HFLM (Hugging Face Language Model) wrapper
+        5. Executing the evaluation through the LM-Eval CLI interface
+
+        Args:
+            model: The language model to evaluate. Can be a Hugging Face model,
+                PyTorch model, or any model compatible with the LM Evaluation Harness.
+                The model will be automatically wrapped with Lightning Fabric for
+                distributed evaluation if not already wrapped.
+            *command_line_args: Additional positional command-line arguments
+                (currently unused but preserved for interface compatibility).
+            **kwargs: Additional keyword arguments that will be converted to
+                command-line flags and passed to the LM Evaluation Harness.
+                Keys will be prefixed with '--' and values converted to strings.
+
+        Returns:
+            None: Results are written to the configured output path and logged.
+
+        Example:
+            ```python
+            >>> taskpool = LMEvalHarnessTaskPool(tasks=["arc_easy"])
+            >>> taskpool.evaluate(model, limit=100, device="cuda")
+            ```
+
+        Note:
+            The method leverages the LM Evaluation Harness's command-line interface
+            internally, which provides standardized evaluation procedures and
+            ensures compatibility with the broader evaluation ecosystem.
+        """
         command_line_args = []
         if self.include_path is not None:
             command_line_args.extend(["--include_path", self.include_path])

fusion_bench/taskpool/nyuv2_taskpool.py

@@ -15,9 +15,37 @@ log = logging.getLogger(__name__)
 
 
 class NYUv2TaskPool(TaskPool):
+    """Task pool for multi-task learning evaluation on the NYUv2 dataset.
+
+    This task pool provides evaluation capabilities for multi-task learning models
+    on the NYU Depth V2 (NYUv2) dataset, which is a popular benchmark for indoor
+    scene understanding. The dataset supports multiple computer vision tasks
+    including semantic segmentation, depth estimation, and surface normal prediction.
+
+    The task pool is designed to work with encoder-decoder architectures where
+    a shared encoder processes input images and task-specific decoders generate
+    predictions for different tasks. It integrates with PyTorch Lightning for
+    streamlined training and evaluation workflows.
+
+    Supported Tasks:
+        - Semantic segmentation
+        - Depth estimation
+        - Surface normal prediction
+    """
+
     _trainer: L.Trainer = None
 
     def __init__(self, taskpool_config: DictConfig):
+        """Initialize the NYUv2 task pool with configuration settings.
+
+        Args:
+            taskpool_config: Configuration object containing all necessary
+                parameters for the task pool, including:
+                - data_dir: Path to the directory containing NYUv2 dataset
+                - tasks: List of tasks to evaluate (e.g., ["semantic", "depth"])
+                - batch_size: Batch size for evaluation data loader
+                - num_workers: Number of worker processes for data loading
+        """
         self.config = taskpool_config
 
     def load_datasets(self):

fusion_bench/utils/__init__.py

@@ -18,4 +18,6 @@ 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

fusion_bench/utils/cache_utils.py

@@ -1,15 +1,30 @@
 import logging
 import os
 import pickle
+import warnings
 from functools import wraps
 from pathlib import Path
 from typing import Any, Callable, Union
 
-__all__ = ["cache_to_disk"]
+from joblib import Memory
+
+__all__ = ["cache_to_disk", "cache_with_joblib", "set_default_cache_dir"]
 
 
 log = logging.getLogger(__name__)
 
+DEFAULT_CACHE_DIR = Path.cwd() / "outputs" / "cache"
+
+
+def set_default_cache_dir(path: str | Path):
+    global DEFAULT_CACHE_DIR
+    if path is None:
+        return
+
+    if isinstance(path, str):
+        path = Path(path)
+    DEFAULT_CACHE_DIR = path
+
 
 def cache_to_disk(file_path: Union[str, Path]) -> Callable:
     """
@@ -17,6 +32,11 @@ def cache_to_disk(file_path: Union[str, Path]) -> Callable:
     the result is loaded from the file. Otherwise, the function is executed and
     the result is saved to the file.
 
+    !!! warning "deprecated"
+        This function is deprecated. Use `cache_with_joblib` instead for better
+        caching capabilities including automatic cache invalidation, better object
+        handling, and memory efficiency.
+
     ## Example usage
 
     ```python
@@ -32,6 +52,13 @@ def cache_to_disk(file_path: Union[str, Path]) -> Callable:
     Returns:
         Callable: The decorated function.
     """
+    warnings.warn(
+        "cache_to_disk is deprecated. Use cache_with_joblib instead for better "
+        "caching capabilities including automatic cache invalidation, better object "
+        "handling, and memory efficiency.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
     if isinstance(file_path, str):
         file_path = Path(file_path)
     assert isinstance(file_path, Path)
@@ -56,3 +83,76 @@ def cache_to_disk(file_path: Union[str, Path]) -> Callable:
         return wrapper
 
     return decorator
+
+
+def cache_with_joblib(
+    cache_dir: Union[str, Path] = None,
+    verbose: int = 0,
+) -> Callable:
+    """
+    A decorator to cache the result of a function using joblib.Memory. This provides
+    more advanced caching capabilities compared to cache_to_disk, including:
+    - Automatic cache invalidation when function arguments change
+    - Better handling of numpy arrays and other complex objects
+    - Memory-efficient storage
+    - Optional verbose output for cache hits/misses
+
+    ## Example usage
+
+    ```python
+    @cache_with_joblib("./cache", verbose=1)
+    def expensive_computation(x: int, y: str) -> Any:
+        # Function implementation
+        return complex_result
+
+    # Or with default settings:
+    @cache_with_joblib()
+    def another_function(x: int) -> int:
+        return x * 2
+    ```
+
+    Args:
+        cache_dir (Union[str, Path]): The directory where cache files should be stored.
+            If `None`, a default directory `outputs/cache` will be used.
+        verbose (int): Verbosity level for joblib.Memory (0=silent, 1=basic, 2++=verbose).
+
+    Returns:
+        Callable: A decorator function that can be applied to functions.
+    """
+
+    if cache_dir is None:
+        cache_dir = DEFAULT_CACHE_DIR
+
+    if isinstance(cache_dir, str):
+        cache_dir = Path(cache_dir)
+    assert isinstance(cache_dir, Path)
+
+    # Create the cache directory if it doesn't exist
+    cache_dir.mkdir(parents=True, exist_ok=True)
+
+    # Create a Memory object for this function
+    memory = Memory(location=cache_dir, verbose=verbose)
+
+    def decorator(func: Callable) -> Callable:
+        nonlocal memory
+
+        # Create the cached version of the function
+        cached_func = memory.cache(func)
+
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            return cached_func(*args, **kwargs)
+
+        # Expose useful methods from joblib.Memory
+        if not (
+            hasattr(cached_func, "clear")
+            or hasattr(cached_func, "call")
+            or hasattr(cached_func, "check_call_in_cache")
+        ):
+            wrapper.clear = cached_func.clear
+            wrapper.call = cached_func.call
+            wrapper.check_call_in_cache = cached_func.check_call_in_cache
+
+        return wrapper
+
+    return decorator

fusion_bench/utils/data.py

@@ -1,6 +1,6 @@
 import pickle
 from pathlib import Path
-from typing import Literal, Optional, Union
+from typing import Any, Literal, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -37,7 +37,9 @@ class InfiniteDataLoader:
         return data
 
 
-def load_tensor_from_file(file_path: Union[str, Path], device=None) -> torch.Tensor:
+def load_tensor_from_file(
+    file_path: Union[str, Path], device: Optional[Union[str, torch.device]] = None
+) -> torch.Tensor:
     """
     Loads a tensor from a file, which can be either a .pt, .pth or .np file.
     If the file is not one of these formats, it will try to load it as a pickle file.
@@ -72,7 +74,7 @@ def train_validation_split(
     validation_size: Optional[int] = None,
     random_seed: Optional[int] = None,
     return_split: Literal["all", "train", "val"] = "both",
-):
+) -> Union[Tuple[Dataset, Dataset], Dataset]:
     """
     Split a dataset into a training and validation set.
 
@@ -134,7 +136,7 @@ def train_validation_test_split(
     test_fraction: float,
     random_seed: Optional[int] = None,
     return_spilt: Literal["all", "train", "val", "test"] = "all",
-):
+) -> Union[Tuple[Dataset, Dataset, Dataset], Dataset]:
     """
     Split a dataset into a training, validation and test set.