fusion-bench 0.2.28__py3-none-any.whl → 0.2.29__py3-none-any.whl

fusion_bench/constants/__init__.py

@@ -5,4 +5,8 @@ from .paths import *
 from .runtime import RuntimeConstants
 
 # fusionbench version
-FUSION_BENCH_VERSION = importlib.metadata.version("fusion-bench")
+try:
+    FUSION_BENCH_VERSION = importlib.metadata.version("fusion-bench")
+except importlib.metadata.PackageNotFoundError:
+    # Fallback when package is not installed (e.g., during development)
+    FUSION_BENCH_VERSION = "0.0.0.dev"

fusion_bench/constants/runtime.py

@@ -1,3 +1,29 @@
+"""
+Runtime Constants Module.
+
+This module provides a thread-safe singleton class for managing runtime configuration
+and constants across the Fusion Bench framework. It centralizes access to runtime
+settings like cache directories, debug flags, and logging preferences.
+
+Example:
+    ```python
+    from fusion_bench.constants.runtime import RuntimeConstants
+
+    # Get the singleton instance
+    runtime = RuntimeConstants()
+
+    # Configure cache directory
+    runtime.cache_dir = "/custom/cache/path"
+
+    # Enable debug mode
+    runtime.debug = True
+
+    # Control function call logging
+    runtime.print_function_call = True
+    ```
+"""
+
+import os
 import threading
 from pathlib import Path
 from typing import Optional, Union
@@ -5,18 +31,46 @@ from typing import Optional, Union
 
 class RuntimeConstants:
     """
-    This class holds constants related to the runtime environment of the Fusion Bench framework.
-    It includes default values for cache directories and other runtime configurations.
+    Thread-safe singleton for managing runtime configuration and constants.
+
+    This class provides centralized access to runtime settings that affect the
+    behavior of the entire Fusion Bench framework. It ensures consistent
+    configuration across all modules and supports thread-safe access in
+    multi-threaded environments.
+
+    Attributes:
+        debug: Global debug flag for enabling verbose logging and debugging features.
+
+    Example:
+        ```python
+        runtime = RuntimeConstants()
+
+        # Configure caching
+        runtime.cache_dir = Path.home() / ".cache" / "fusion_bench"
+
+        # Enable debugging
+        runtime.debug = True
+        runtime.print_function_call = True
+        ```
 
-    Implemented as a thread-safe singleton to ensure consistent runtime configuration
-    across the entire application.
+    Note:
+        This class implements the singleton pattern with thread-safe initialization.
+        Multiple calls to the constructor will return the same instance.
     """
 
     _instance: Optional["RuntimeConstants"] = None
     _lock = threading.Lock()
 
     def __new__(cls) -> "RuntimeConstants":
-        """Create a new instance using singleton pattern with thread safety."""
+        """
+        Create or return the singleton instance with thread safety.
+
+        Uses double-check locking pattern to ensure thread-safe singleton creation
+        while minimizing synchronization overhead.
+
+        Returns:
+            The singleton RuntimeConstants instance.
+        """
         with cls._lock:
             # Double-check locking pattern
             if cls._instance is None:
@@ -25,33 +79,83 @@ class RuntimeConstants:
             return cls._instance
 
     def __init__(self):
-        """Initialize the singleton instance only once."""
+        """
+        Initialize the singleton instance only once.
+
+        Subsequent calls to __init__ are no-ops to maintain singleton behavior.
+        """
         if not self._initialized:
-            # Add your runtime constants here
+            # Initialize default values
             self._initialized = True
 
     debug = False
+    """Global debug flag for enabling verbose logging and debugging features."""
 
     @property
     def cache_dir(self) -> Path:
+        """
+        Get the default cache directory for models and datasets.
+
+        Returns:
+            Path object pointing to the cache directory.
+
+        Example:
+            ```python
+            runtime = RuntimeConstants()
+            print(f"Cache directory: {runtime.cache_dir}")
+            ```
+        """
         from fusion_bench.utils.cache_utils import DEFAULT_CACHE_DIR
 
         return DEFAULT_CACHE_DIR
 
     @cache_dir.setter
     def cache_dir(self, path: Union[str, Path]) -> None:
