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

fusion_bench/metrics/nyuv2/normal.py

@@ -7,14 +7,36 @@ from torchmetrics import Metric
 
 
 class NormalMetric(Metric):
+    """
+    Metric for evaluating surface normal prediction on NYUv2 dataset.
+
+    This metric computes angular error statistics between predicted and ground truth
+    surface normals, including mean, median, and percentage of predictions within
+    specific angular thresholds (11.25°, 22.5°, 30°).
+
+    Attributes:
+        metric_names: List of metric names ["mean", "median", "<11.25", "<22.5", "<30"].
+        record: List storing angular errors (in degrees) for all pixels across batches.
+    """
+
     metric_names = ["mean", "median", "<11.25", "<22.5", "<30"]
 
     def __init__(self):
+        """Initialize the NormalMetric with state for recording angular errors."""
         super(NormalMetric, self).__init__()
 
         self.add_state("record", default=[], dist_reduce_fx="cat")
 
     def update(self, preds, target):
+        """
+        Update metric state with predictions and targets from a batch.
+
+        Args:
+            preds: Predicted surface normals of shape (batch_size, 3, height, width).
+                   Will be L2-normalized before computing errors.
+            target: Ground truth surface normals of shape (batch_size, 3, height, width).
+                   Already normalized on NYUv2 dataset. Pixels with sum of 0 are invalid.
+        """
         # gt has been normalized on the NYUv2 dataset
         preds = preds / torch.norm(preds, p=2, dim=1, keepdim=True)
         binary_mask = torch.sum(target, dim=1) != 0
@@ -33,7 +55,18 @@ class NormalMetric(Metric):
 
     def compute(self):
         """
-        returns mean, median, and percentage of pixels with error less than 11.25, 22.5, and 30 degrees ("mean", "median", "<11.25", "<22.5", "<30")
+        Compute final metric values from all recorded angular errors.
+
+        Returns:
+            List[Tensor]: A list containing five metrics:
+                - mean: Mean angular error in degrees.
+                - median: Median angular error in degrees.
+                - <11.25: Percentage of pixels with error < 11.25°.
+                - <22.5: Percentage of pixels with error < 22.5°.
+                - <30: Percentage of pixels with error < 30°.
+
+        Note:
+            Returns zeros if no data has been recorded.
         """
         if self.record is None:
             return torch.asarray([0.0, 0.0, 0.0, 0.0, 0.0])

fusion_bench/metrics/nyuv2/segmentation.py

@@ -6,9 +6,28 @@ from torchmetrics import Metric
 
 
 class SegmentationMetric(Metric):
+    """
+    Metric for evaluating semantic segmentation on NYUv2 dataset.
+
+    This metric computes mean Intersection over Union (mIoU) and pixel accuracy
+    for multi-class segmentation tasks.
+
+    Attributes:
+        metric_names: List of metric names ["mIoU", "pixAcc"].
+        num_classes: Number of segmentation classes (default: 13 for NYUv2).
+        record: Confusion matrix of shape (num_classes, num_classes) tracking
+                predictions vs ground truth.
+    """
+
     metric_names = ["mIoU", "pixAcc"]
 
     def __init__(self, num_classes=13):
+        """
+        Initialize the SegmentationMetric.
+
+        Args:
+            num_classes: Number of segmentation classes. Default is 13 for NYUv2 dataset.
+        """
         super().__init__()
 
         self.num_classes = num_classes
@@ -21,9 +40,19 @@ class SegmentationMetric(Metric):
         )
 
     def reset(self):
+        """Reset the confusion matrix to zeros."""
         self.record.zero_()
 
     def update(self, preds: Tensor, target: Tensor):
+        """
+        Update the confusion matrix with predictions and targets from a batch.
+
+        Args:
+            preds: Predicted segmentation logits of shape (batch_size, num_classes, height, width).
+                   Will be converted to class predictions via softmax and argmax.
+            target: Ground truth segmentation labels of shape (batch_size, height, width).
+                   Pixels with negative values or values >= num_classes are ignored.
+        """
         preds = preds.softmax(1).argmax(1).flatten()
         target = target.long().flatten()
 
@@ -35,7 +64,12 @@ class SegmentationMetric(Metric):
 
     def compute(self):
         """
-        return mIoU and pixel accuracy
+        Compute mIoU and pixel accuracy from the confusion matrix.
+
+        Returns:
+            List[Tensor]: A list containing [mIoU, pixel_accuracy]:
+                - mIoU: Mean Intersection over Union across all classes.
+                - pixel_accuracy: Overall pixel classification accuracy.
         """
         h = cast(Tensor, self.record).float()
         iu = torch.diag(h) / (h.sum(1) + h.sum(0) - torch.diag(h))

