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

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/mixins/lightning_fabric.py

@@ -10,6 +10,7 @@ from lightning.fabric.loggers import TensorBoardLogger
 from lightning.fabric.utilities.rank_zero import rank_zero_only
 from omegaconf import DictConfig, OmegaConf
 
+from fusion_bench.constants import RuntimeConstants
 from fusion_bench.utils import import_object
 from fusion_bench.utils.instantiate_utils import instantiate
 
@@ -206,14 +207,7 @@ class LightningFabricMixin:
         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(
-            "fast_dev_run", False
-        ):
-            return True
-        else:
-            return False
+        return RuntimeConstants().debug
 
     def log(self, name: str, value: Any, step: Optional[int] = None):
         """

fusion_bench/mixins/openclip_classification.py

@@ -1,11 +1,165 @@
+import functools
 import logging
+from typing import TYPE_CHECKING, Callable, Dict, Iterator, List, Literal, Optional
 
+import torch
+from omegaconf import DictConfig
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from fusion_bench.dataset.clip_dataset import CLIPDataset
 from fusion_bench.mixins import LightningFabricMixin
-from fusion_bench.models.open_clip import ImageClassifier, ImageEncoder
+from fusion_bench.modelpool import OpenCLIPVisionModelPool
+from fusion_bench.models.open_clip import (
+    ClassificationHead,
+    ImageClassifier,
+    ImageEncoder,
+)
+from fusion_bench.utils.data import InfiniteDataLoader
 
 log = logging.getLogger(__name__)
 
 
 class OpenCLIPClassificationMixin(LightningFabricMixin):
+
     _train_processor = None
     _test_processor = None
+    dataloader_kwargs: DictConfig
+    modelpool: OpenCLIPVisionModelPool
+    zero_shot_heads: Dict[str, ClassificationHead] = {}
+
+    def _init_processor(self, encoder: Optional["ImageEncoder"] = None):
+        """
+        Initialize the CLIP processors for training and testing.
+        """
+        if encoder is None:
+            encoder: "ImageEncoder" = self.modelpool.load_pretrained_or_first_model()
+        self._train_processor = encoder.train_preprocess
+        self._test_processor = encoder.val_preprocess
+        return self._train_processor, self._test_processor
+
+    def get_clip_processor(self, stage: Literal["train", "test"]):
+        """
+        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 stage == "train":
+            if self._train_processor is None:
+                self._init_processor()
+            return self._train_processor
+        elif stage == "test":
+            if self._test_processor is None:
+                self._init_processor()
+            return self._test_processor
+        else:
+            raise ValueError(f"Invalid stage: {stage}")
+
+    def setup_zero_shot_classification_head(
+        self,
+        task_names: Optional[List[str]] = None,
+        freeze: bool = True,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        # check task names consistency across 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.")
+
+        for task in tqdm(
+            self.modelpool.model_names if task_names is None else task_names,
+            "Setting up zero-shot classification head",
+            disable=not self.fabric.is_global_zero,
+        ):
+            head = self.modelpool.load_classification_head(task)
+            if freeze:
+                head.requires_grad_(False)
+            if dtype is not None:
+                head = head.to(dtype=dtype)
+            self.zero_shot_heads[task] = self.to_device(head)
+
+    def set_clip_processor(self, stage: Literal["train", "test"], processor: Callable):
+        """
+        Set the CLIP processor for a specific stage.
+
+        Args:
+            stage (Literal["train", "test"]): The stage for which to set the processor.
+            processor (Callable): The CLIP processor to set.
+        """
+        if stage == "train":
+            self._train_processor = processor
+        elif stage == "test":
+            self._test_processor = processor
+        else:
+            raise ValueError(f"Invalid stage: {stage}")
+
+    @functools.cache
+    def get_shuffled_test_loader_iter(
+        self,
+        task: str,
+        batch_size: Optional[int] = None,
+        num_workers: Optional[int] = None,
+        **loader_kwargs,
+    ) -> Iterator:
+        """
+        Get an iterator for a shuffled test DataLoader.
+
+        This method creates a DataLoader for the test dataset of the specified task,
+        with shuffling enabled. It allows for optional customization of batch size,
+        number of workers, and other DataLoader keyword arguments.
+
+        Args:
+            task (str): The task identifier for which the test dataset is to be loaded.
+            batch_size (Optional[int]): The batch size to use for the DataLoader. If None, the default batch size is used.
+            num_workers (Optional[int]): The number of worker processes to use for data loading. If None, the default number of workers is used.
+            **loader_kwargs: Additional keyword arguments to pass to the DataLoader.
+
+        Returns:
+            Iterator: An iterator over the shuffled test DataLoader.
+        """
+        # get dataloader kwargs
+        dataloader_kwargs = self.dataloader_kwargs.copy()
+        dataloader_kwargs["shuffle"] = True
+        if batch_size is not None:
+            dataloader_kwargs["batch_size"] = batch_size
+        if num_workers is not None:
+            dataloader_kwargs["num_workers"] = num_workers
+        dataloader_kwargs.update(loader_kwargs)
+
+        # get the test dataset
+        clip_dataset = CLIPDataset(
+            self.modelpool.load_test_dataset(task),
+            processor=self.get_clip_processor(stage="test"),
+        )
+        # create the dataloader
+        loader = DataLoader(clip_dataset, **dataloader_kwargs)
+        loader = self.fabric.setup_dataloaders(loader)
+        return iter(InfiniteDataLoader(loader))
+
+    def compute_logits(
+        self,
+        module: ImageClassifier,
+        images,
+        task: str,
+    ):
+        """
+        Compute the logits for a batch of images using the provided module and task.
+
+        Args:
+            module (ImageClassifier): The image classification module to use for computing logits.
+            images (torch.Tensor): The batch of images for which to compute logits.
+            task (str): The task identifier to specify which classification head to use.
+
+        Returns:
+            torch.Tensor: The computed logits for the input images.
+        """
+        if len(self.zero_shot_heads) == 0:
+            self.setup_zero_shot_classification_head()
+        task_head = self.zero_shot_heads[task]
+        features = module(images)
+        logits = task_head(features)
+        return logits

fusion_bench/modelpool/base_pool.py

@@ -7,6 +7,7 @@ from omegaconf import DictConfig, OmegaConf, UnsupportedValueType
 from torch import nn
 from torch.utils.data import Dataset
 
+from fusion_bench import TorchModelType
 from fusion_bench.mixins import BaseYAMLSerializable, HydraConfigMixin
 from fusion_bench.utils import (
     ValidationError,

fusion_bench/modelpool/openclip_vision/modelpool.py

@@ -1,7 +1,7 @@
 import logging
 import pickle
 import sys
-from typing import Callable, Optional, Union, cast
+from typing import Callable, Optional, Union, cast, override
 
 import torch
 from datasets import load_dataset
@@ -41,8 +41,8 @@ def _check_and_redirect_open_clip_modeling():
             )
 
     try:
-        import src
-        import src.modeling
+        import src  # type: ignore
+        import src.modeling  # type: ignore
     except ImportError:
         if "src" not in sys.modules:
             # redirect the import of `src` to `fusion_bench.models.open_clip`
@@ -114,6 +114,7 @@ class OpenCLIPVisionModelPool(BaseModelPool):
             self._test_processor = encoder.val_preprocess
         return self._test_processor
 
+    @override
     def load_model(
         self, model_name_or_config: Union[str, DictConfig], *args, **kwargs
     ) -> ImageEncoder:
@@ -210,6 +211,8 @@ class OpenCLIPVisionModelPool(BaseModelPool):
         - A string, which is the path to the model checkpoint in pickle format. Load directly using `torch.load`.
         - Default, load the model using `instantiate` from hydra.
         """
+        if self._classification_heads is None:
+            raise ValueError("No classification heads are defined in the model pool.")
         if (
             isinstance(model_name_or_config, str)
             and model_name_or_config in self._classification_heads
@@ -222,6 +225,8 @@ class OpenCLIPVisionModelPool(BaseModelPool):
         return head
 
     def load_train_dataset(self, dataset_name: str, *args, **kwargs):
+        if self._train_datasets is None:
+            raise ValueError("No train datasets are defined in the model pool.")
         dataset_config = self._train_datasets[dataset_name]
         if isinstance(dataset_config, str):
             log.info(
@@ -233,6 +238,8 @@ class OpenCLIPVisionModelPool(BaseModelPool):
         return dataset
 
     def load_val_dataset(self, dataset_name: str, *args, **kwargs):
+        if self._val_datasets is None:
+            raise ValueError("No val datasets are defined in the model pool.")
         dataset_config = self._val_datasets[dataset_name]
         if isinstance(dataset_config, str):
             log.info(
@@ -244,6 +251,8 @@ class OpenCLIPVisionModelPool(BaseModelPool):
         return dataset
 
     def load_test_dataset(self, dataset_name: str, *args, **kwargs):
+        if self._test_datasets is None:
+            raise ValueError("No test datasets are defined in the model pool.")
         dataset_config = self._test_datasets[dataset_name]
         if isinstance(dataset_config, str):
             log.info(

fusion_bench/models/masks/mask_model.py

@@ -113,21 +113,27 @@ class MaskModel(ParameterDictModel):
     def get_distribution(
         self,
         mask_type: Literal["discrete", "continuous"],
+        temperature: float = 0.5,
         **kwargs,
     ):
         return {
-            name: self._param_to_distribution(param, mask_type=mask_type, **kwargs)
+            name: self._param_to_distribution(
+                param, mask_type=mask_type, temperature=temperature, **kwargs
+            )
             for name, param in self.named_parameters()
         }
 
     def sample_mask(
         self,
         mask_type: Literal["discrete", "continuous"] = "discrete",
+        temperature: float = 0.5,
         **kwargs,
     ):
         mask = {}
         for name, param in self.named_parameters():
-            dist = self._param_to_distribution(param, mask_type, **kwargs)
+            dist = self._param_to_distribution(
+                param, mask_type, temperature=temperature, **kwargs
+            )
             if mask_type == "discrete":
                 mask[name] = dist.sample()
             elif mask_type == "continuous":

fusion_bench/models/open_clip/modeling.py

@@ -1,3 +1,25 @@
+"""
+OpenCLIP model wrappers used by FusionBench.
+
+This module provides lightweight `torch.nn.Module` wrappers around OpenCLIP
+components that are commonly used throughout FusionBench experiments:
+
+- `ImageEncoder`: loads an OpenCLIP image encoder and exposes `encode_image`.
+- `ClassificationHead`: a linear head optionally normalizing inputs.
+- `ImageClassifier` / `MultiHeadImageClassifier`: convenience compositions.
+
+Note:
+    This module requires the optional dependency `open_clip_torch`.
+"""
+
+from fusion_bench.utils.packages import is_open_clip_available
+
+if not is_open_clip_available():
+    raise ImportError(
+        "open_clip is not installed. Please install it with `pip install open_clip_torch`."
+    )
+
+from pathlib import Path
 from typing import Callable, List
 
 import open_clip
@@ -10,6 +32,19 @@ from .variables_and_paths import CACHEDIR, MODELS, OPENCLIP_CACHEDIR
 
 class ImageEncoder(torch.nn.Module):
     R"""
+    OpenCLIP image encoder wrapper.
+
+    This class loads an OpenCLIP model by name and exposes a forward pass that
+    returns image embeddings via `model.encode_image`.
+
+    Args:
+        model_name: A model name supported by `open_clip`. FusionBench also
+            supports suffixes:
+            - ``"__pretrained__<tag>"`` to select a specific pretrained weights tag.
+            - ``"__init__"`` to use random initialization.
+        keep_lang: If False (default), removes the text encoder (when present)
+            to reduce memory usage.
+
     Examples:
 
         load the image encoder for a given model name
@@ -18,7 +53,7 @@ class ImageEncoder(torch.nn.Module):
         >>> image_encoder = ImageEncoder(model_name="ViT-B-32")
     """
 
-    def __init__(self, model_name: str, keep_lang=False):
+    def __init__(self, model_name: str, keep_lang: bool = False):
         super().__init__()
         assert (
             model_name in MODELS
@@ -42,22 +77,26 @@ class ImageEncoder(torch.nn.Module):
 
         self.cache_dir = CACHEDIR
 
+        # if `keep_lang` is False, remove the text encoder to save memory
         if not keep_lang and hasattr(self.model, "transformer"):
             delattr(self.model, "transformer")
 
-    def forward(self, images):
+    def forward(self, images: Tensor) -> Tensor:
+        """Encode a batch of images into embedding vectors."""
         assert self.model is not None
         return self.model.encode_image(images)
 
-    def __call__(self, inputs):
+    def __call__(self, inputs: Tensor) -> Tensor:
         return self.forward(inputs)
 
-    def save(self, filename):
+    def save(self, filename: str) -> None:
+        """Serialize this module to disk."""
         print(f"Saving image encoder to {filename}")
         utils.torch_save(self, filename)
 
     @classmethod
-    def load(cls, model_name, filename):
+    def load(cls, model_name: str, filename: str | Path):
+        """Load a saved encoder state dict into a freshly constructed encoder."""
         print(f"Loading image encoder from {filename}")
 
         state_dict = torch.load(filename, map_location="cpu")
@@ -68,6 +107,15 @@ class ImageEncoder(torch.nn.Module):
 
 
 class ClassificationHead(torch.nn.Linear):
+    """A linear classification head with optional input normalization.
+
+    Args:
+        normalize: If True, L2-normalize inputs along the last dimension before
+            applying the linear projection.
+        weights: Weight matrix of shape (num_classes, feature_dim).
+        biases: Optional bias vector of shape (num_classes,).
+    """
+
     def __init__(
         self,
         normalize: bool,
@@ -85,6 +133,7 @@ class ClassificationHead(torch.nn.Linear):
             self.bias = torch.nn.Parameter(torch.zeros_like(self.bias))
 
     def forward(self, inputs: Tensor):
+        """Compute logits from input features."""
         if self.normalize:
             inputs = inputs / inputs.norm(dim=-1, keepdim=True)
         return super().forward(inputs)
@@ -93,11 +142,13 @@ class ClassificationHead(torch.nn.Linear):
         return self.forward(inputs)
 
     def save(self, filename):
+        """Serialize this head to disk."""
         print(f"Saving classification head to {filename}")
         utils.torch_save(self, filename, save_state_dict=False)
 
     @classmethod
     def load(cls, filename):
+        """Load a serialized `ClassificationHead` instance from disk."""
         # print(f"Loading classification head from {filename}")
         return utils.torch_load(filename)
 
@@ -106,6 +157,8 @@ class ImageClassifier(torch.nn.Module):
     train_preprocess: Callable
     val_preprocess: Callable
 
+    """Convenience module combining an `ImageEncoder` and a `ClassificationHead`."""
+
     def __init__(
         self,
         image_encoder: ImageEncoder,
@@ -119,10 +172,12 @@ class ImageClassifier(torch.nn.Module):
             self.val_preprocess = self.image_encoder.val_preprocess
 
     def freeze_head(self):
+        """Disable gradient computation for the classification head."""
         self.classification_head.weight.requires_grad_(False)
         self.classification_head.bias.requires_grad_(False)
 
     def forward(self, inputs: Tensor):
+        """Run encoder then head and return logits."""
         features = self.image_encoder(inputs)
         outputs = self.classification_head(features)
         return outputs
@@ -131,16 +186,20 @@ class ImageClassifier(torch.nn.Module):
         return self.forward(inputs)
 
     def save(self, filename):
+        """Serialize this module to disk."""
         print(f"Saving image classifier to {filename}")
         utils.torch_save(self, filename)
 
     @classmethod
     def load(cls, filename):
+        """Load a serialized `ImageClassifier` instance from disk."""
         print(f"Loading image classifier from {filename}")
         return utils.torch_load(filename)
 
 
 class MultiHeadImageClassifier(torch.nn.Module):
+    """Image encoder with multiple task-specific classification heads."""
+
     def __init__(
         self,
         image_encoder: ImageEncoder,
@@ -154,11 +213,13 @@ class MultiHeadImageClassifier(torch.nn.Module):
             self.val_preprocess = self.image_encoder.val_preprocess
 
     def freeze_head(self):
+        """Disable gradient computation for all heads."""
         for idx in range(len(self.classification_heads)):
             self.classification_heads[idx].weight.requires_grad_(False)
             self.classification_heads[idx].bias.requires_grad_(False)
 
     def forward(self, inputs, head_idx):
+        """Run encoder then the selected head and return logits."""
         features = self.image_encoder(inputs)
         outputs = self.classification_heads[head_idx](features)
         return outputs
@@ -167,10 +228,12 @@ class MultiHeadImageClassifier(torch.nn.Module):
         return self.forward(inputs, head_idx)
 
     def save(self, filename):
+        """Serialize this module to disk."""
         print(f"Saving image classifier to {filename}")
         utils.torch_save(self, filename)
 
     @classmethod
     def load(cls, filename):
+        """Load a serialized `MultiHeadImageClassifier` instance from disk."""
         print(f"Loading image classifier from {filename}")
         return utils.torch_load(filename)