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

fusion_bench/metrics/model_kinship/calculate_split.py

@@ -0,0 +1,171 @@
+import logging
+from typing import Dict, List
+
+import numpy
+import torch
+from tqdm import tqdm
+
+from .utility import Metric, load_model_state_dict, quantize_8bit
+
+
+def cosine_similarity(a, b):
+    similarity = numpy.sqrt(numpy.dot(a, b) ** 2 / (numpy.dot(a, a) * numpy.dot(b, b)))
+    return similarity
+
+
+def calculate_model_kinship_split(
+    model_1_name: str,
+    model_2_name: str,
+    model_base_name: str,
+    low_precision: bool,
+    metrics: List[str],
+    device: str = "cuda" if torch.cuda.is_available() else "cpu",
+) -> dict:
+
+    # Extract state dictionaries from models
+    state_dict_1 = load_model_state_dict(model_1_name, device)
+    state_dict_2 = load_model_state_dict(model_2_name, device)
+    state_dict_base = load_model_state_dict(model_base_name, device)
+    results = {}
+
+    # Validate metrics before processing
+    valid_metrics = Metric.list()
+    for metric in metrics:
+        try:
+            if metric not in valid_metrics:
+                raise ValueError(
+                    f"Unsupported metric: {metric}. Valid metrics are: {', '.join(valid_metrics)}"
+                )
+            results[metric] = calculate_metrics_by_split(
+                state_dict_1, state_dict_2, state_dict_base, low_precision, metric
+            )
+        except Exception as e:
+            logging.error(f"Error calculating {metric}: {str(e)}")
+            results[metric] = f"Error calculating {metric}: {str(e)}"
+
+    return results
+
+
+def calculate_metrics_by_split(
+    state_dict_1: dict,
+    state_dict_2: dict,
+    state_dict_base: dict,
+    low_precision: bool,
+    metric: str,
+) -> str:
+    """
+    Calculate metrics for each key and integrate results.
+
+    Args:
+        state_dict_1 (dict): State dictionary of first model
+        state_dict_2 (dict): State dictionary of second model
+        state_dict_base (dict): State dictionary of base model
+        low_precision (bool): Whether to use 8-bit quantization
+        metric (str): Metric to calculate ('pcc', 'ed', 'cs')
+
+    Returns:
+        str: Integrated metric result as formatted string
+    """
+    total_similarity = 0.0
+    total_weight = 0.0
+    split_results = {}
+
+    # Determine the number of layers
+    num_layers = state_dict_base["lm_head.weight"].shape[0]
+
+    # Check architectures
+    if (
+        state_dict_1["lm_head.weight"].shape[0]
+        != state_dict_2["lm_head.weight"].shape[0]
+    ):
+        shape_1 = state_dict_1["lm_head.weight"].shape
+        shape_2 = state_dict_2["lm_head.weight"].shape
+        logging.warning(
+            f"Warning: Model architectures do not match. "
+            f"Using sub weight space instead.\n"
+            f"Vocab sizes in model 1: {shape_1[0]}, "
+            f"Vocab sizes in model 2: {shape_2[0]}"
+        )
+
+    # Process each key
+    for key, base_params in tqdm(
+        state_dict_base.items(), desc=f"Processing {metric.upper()} by key"
+    ):
+        try:
+            if key not in state_dict_1 or key not in state_dict_2:
+                logging.warning(f"Key {key} not found in one of the models")
+                continue
+
+            # Get parameters and calculate deltas
+            params_1 = state_dict_1[key][:num_layers]
+            params_2 = state_dict_2[key][:num_layers]
+
+            delta_1 = (params_1 - base_params).view(-1)
+            delta_2 = (params_2 - base_params).view(-1)
+
+            if low_precision:
+                delta_1 = quantize_8bit(delta_1)
+                delta_2 = quantize_8bit(delta_2)
+
+            # Calculate weight based on parameter count
+            weight = delta_1.numel()
+
+            # Calculate metric for current key
+            if metric == "pcc":
+                stack = torch.stack((delta_1, delta_2), dim=0)
+                split_similarity = torch.corrcoef(stack)[0, 1].item()
+            elif metric == "ed":
+                split_similarity = torch.dist(delta_1, delta_2).item()
+            elif metric == "cs":
+                split_similarity = cosine_similarity(delta_1, delta_2)
+            else:
+                raise ValueError(f"Unsupported metric: {metric}")
+
+            # Skip NaN values
+            if torch.isnan(torch.tensor(split_similarity)):
+                logging.warning(f"Skipping key {key} due to NaN result")
+                continue
+
+            # Store valid result
+            split_results[key] = split_similarity
+
+            # Update weighted average only for valid results
+            weight = delta_1.numel()
+            total_similarity += split_similarity * weight
+            total_weight += weight
+
+            # Log progress for large layers
+            if weight > 1000000:
+                logging.info(
+                    f"Layer {key}: {metric.upper()} = {split_similarity:.4f}, parameters = {weight}"
+                )
+
+            # Free memory
+            del delta_1, delta_2
+
+        except Exception as e:
+            logging.error(f"Error processing key {key}: {str(e)}")
+            continue
+
+    # Calculate final weighted average
+    if total_weight > 0:
+        final_result = total_similarity / total_weight
+
+        # Log summary statistics
+        logging.info(f"\nSummary for {metric.upper()}:")
+        logging.info(f"Total parameters: {total_weight}")
+
+        # Log detailed results for valid splits
+        logging.info(f"\nDetailed {metric.upper()} results by key:")
+        for key, value in split_results.items():
+            logging.info(f"{key}: {value:.4f}")
+
+        metric_names = {
+            "pcc": "Pearson Correlation Coefficient",
+            "ed": "Euclidean Distance",
+            "cs": "Cosine Similarity",
+        }
+
+        return f"Model Kinship based on {metric_names[metric]} (weighted average): {final_result:.4f}"
+    else:
+        return f"Error: No valid parameters found for {metric.upper()} calculation"