+        """
+        Set the default cache directory for models and datasets.
+
+        Args:
+            path: New cache directory path as string or Path object.
+
+        Example:
+            ```python
+            runtime = RuntimeConstants()
+            runtime.cache_dir = "/data/fusion_bench_cache"
+            ```
+        """
         from fusion_bench.utils.cache_utils import set_default_cache_dir
 
         set_default_cache_dir(path)
 
     @property
     def print_function_call(self) -> bool:
+        """
+        Get whether function calls are printed during instantiation.
+
+        Returns:
+            True if function call printing is enabled, False otherwise.
+        """
         from fusion_bench.utils.instantiate_utils import PRINT_FUNCTION_CALL
 
         return PRINT_FUNCTION_CALL
 
     @print_function_call.setter
     def print_function_call(self, enable: bool) -> None:
+        """
+        Set whether to print function calls during instantiation.
+
+        Useful for debugging to see which functions are being called
+        when instantiating objects from configuration.
+
+        Args:
+            enable: True to enable printing, False to disable.
+
+        Example:
+            ```python
+            runtime = RuntimeConstants()
+            runtime.print_function_call = True  # Enable verbose logging
+            ```
+        """
         from fusion_bench.utils.instantiate_utils import set_print_function_call
 
         set_print_function_call(enable)

fusion_bench/dataset/gsm8k.py