fusion_bench/mixins/clip_classification.py

@@ -59,6 +59,15 @@ class CLIPClassificationMixin(LightningFabricMixin):
 
     @property
     def clip_processor(self):
+        """
+        Get the CLIP processor, loading it from the model pool if necessary.
+
+        Returns:
+            CLIPProcessor: The CLIP processor for image and text preprocessing.
+
+        Raises:
+            AssertionError: If the model pool is not set.
+        """
         if self._clip_processor is None:
             assert self.modelpool is not None, "Model pool is not set"
             self._clip_processor = self.modelpool.load_processor()
@@ -125,6 +134,11 @@ class CLIPClassificationMixin(LightningFabricMixin):
             clip_model (Optional[CLIPModel]): The CLIP model to use. If not provided, a pretrained model is loaded from the model pool.
             task_names (Optional[List[str]]): A list of task names to set up the classification head for. If not provided, all models in the model pool will be used.
         """
+        # make sure the task names are equal across all processes
+        _task_names = self.fabric.broadcast(task_names, src=0)
+        if not self.fabric.is_global_zero and task_names != _task_names:
+            raise ValueError("The `task_names` must be the same across all processes.")
+
         self.whether_setup_zero_shot_classification_head = True
         # load clip model if not provided
         if clip_model is None:
@@ -147,7 +161,10 @@ class CLIPClassificationMixin(LightningFabricMixin):
         self.logit_scale_exp = self.fabric.to_device(self.logit_scale_exp)
 
         @cache_with_joblib()
-        def construct_classification_head(task: str):
+        def construct_classification_head(task: str, model_name: str):
+            log.info(
+                f"Constructing zero-shot classification head for task: {task} using model: {model_name}"
+            )
             nonlocal clip_classifier
 
             classnames, templates = get_classnames_and_templates(task)
@@ -163,7 +180,18 @@ class CLIPClassificationMixin(LightningFabricMixin):
         ):
             zeroshot_weights = None
             if self.fabric.is_global_zero:
-                zeroshot_weights = construct_classification_head(task)
+                if hasattr(clip_model, "config") and hasattr(
+                    clip_model.config, "_name_or_path"
+                ):
+                    model_name = clip_model.config._name_or_path
+                else:
+                    model_name = "unknown_model"
+                    log.warning(
+                        "CLIP model config does not have `_name_or_path` attribute. Using 'unknown_model' as model name."
+                    )
+                zeroshot_weights = construct_classification_head(
+                    task, model_name=model_name
+                )
 
             self.fabric.barrier()
             self.zeroshot_weights[task] = self.fabric.broadcast(zeroshot_weights, src=0)

fusion_bench/mixins/lightning_fabric.py

@@ -1,7 +1,7 @@
 import functools
 import logging
 import os
-from typing import TYPE_CHECKING, Any, List, Optional, TypeVar
+from typing import TYPE_CHECKING, Any, List, Mapping, Optional, TypeVar
 
 import lightning as L
 import torch
@@ -96,12 +96,24 @@ class LightningFabricMixin:
 
     @property
     def fabric(self):
+        """
+        Get the Lightning Fabric instance, initializing it if necessary.
+
+        Returns:
+            L.Fabric: The Lightning Fabric instance for distributed computing.
+        """
         if self._fabric_instance is None:
             self.setup_lightning_fabric(getattr(self, "config", DictConfig({})))
         return self._fabric_instance
 
     @fabric.setter
     def fabric(self, instance: L.Fabric):
+        """
+        Set the Lightning Fabric instance.
+
+        Args:
+            instance: The Lightning Fabric instance to use.
+        """
         self._fabric_instance = instance
 
     @property
@@ -172,6 +184,15 @@ class LightningFabricMixin:
     def tensorboard_summarywriter(
         self,
     ) -> "lightning.fabric.loggers.tensorboard.SummaryWriter":
+        """
+        Get the TensorBoard SummaryWriter for detailed logging.
+
+        Returns:
+            SummaryWriter: The TensorBoard SummaryWriter instance.
+
+        Raises:
+            AttributeError: If the logger is not a TensorBoardLogger.
+        """
         if isinstance(self.fabric.logger, TensorBoardLogger):
             return self.fabric.logger.experiment
         else:
@@ -179,6 +200,12 @@ class LightningFabricMixin:
 
     @property
     def is_debug_mode(self):
+        """
+        Check if the program is running in debug mode (fast_dev_run).
+
+        Returns:
+            bool: True if fast_dev_run is enabled, False otherwise.
+        """
         if hasattr(self, "config") and self.config.get("fast_dev_run", False):
             return True
         elif hasattr(self, "_program") and self._program.config.get(
@@ -190,13 +217,22 @@ class LightningFabricMixin:
 
     def log(self, name: str, value: Any, step: Optional[int] = None):
         """