fusion_bench/metrics/model_kinship/utility.py

@@ -0,0 +1,184 @@
+import logging
+from enum import Enum
+from typing import List
+
+import click
+import torch
+from tqdm import tqdm
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    PretrainedConfig,
+)
+
+
+class Metric(str, Enum):
+    """Enumeration of supported metrics"""
+
+    PCC = "pcc"
+    ED = "ed"
+    CS = "cs"
+
+    @classmethod
+    def list(cls) -> List[str]:
+        """Return list of supported metric values"""
+        return [metric.value for metric in cls]
+
+
+def get_config(model: str, trust_remote_code: bool = False) -> PretrainedConfig:
+    """
+    Fetch the configuration of a pretrained model from HuggingFace.
+
+    Args:
+        model (str): The name or path of the model to load configuration for.
+        trust_remote_code (bool, optional): Whether to trust remote code during loading.
+                                            Defaults to False.
+
+    Returns:
+        PretrainedConfig: The configuration object of the specified model.
+    """
+    # Fetch the configuration from HuggingFace's model hub.
+    config = AutoConfig.from_pretrained(
+        model,
+        trust_remote_code=trust_remote_code,  # Whether to allow remote code execution.
+    )
+    return config
+
+
+def validate_models(model_1: str, model_2: str, base_model: str) -> None:
+    """
+    Validate model names to ensure they are different and exist.
+
+    Args:
+        model_1: Name of the first model
+        model_2: Name of the second model
+        base_model: Name of the base model
+
+    Raises:
+        click.BadParameter: If validation fails
+    """
+    if model_1 == model_2 or model_1 == base_model or model_2 == base_model:
+        raise click.BadParameter("All model names must be different")
+
+
+def quantize_8bit(x: torch.Tensor) -> torch.Tensor:
+    # Get absolute min and max values
+    abs_max = torch.max(torch.abs(x))
+
+    # Scale to [-127, 127] range for 8-bit signed integers
+    # Using 127 instead of 128 to keep zero exactly representable
+    scaled = 127 * (x / abs_max)
+
+    # Round to nearest integer
+    quantized = torch.round(scaled)
+
+    # Clamp values to ensure they stay in valid range
+    quantized = torch.clamp(quantized, -127, 127)
+
+    return quantized
+
+
+def load_model_state_dict(model_name: str, device: str) -> dict:
+    """
+    Load a model and return its state dictionary.
+
+    Args:
+        model_name (str): Name or path of the model to load
+        device (str): Device to load the model on ('cuda' or 'cpu')
+
+    Returns:
+        dict: State dictionary of the loaded model
+    """
+    logging.info(f"Loading model: {model_name}")
+    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
+    state_dict = model.state_dict()
+    del model  # Free memory
+    return state_dict
+
+
+def extract_delta_parameters(
+    model_1_name: str,
+    model_2_name: str,
+    model_base_name: str,
+    low_precision: bool,
+    device: str = "cuda" if torch.cuda.is_available() else "cpu",
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Extract the delta parameters (weight differences) between two models
+    relative to a base model.
+
+    Args:
+        model_1_name (str): Name or path of the first model.
+        model_2_name (str): Name or path of the second model.
+        model_base_name (str): Name or path of the base model for comparison.
+        low_precision (bool): Whether to use low precision weights
+
+    Returns:
+        (torch.Tensor, torch.Tensor): Delta parameters of model_1 and model_2 relative to base model.
+    """
+
+    # Extract state dictionaries from models
+    state_dict_1 = load_model_state_dict(model_1_name, device)
+    state_dict_2 = load_model_state_dict(model_2_name, device)
+    state_dict_base = load_model_state_dict(model_base_name, device)
+
+    # Determine the number of layers
+    num_layers = state_dict_base["lm_head.weight"].shape[0]
+
+    # Check if model architectures match, log a warning if not
+    if (
+        state_dict_1["lm_head.weight"].shape[0]
+        != state_dict_2["lm_head.weight"].shape[0]
+    ):
+        shape_1 = state_dict_1["lm_head.weight"].shape
+        shape_2 = state_dict_2["lm_head.weight"].shape
+        logging.warning(
+            f"Warning: Model architectures do not match. "
+            f"Using sub weight space instead.\n"
+            f"Vocab sizes in model 1: {shape_1[0]}, "
+            f"Vocab sizes in model 2: {shape_2[0]}"
+        )
+
+    # Initialize lists to store delta parameters for both models
+    d_vector_1, d_vector_2 = [], []
+
+    # Iterate over keys in the base model's state dictionary with tqdm
+    for key, base_params in tqdm(
+        state_dict_base.items(), desc="Processing keys", unit="key"
+    ):
+        # Only proceed if key exists in both models
+        try:
+            if key not in state_dict_1 or key not in state_dict_2:
+                logging.warning(f"Key {key} not found in one of the models")
+                continue
+        except Exception as e:
+            logging.error(f"Error processing key {key}: {str(e)}")
+
+        # Get the parameters for each model (truncate to num_layers for consistency)
+        params_1 = state_dict_1[key][:num_layers]
+        params_2 = state_dict_2[key][:num_layers]
+
+        # Compute the deltas relative to the base model
+        delta_1 = (params_1 - base_params).view(-1)
+        delta_2 = (params_2 - base_params).view(-1)
+
+        # Accumulate deltas
+        d_vector_1.append(delta_1)
+        d_vector_2.append(delta_2)
+
+    # Clear memory
+    del state_dict_1, state_dict_2, state_dict_base
+
+    logging.info("Concatenating delta vectors...")
+
+    d_vector_1 = torch.cat(d_vector_1)
+    d_vector_2 = torch.cat(d_vector_2)
+
+    if low_precision:
+        logging.info("Quantizing delta vectors to 8-bit precision...")
+        d_vector_1 = quantize_8bit(d_vector_1)
+        d_vector_2 = quantize_8bit(d_vector_2)
+        logging.info("Quantization complete")
+
+    return d_vector_1, d_vector_2

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]

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))