@@ -13,8 +13,12 @@ def load_gsm8k_question_label_data(
 
     An example in the dataset:
 
-    {'question': 'Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?',
-     'answer': 'Natalia sold 48/2 = <<48/2=24>>24 clips in May.\nNatalia sold 48+24 = <<48+24=72>>72 clips altogether in April and May.\n#### 72'}
+    ```python
+    {
+        'question': 'Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?',
+        'answer': 'Natalia sold 48/2 = <<48/2=24>>24 clips in May.\nNatalia sold 48+24 = <<48+24=72>>72 clips altogether in April and May.\n#### 72'
+    }
+    ```
 
     Args:
         dataset_name (Literal["train", "test", "train_socratic", "test_socratic"]): The name of the dataset to load.

fusion_bench/dataset/image_corruption/make_corruption.py

@@ -1,4 +1,31 @@
 # -*- coding: utf-8 -*-
+"""
+Image Corruption Module for Robustness Testing.
+
+This module provides various image corruption functions to test model robustness.
+It implements common corruptions such as noise, blur, compression artifacts, and
+weather effects. These corruptions are commonly used in benchmark datasets like
+ImageNet-C and CIFAR-10-C.
+
+The corruptions can be applied at different severity levels (1-5), where higher
+levels indicate stronger corruption effects.
+
+Example:
+    ```python
+    from PIL import Image
+    from fusion_bench.dataset.image_corruption.make_corruption import gaussian_noise, motion_blur
+
+    # Load an image
+    img = Image.open("example.jpg")
+
+    # Apply gaussian noise at severity level 3
+    corrupted_img = gaussian_noise(img, severity=3)
+
+    # Apply motion blur at severity level 2
+    blurred_img = motion_blur(img, severity=2)
+    ```
+"""
+
 import logging
 
 logger = logging.getLogger(__name__)
@@ -37,11 +64,39 @@ warnings.simplefilter("ignore", UserWarning)
 
 # /////////////// Distortions ///////////////
 class MotionImage(WandImage):
+    """
+    Extended WandImage class with motion blur capability.
+
+    This class wraps ImageMagick's motion blur functionality through the Wand library.
+    """
+
     def motion_blur(self, radius=0.0, sigma=0.0, angle=0.0):
+        """
+        Apply motion blur effect to the image.
+
+        Args:
+            radius: The radius of the Gaussian, in pixels, not counting the center pixel.
+            sigma: The standard deviation of the Gaussian, in pixels.
+            angle: Apply the effect along this angle in degrees.
+        """
         wandlibrary.MagickMotionBlurImage(self.wand, radius, sigma, angle)
 
 
 def gaussian_noise(x, severity=1):
+    """
+    Apply Gaussian noise corruption to an image.
+
+    Adds random Gaussian noise to the image, simulating sensor noise or
+    environmental interference.
+
+    Args:
+        x: Input image as PIL Image or numpy array. If numpy array, should be in
+           range [0, 255].
+        severity: Corruption severity level from 1 (mild) to 5 (severe).
+
+    Returns:
+        numpy.ndarray: Corrupted image as numpy array in range [0, 255].
+    """
     c = [0.04, 0.06, 0.08, 0.09, 0.10][severity - 1]
 
     x = np.array(x) / 255.0
@@ -49,6 +104,20 @@ def gaussian_noise(x, severity=1):
 
 
 def impulse_noise(x, severity=1):
+    """
+    Apply impulse (salt-and-pepper) noise corruption to an image.
+
+    Randomly replaces pixels with either maximum or minimum intensity values,
+    simulating transmission errors or faulty pixels.
+
+    Args:
+        x: Input image as PIL Image or numpy array. If numpy array, should be in
+           range [0, 255].
+        severity: Corruption severity level from 1 (mild) to 5 (severe).
+
+    Returns:
+        numpy.ndarray: Corrupted image as numpy array in range [0, 255].
+    """
     c = [0.01, 0.02, 0.03, 0.05, 0.07][severity - 1]
 
     x = sk.util.random_noise(np.array(x) / 255.0, mode="s&p", amount=c)
@@ -56,6 +125,21 @@ def impulse_noise(x, severity=1):
 
 
 def motion_blur(x, severity=1):
+    """
+    Apply motion blur corruption to an image.
+
+    Simulates camera shake or object motion during image capture by applying
+    directional blur at a random angle.
+
+    Args:
+        x: Input PIL Image.
+        severity: Corruption severity level from 1 (mild) to 5 (severe).
+           Higher severity increases blur radius and sigma.
+
+    Returns:
+        numpy.ndarray: Corrupted image as numpy array in range [0, 255].
+           Returns RGB image regardless of input format.
+    """
     c = [(6, 1), (6, 1.5), (6, 2), (8, 2), (9, 2.5)][severity - 1]
 
     output = BytesIO()
@@ -73,6 +157,21 @@ def motion_blur(x, severity=1):
 
 
 def spatter(x, severity=1):
+    """
+    Apply spatter corruption to an image.
+
+    Simulates liquid splatter effects (water or mud) on the image, creating
+    realistic occlusions similar to raindrops or dirt on a camera lens.
+
+    Args:
+        x: Input image as PIL Image or numpy array. If numpy array, should be in
+           range [0, 255].
+        severity: Corruption severity level from 1 (mild) to 5 (severe).
+           Levels 1-3 simulate water splatter, levels 4-5 simulate mud splatter.
+
+    Returns:
+        numpy.ndarray: Corrupted image as numpy array in range [0, 255].
+    """
     c = [
         (0.62, 0.1, 0.7, 0.7, 0.5, 0),
         (0.65, 0.1, 0.8, 0.7, 0.5, 0),
@@ -140,6 +239,21 @@ def spatter(x, severity=1):
 
 
 def contrast(x, severity=1):
+    """
+    Apply contrast reduction corruption to an image.
+
+    Reduces image contrast by blending pixels toward their mean values,
+    simulating poor lighting conditions or low-quality image sensors.
+
+    Args:
+        x: Input image as PIL Image or numpy array. If numpy array, should be in
+           range [0, 255].
+        severity: Corruption severity level from 1 (mild) to 5 (severe).
+           Higher severity results in lower contrast.
+
+    Returns:
+        numpy.ndarray: Corrupted image as numpy array in range [0, 255].
+    """
     c = [0.75, 0.5, 0.4, 0.3, 0.15][severity - 1]
 
     x = np.array(x) / 255.0
@@ -148,6 +262,20 @@ def contrast(x, severity=1):
 
 
 def jpeg_compression(x, severity=1):
+    """
+    Apply JPEG compression artifacts to an image.
+
+    Simulates compression artifacts from lossy JPEG encoding at various
+    quality levels, commonly seen in heavily compressed images.
+
+    Args:
+        x: Input PIL Image.
+        severity: Corruption severity level from 1 (mild) to 5 (severe).
+           Lower severity uses higher JPEG quality (less compression).
+
+    Returns:
+        PIL.Image: Corrupted image as PIL Image.
+    """
     c = [80, 65, 58, 50, 40][severity - 1]
 
     output = BytesIO()
@@ -158,6 +286,23 @@ def jpeg_compression(x, severity=1):
 
 
 def pixelate(x, severity=1):
+    """
+    Apply pixelation corruption to an image.
+
+    Reduces image resolution by downsampling and then upsampling,
+    creating a blocky, pixelated appearance.
+
+    Args:
+        x: Input PIL Image with size (32, 32).
+        severity: Corruption severity level from 1 (mild) to 5 (severe).
+           Higher severity results in more pixelation.
+
+    Returns:
+        PIL.Image: Corrupted image as PIL Image with original size (32, 32).
+
+    Note:
+        This function is designed for 32x32 images (e.g., CIFAR-10).
+    """
     c = [0.95, 0.9, 0.85, 0.75, 0.65][severity - 1]
 
     x = x.resize((int(32 * c), int(32 * c)), PILImage.BOX)
@@ -170,6 +315,29 @@ def pixelate(x, severity=1):
 
 
 distortion_methods = collections.OrderedDict()
+"""
+OrderedDict mapping corruption names to their corresponding functions.
+
+Available corruptions:
+    - "Gaussian Noise": Additive Gaussian noise
+    - "Impulse Noise": Salt-and-pepper noise
+    - "Motion Blur": Directional motion blur
+    - "Contrast": Reduced contrast
+    - "Pixelate": Resolution reduction
+    - "JPEG": JPEG compression artifacts
+    - "Spatter": Water or mud splatter effects
+
+Example:
+    ```python
+    from PIL import Image
+    from fusion_bench.dataset.image_corruption.make_corruption import distortion_methods
+
+    img = Image.open("example.jpg")
+    for name, corruption_fn in distortion_methods.items():
+        corrupted = corruption_fn(img, severity=3)
+        # Process corrupted image
+    ```
+"""
 distortion_methods["Gaussian Noise"] = gaussian_noise
 distortion_methods["Impulse Noise"] = impulse_noise
 distortion_methods["Motion Blur"] = motion_blur

fusion_bench/method/__init__.py

@@ -167,6 +167,7 @@ if TYPE_CHECKING:
     from .dawe import DataAdaptiveWeightEnsemblingForCLIP
     from .depth_upscaling import DepthUpscalingAlgorithm, DepthUpscalingForLlama
     from .doge_ta import DOGE_TA_Algorithm
+    from .dop import ContinualDOPForCLIP
     from .dummy import DummyAlgorithm
     from .ensemble import (
         MaxModelPredictorAlgorithm,
@@ -215,7 +216,6 @@ if TYPE_CHECKING:
     from .model_recombination import ModelRecombinationAlgorithm
     from .model_stock import ModelStock
     from .opcm import OPCMForCLIP
-    from .dop import ContinualDOPForCLIP
     from .pruning import (
         MagnitudeDiffPruningAlgorithm,
         MagnitudePruningForLlama,

fusion_bench/method/classification/image_classification_finetune.py

@@ -15,7 +15,7 @@ from lightning_utilities.core.rank_zero import rank_zero_only
 from lit_learn.lit_modules import ERM_LitModule
 from omegaconf import DictConfig
 from torch import nn
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, random_split
 from torchmetrics.classification import Accuracy
 
 from fusion_bench import (
@@ -29,7 +29,6 @@ from fusion_bench import (
 from fusion_bench.dataset import CLIPDataset
 from fusion_bench.modelpool import ResNetForImageClassificationPool
 from fusion_bench.tasks.clip_classification import get_num_classes
-from torch.utils.data import random_split
 
 log = get_rankzero_logger(__name__)
 

fusion_bench/method/gossip/clip_task_wise_gossip.py

@@ -13,41 +13,13 @@ from fusion_bench.modelpool import CLIPVisionModelPool
 from fusion_bench.models.hf_clip import HFCLIPClassifier
 from fusion_bench.tasks.clip_classification import get_classnames_and_templates
 from fusion_bench.utils import timeit_context
+from fusion_bench.utils.data import InfiniteDataLoader
 
 from .task_wise_gossip import TaskWiseGossipAlgorithm
 
 log = logging.getLogger(__name__)
 
 
-class InfiniteDataLoader:
-    """
-    A wrapper class for DataLoader to create an infinite data loader.
-    This is useful in case we are only interested in the number of steps and not the number of epochs.
-
-    This class wraps a DataLoader and provides an iterator that resets
-    when the end of the dataset is reached, creating an infinite loop.
-
-    Attributes:
-        data_loader (DataLoader): The DataLoader to wrap.
-        data_iter (iterator): An iterator over the DataLoader.
-    """
-
-    def __init__(self, data_loader):
-        self.data_loader = data_loader
-        self.data_iter = iter(data_loader)
-
-    def __iter__(self):
-        return self
-
-    def __next__(self):
-        try:
-            data = next(self.data_iter)
-        except StopIteration:
-            self.data_iter = iter(self.data_loader)  # Reset the data loader
-            data = next(self.data_iter)
-        return data
-
-
 class CLIPTaskWiseGossipAlgorithm(TaskWiseGossipAlgorithm):
     """
     A class for task-wise adaptive merging of CLIP models.

fusion_bench/metrics/nyuv2/__init__.py

@@ -1,3 +1,34 @@
+"""
+NYUv2 Dataset Metrics Module.
+
+This module provides metric classes and loss functions for evaluating multi-task learning
+models on the NYUv2 dataset. NYUv2 is a popular indoor scene understanding dataset that
+includes multiple tasks: semantic segmentation, depth estimation, and surface normal prediction.
+
+Available Metrics:
+    - SegmentationMetric: Computes mIoU and pixel accuracy for semantic segmentation.
+    - DepthMetric: Computes absolute and relative errors for depth estimation.
+    - NormalMetric: Computes angular errors for surface normal prediction.
+    - NoiseMetric: Placeholder metric for noise evaluation.
+
+Usage:
+    ```python
+    from fusion_bench.metrics.nyuv2 import SegmentationMetric, DepthMetric
+
+    # Initialize metrics
+    seg_metric = SegmentationMetric(num_classes=13)
+    depth_metric = DepthMetric()
+
+    # Update with predictions and targets
+    seg_metric.update(seg_preds, seg_targets)
+    depth_metric.update(depth_preds, depth_targets)
+
+    # Compute final metrics
+    miou, pix_acc = seg_metric.compute()
+    abs_err, rel_err = depth_metric.compute()
+    ```
+"""
+
 from .depth import DepthMetric
 from .noise import NoiseMetric
 from .normal import NormalMetric

fusion_bench/metrics/nyuv2/depth.py

@@ -7,9 +7,23 @@ from torchmetrics import Metric
 
 
 class DepthMetric(Metric):
+    """
+    Metric for evaluating depth estimation performance on NYUv2 dataset.
+
+    This metric computes absolute error and relative error for depth predictions,
+    properly handling the binary mask to exclude invalid depth regions.
+
+    Attributes:
+        metric_names: List of metric names ["abs_err", "rel_err"].
+        abs_record: List storing absolute error values for each batch.
+        rel_record: List storing relative error values for each batch.
+        batch_size: List storing batch sizes for weighted averaging.
+    """
+
     metric_names = ["abs_err", "rel_err"]
 
     def __init__(self):
+        """Initialize the DepthMetric with state variables for tracking errors."""
         super().__init__()
 
         self.add_state("abs_record", default=[], dist_reduce_fx="cat")
@@ -17,11 +31,20 @@ class DepthMetric(Metric):
         self.add_state("batch_size", default=[], dist_reduce_fx="cat")
 
     def reset(self):
+        """Reset all metric states to empty lists."""
         self.abs_record = []
         self.rel_record = []
         self.batch_size = []
 
     def update(self, preds: Tensor, target: Tensor):
+        """
+        Update metric states with predictions and targets from a batch.
+
+        Args:
+            preds: Predicted depth values of shape (batch_size, 1, height, width).
+            target: Ground truth depth values of shape (batch_size, 1, height, width).
+                   Pixels with sum of 0 are considered invalid and masked out.
+        """
         binary_mask = (torch.sum(target, dim=1) != 0).unsqueeze(1)
         preds = preds.masked_select(binary_mask)
         target = target.masked_select(binary_mask)
@@ -38,6 +61,13 @@ class DepthMetric(Metric):
         self.batch_size.append(torch.asarray(preds.size(0), device=preds.device))
 
     def compute(self):
+        """
+        Compute the final metric values across all batches.
+
+        Returns:
+            List[Tensor]: A list containing [absolute_error, relative_error],
+                         where each value is the weighted average across all batches.
+        """
         records = torch.stack(
             [torch.stack(self.abs_record), torch.stack(self.rel_record)]
         )

fusion_bench/metrics/nyuv2/loss.py

@@ -3,10 +3,35 @@ from torch import Tensor, nn
 
 
 def segmentation_loss(pred: Tensor, gt: Tensor):
+    """
+    Compute cross-entropy loss for semantic segmentation.
+
+    Args:
+        pred: Predicted segmentation logits of shape (batch_size, num_classes, height, width).
+        gt: Ground truth segmentation labels of shape (batch_size, height, width).
+            Pixels with value -1 are ignored in the loss computation.
+
+    Returns:
+        Tensor: Scalar loss value.
+    """
     return nn.functional.cross_entropy(pred, gt.long(), ignore_index=-1)
 
 
 def depth_loss(pred: Tensor, gt: Tensor):
+    """
+    Compute L1 loss for depth estimation with binary masking.
+
+    This loss function calculates the absolute error between predicted and ground truth
+    depth values, but only for valid pixels (where ground truth depth is non-zero).
+
+    Args:
+        pred: Predicted depth values of shape (batch_size, 1, height, width).
+        gt: Ground truth depth values of shape (batch_size, 1, height, width).
+            Pixels with sum of 0 across channels are considered invalid and masked out.
+
+    Returns:
+        Tensor: Scalar loss value averaged over valid pixels.
+    """
     binary_mask = (torch.sum(gt, dim=1) != 0).float().unsqueeze(1).to(pred.device)
     loss = torch.sum(torch.abs(pred - gt) * binary_mask) / torch.nonzero(
         binary_mask, as_tuple=False
@@ -15,6 +40,21 @@ def depth_loss(pred: Tensor, gt: Tensor):
 
 
 def normal_loss(pred: Tensor, gt: Tensor):
+    """
+    Compute cosine similarity loss for surface normal prediction.
+
+    This loss measures the angular difference between predicted and ground truth
+    surface normals using normalized cosine similarity (1 - dot product).
+
+    Args:
+        pred: Predicted surface normals of shape (batch_size, 3, height, width).
+              Will be L2-normalized before computing loss.
+        gt: Ground truth surface normals of shape (batch_size, 3, height, width).
+            Already normalized on NYUv2 dataset. Pixels with sum of 0 are invalid.
+
+    Returns:
+        Tensor: Scalar loss value (1 - mean cosine similarity) over valid pixels.
+    """
     # gt has been normalized on the NYUv2 dataset
     pred = pred / torch.norm(pred, p=2, dim=1, keepdim=True)
     binary_mask = (torch.sum(gt, dim=1) != 0).float().unsqueeze(1).to(pred.device)

fusion_bench/metrics/nyuv2/noise.py

@@ -6,11 +6,35 @@ from torchmetrics import Metric
 
 
 class NoiseMetric(Metric):
+    """
+    A placeholder metric for noise evaluation on NYUv2 dataset.
+
+    This metric currently serves as a placeholder and always returns a value of 1.
+    It can be extended in the future to include actual noise-related metrics.
+
+    Note:
+        This is a dummy implementation that doesn't perform actual noise measurements.
+    """
+
     def __init__(self):
+        """Initialize the NoiseMetric."""
         super().__init__()
 
     def update(self, preds: Tensor, target: Tensor):
+        """
+        Update metric state (currently a no-op).
+
+        Args:
+            preds: Predicted values (unused).
+            target: Ground truth values (unused).
+        """
         pass
 
     def compute(self):
+        """
+        Compute the metric value.
+
+        Returns:
+            List[int]: A list containing [1] as a placeholder value.
+        """
         return [1]