-        Logs the metric to the fabric's logger.
+        Logs a single metric to the fabric's logger.
+
+        Args:
+            name: The name of the metric to log.
+            value: The value of the metric.
+            step: Optional step number for the metric.
         """
         self.fabric.log(name, value, step=step)
 
-    def log_dict(self, metrics: dict, step: Optional[int] = None):
+    def log_dict(self, metrics: Mapping[str, Any], step: Optional[int] = None):
         """
-        Logs the metrics to the fabric's logger.
+        Logs multiple metrics to the fabric's logger.
+
+        Args:
+            metrics: Dictionary of metric names and values.
+            step: Optional step number for the metrics.
         """
         self.fabric.log_dict(metrics, step=step)
 
@@ -207,7 +243,12 @@ class LightningFabricMixin:
         name_template: str = "train/lr_group_{0}",
     ):
         """
-        Logs the learning rate of the optimizer to the fabric's logger.
+        Logs the learning rate of each parameter group in the optimizer.
+
+        Args:
+            optimizer: The optimizer whose learning rates should be logged.
+            step: Optional step number for the log entry.
+            name_template: Template string for the log name. Use {0} as placeholder for group index.
         """
         for i, param_group in enumerate(optimizer.param_groups):
             self.fabric.log(name_template.format(i), param_group["lr"], step=step)

fusion_bench/mixins/rich_live.py

@@ -2,20 +2,96 @@ from rich.live import Live
 
 
 class RichLiveMixin:
+    """
+    A mixin class that provides Rich Live display capabilities.
+
+    This mixin integrates Rich's Live display functionality, allowing for
+    dynamic, auto-refreshing console output. It's particularly useful for
+    displaying real-time updates, progress information, or continuously
+    changing data without cluttering the terminal.
+
+    Attributes:
+        _rich_live (Live): The internal Rich Live instance for live display updates.
+
+    Example:
+        ```python
+        class MyTask(RichLiveMixin):
+            def run(self):
+                self.start_rich_live()
+                for i in range(100):
+                    self.rich_live_print(f"Processing item {i}")
+                    time.sleep(0.1)
+                self.stop_rich_live()
+        ```
+    """
+
     _rich_live: Live = None
 
     @property
     def rich_live(self) -> Live:
+        """
+        Get the Rich Live instance, creating it if necessary.
+
+        Returns:
+            Live: The Rich Live instance for dynamic console output.
+        """
         if self._rich_live is None:
             self._rich_live = Live()
         return self._rich_live
 
     def start_rich_live(self):
+        """
+        Start the Rich Live display context.
+
+        This method enters the Rich Live context, enabling dynamic console output.
+        Must be paired with stop_rich_live() to properly clean up resources.
+
+        Returns:
+            The Rich Live instance in its started state.
+
+        Example:
+            ```python
+            self.start_rich_live()
+            # Display dynamic content
+            self.rich_live_print("Dynamic output")
+            self.stop_rich_live()
+            ```
+        """
         return self.rich_live.__enter__()
 
     def stop_rich_live(self):
+        """
+        Stop the Rich Live display context and clean up resources.
+
+        This method exits the Rich Live context and resets the internal Live instance.
+        Should be called after start_rich_live() when dynamic display is complete.
+
+        Example:
+            ```python
+            self.start_rich_live()
+            # ... display content ...
+            self.stop_rich_live()
+            ```
+        """
         self.rich_live.__exit__(None, None, None)
         self._rich_live = None
 
     def rich_live_print(self, msg):
+        """
+        Print a message to the Rich Live console.
+
+        This method displays the given message through the Rich Live console,
+        allowing for formatted, dynamic output that updates in place.
+
+        Args:
+            msg: The message to display. Can be a string or any Rich renderable object.
+
+        Example:
+            ```python
+            self.start_rich_live()
+            self.rich_live_print("[bold green]Success![/bold green]")
+            self.rich_live_print(Panel("Status: Running"))
+            self.stop_rich_live()
+            ```
+        """
         self.rich_live.console.print(msg)

fusion_bench/modelpool/base_pool.py

@@ -1,6 +1,6 @@
 import logging
 from copy import deepcopy
-from typing import Dict, Generator, List, Optional, Tuple, Union
+from typing import Any, Dict, Generator, List, Optional, Tuple, Union
 
 import torch
 from omegaconf import DictConfig, OmegaConf, UnsupportedValueType
@@ -8,7 +8,12 @@ from torch import nn
 from torch.utils.data import Dataset
 
 from fusion_bench.mixins import BaseYAMLSerializable, HydraConfigMixin
-from fusion_bench.utils import instantiate, timeit_context
+from fusion_bench.utils import (
+    ValidationError,
+    instantiate,
+    timeit_context,
+    validate_model_name,
+)
 
 __all__ = ["BaseModelPool"]
 
@@ -59,6 +64,16 @@ class BaseModelPool(
             except UnsupportedValueType:
                 pass
 
+        if not models:
+            log.warning("Initialized BaseModelPool with empty models dictionary.")
+        else:
+            # Validate model names
+            for model_name in models.keys():
+                try:
+                    validate_model_name(model_name, allow_special=True)
+                except ValidationError as e:
+                    log.warning(f"Invalid model name '{model_name}': {e}")
+
         self._models = models
         self._train_datasets = train_datasets
         self._val_datasets = val_datasets
@@ -147,7 +162,9 @@ class BaseModelPool(
         """
         return model_name.startswith("_") and model_name.endswith("_")
 
-    def get_model_config(self, model_name: str, return_copy: bool = True) -> DictConfig:
+    def get_model_config(
+        self, model_name: str, return_copy: bool = True
+    ) -> Union[DictConfig, str, Any]:
         """
         Get the configuration for the specified model.
 
@@ -155,10 +172,36 @@ class BaseModelPool(
             model_name (str): The name of the model.
 
         Returns:
-            DictConfig: The configuration for the specified model.
+            Union[DictConfig, str, Any]: The configuration for the specified model, which may be a DictConfig, string path, or other type.
+
+        Raises:
+            ValidationError: If model_name is invalid.
+            KeyError: If model_name is not found in the pool.
         """
+        # Validate model name
+        validate_model_name(model_name, allow_special=True)
+
+        # raise friendly error if model not found in the pool
+        if model_name not in self._models:
+            available_models = list(self._models.keys())
+            raise KeyError(
+                f"Model '{model_name}' not found in model pool. "
+                f"Available models: {available_models}"
+            )
+
         model_config = self._models[model_name]
+        if isinstance(model_config, nn.Module):
+            log.warning(
+                f"Model configuration for '{model_name}' is a pre-instantiated model. "
+                "Returning the model instance instead of configuration."
+            )
+
         if return_copy:
+            if isinstance(model_config, nn.Module):
+                # raise performance warning
+                log.warning(
+                    f"Furthermore, returning a copy of the pre-instantiated model '{model_name}' may be inefficient."
+                )
             model_config = deepcopy(model_config)
         return model_config
 
@@ -171,12 +214,28 @@ class BaseModelPool(
 
         Returns:
             str: The path for the specified model.
+
+        Raises:
+            ValidationError: If model_name is invalid.
+            KeyError: If model_name is not found in the pool.
+            ValueError: If model configuration is not a string path.
         """
+        # Validate model name
+        validate_model_name(model_name, allow_special=True)
+
+        if model_name not in self._models:
+            available_models = list(self._models.keys())
+            raise KeyError(
+                f"Model '{model_name}' not found in model pool. "
+                f"Available models: {available_models}"
+            )
+
         if isinstance(self._models[model_name], str):
             return self._models[model_name]
         else:
             raise ValueError(
-                "Model path is not a string. Try to override this method in derived modelpool class."
+                f"Model configuration for '{model_name}' is not a string path. "
+                "Try to override this method in derived modelpool class."
             )
 
     def load_model(
@@ -357,3 +416,25 @@ class BaseModelPool(
         """
         with timeit_context(f"Saving the state dict of model to {path}"):
             torch.save(model.state_dict(), path)
+
+    def __contains__(self, model_name: str) -> bool:
+        """
+        Check if a model with the given name exists in the model pool.
+
+        Examples:
+            >>> modelpool = BaseModelPool(models={"modelA": ..., "modelB": ...})
+            >>> "modelA" in modelpool
+            True
+            >>> "modelC" in modelpool
+            False
+
+        Args:
+            model_name (str): The name of the model to check.
+
+        Returns:
+            bool: True if the model exists, False otherwise.
+        """
+        if self._models is None:
+            raise RuntimeError("Model pool is not initialized")
+        validate_model_name(model_name, allow_special=True)
+        return model_name in self._models