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

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/hf_clip.py

@@ -62,16 +62,36 @@ class HFCLIPClassifier(nn.Module):
             persistent=False,
         )
 
+    # NOTE:
+    # The property setters seems not to work properly with `nn.Module` attributes.
+    # So avoid using them in practice.
+    # To set the text or vision model, directly access the attributes.
+    # For example:
+    #    classifier.clip_model.text_model = new_text_model
+    # or
+    #    classifier.clip_model.vision_model = new_vision_model
+    # reference: https://github.com/pytorch/pytorch/issues/52664
+
     @property
     def text_model(self):
         """Get the text model component of CLIP."""
         return self.clip_model.text_model
 
+    @text_model.setter
+    def text_model(self, model: nn.Module):
+        """Set the text model component of CLIP."""
+        self.clip_model.text_model = model
+
     @property
     def vision_model(self):
         """Get the vision model component of CLIP."""
         return self.clip_model.vision_model
 
+    @vision_model.setter
+    def vision_model(self, model: nn.Module):
+        """Set the vision model component of CLIP."""
+        self.clip_model.vision_model = model
+
     def set_classification_task(
         self,
         classnames: List[str],

fusion_bench/models/modulator/__init__.py

@@ -0,0 +1 @@
+from .base import ModulatedModel, TaskModulator

fusion_bench/models/modulator/base.py

@@ -0,0 +1,123 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Generic, Optional
+
+import torch
+from torch import nn
+
+from fusion_bench import TorchModelType
+
+log = logging.getLogger(__name__)
+
+
+class ModulatedModel(nn.Module, Generic[TorchModelType]):
+    """
+    A model wrapper that uses task-specific modulators to adapt a shared backbone
+    for different tasks.
+
+    The model maintains a shared backbone and task-specific modulators. During forward pass,
+    the appropriate modulator is applied based on the current task.
+    """
+
+    _current_task: Optional[str] = None
+
+    def __init__(
+        self,
+        backbone: TorchModelType,
+        modulators: Dict[str, "TaskModulator[TorchModelType]"],
+    ):
+        super().__init__()
+        self.backbone = backbone
+        self.modulators = nn.ModuleDict(modulators)
+
+    def add_modulator(self, task_name: str, modulator: "TaskModulator[TorchModelType]"):
+        """Add a new task-specific modulator."""
+        if task_name in self.modulators:
+            raise ValueError(f"Modulator for task '{task_name}' already exists.")
+        self.modulators[task_name] = modulator
+
+    def remove_modulator(self, task_name: str):
+        """Remove an existing task-specific modulator."""
+        if task_name not in self.modulators:
+            raise ValueError(f"Modulator for task '{task_name}' does not exist.")
+        if self._current_task == task_name:
+            log.warning(
+                f"Removing modulator for current task '{task_name}'. "
+                "This will make unset the current task unpredictable."
+            )
+        del self.modulators[task_name]
+
+    def set_task(self, task_name: str):
+        """Set the current task for inference."""
+        if task_name not in self.modulators:
+            raise ValueError(
+                f"Task '{task_name}' not found in modulators. Available tasks: {list(self.modulators.keys())}"
+            )
+        if self._current_task == task_name:
+            return
+
+        # unset previous task
+        if self._current_task is not None:
+            self.modulators[self._current_task].remove(self)
+            assert (
+                self._current_task is None
+            ), "Current task should be None after removal."
+
+        # set new task
+        self.modulators[task_name].apply(self)
+        self._current_task = task_name
+
+    @property
+    def current_task(self) -> Optional[str]:
+        """Get the current task name."""
+        return self._current_task
+
+    def forward(self, *args, **kwargs) -> Any:
+        """
+        Forward pass with task-specific modulation.
+
+        Args:
+            *args: Positional arguments for the backbone model
+            **kwargs: Keyword arguments for the backbone model
+
+        Returns:
+            Model output after applying task-specific modulation
+        """
+        if self._current_task is None:
+            raise ValueError(
+                "No task specified. Set current_task or provide 'task' argument."
+            )
+
+        return self.backbone(*args, **kwargs)
+
+
+class TaskModulator(nn.Module, Generic[TorchModelType], ABC):
+    """
+    Lightweight, task-specific parameterization that modulates
+    a shared representation.
+
+    This is the base class for all task modulators. Subclasses should implement
+    the `apply` method to define how the modulator adapts the backbone model
+    for a specific task.
+    """
+
+    @abstractmethod
+    def apply(self, modulated_model: "ModulatedModel[TorchModelType]"):
+        """
+        Apply task-specific modulation to the backbone model.
+
+        Args:
+            modulated_model: The modulated model
+        """
+        raise NotImplementedError("Subclasses must implement the apply method.")
+
+    @abstractmethod
+    def remove(self, modulated_model: "ModulatedModel[TorchModelType]"):
+        """
+        Remove task-specific modulation from the backbone model.
+        This is called when switching tasks.
+
+        Args:
+            modulated_model: The modulated model
+        """
+        raise NotImplementedError("Subclasses must implement the remove method.")

fusion_bench/models/open_clip/modeling.py

@@ -1,3 +1,17 @@
+"""
+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():
@@ -5,6 +19,7 @@ if not is_open_clip_available():
         "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
@@ -17,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
@@ -25,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
@@ -49,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")
@@ -75,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,
@@ -92,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)
@@ -100,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)
 
@@ -113,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,
@@ -126,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
@@ -138,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,
@@ -161,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
@@ -174,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)

fusion_bench/models/open_clip/utils.py

@@ -77,7 +77,16 @@ def torch_load_old(save_path: str, device=None):
     return classifier
 
 
-def torch_save(model, save_path, save_state_dict=True):
+def torch_save(model: torch.nn.Module, save_path: str, save_state_dict: bool = True):
+    """
+    Save a model to disk.
+
+    Args:
+        model: The model to save.
+        save_path (str): The path to save the model to.
+        save_state_dict (bool): Whether to save the state dict of the model (weights only).
+            If False, the entire model object is saved. Default is True.
+    """
     # TODO: hacky way to save state dict
     if save_state_dict and isinstance(model, torch.nn.Module):
         model = model.state_dict()
@@ -86,7 +95,9 @@ def torch_save(model, save_path, save_state_dict=True):
     torch.save(model, save_path)
 
 
-def torch_load(save_path, device=None):
+def torch_load(
+    save_path: str, device: Optional[torch.device] = None
+) -> torch.nn.Module:
     model = torch.load(save_path, map_location="cpu")
     if device is not None:
         model = model.to(device)

fusion_bench/models/parameter_dict.py

@@ -1,12 +1,12 @@
-from typing import List, Mapping, Optional, Tuple
+from typing import Iterator, List, Mapping, Optional, Tuple, Union
 
 import torch
 from torch import nn
 
-__all__ = "ParamterDictModel"
+__all__ = ["ParameterDictModel"]
 
 
-def _set_attr(
+def set_nested_attr(
     obj,
     names: List[str],
     val,
@@ -27,7 +27,7 @@ def _set_attr(
     else:
         if check_parent and not hasattr(obj, names[0]):
             setattr(obj, names[0], parent_builder())
-        _set_attr(
+        set_nested_attr(
             getattr(obj, names[0]),
             names[1:],
             val,
@@ -36,7 +36,7 @@ def _set_attr(
         )
 
 
-def has_attr(obj, names: List[str]):
+def has_nested_attr(obj, names: List[str]):
     """
     Checks if an attribute exists in an object recursively.
 
@@ -50,26 +50,49 @@ def has_attr(obj, names: List[str]):
     if len(names) == 1:
         return hasattr(obj, names[0])
     else:
-        return has_attr(getattr(obj, names[0]), names[1:])
+        if not hasattr(obj, names[0]):
+            return False
+        return has_nested_attr(getattr(obj, names[0]), names[1:])
 
 
 class ParameterDictModel(nn.Module):
     """
-    This model is used to create a model with parameters from a dictionary.
-    It behaves like a normal `nn.ParameterDict`, but support keys with dots.
+    A module that stores parameters in a nested dictionary structure.
+
+    This model behaves similarly to `nn.ParameterDict`, but supports hierarchical keys
+    with dots (e.g., "layer1.weight"). Parameters are stored as nested attributes,
+    allowing for structured parameter access and manipulation.
+
+    Example:
+        >>> params = {
+        ...     "encoder.weight": nn.Parameter(torch.randn(10, 5)),
+        ...     "decoder.bias": nn.Parameter(torch.randn(5)),
+        ... }
+        >>> model = ParameterDictModel(params)
+        >>> model["encoder.weight"].shape
+        torch.Size([10, 5])
+        >>> "encoder.weight" in model
+        True
     """
 
     def __init__(
         self,
-        parameters: Optional[Mapping[str, nn.Parameter]] = None,
-    ):
+        parameters: Optional[Mapping[str, Union[nn.Parameter, torch.Tensor]]] = None,
+    ) -> None:
+        """
+        Args:
+            parameters: Optional mapping of parameter names to parameter tensors.
+                Keys can contain dots to create nested structures.
+                Values must be `nn.Parameter` or `nn.Buffer` instances.
+        """
+
         super().__init__()
         if parameters is not None:
             for name, param in parameters.items():
                 assert isinstance(
                     param, (nn.Parameter, nn.Buffer)
                 ), f"{name} is not a nn.Parameter or nn.Buffer"
-                _set_attr(
+                set_nested_attr(
                     self,
                     name.split("."),
                     param,
@@ -77,12 +100,13 @@ class ParameterDictModel(nn.Module):
                     parent_builder=__class__,
                 )
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """
         Generate a string representation of the model's parameters.
 
         Returns:
-            str: A string representation of the model's parameters.
+            A string representation of the model's parameters in the format:
+            "ParameterDictModel(name1: shape1, name2: shape2, ...)"
         """
         param_reprs = []
         for name, param in self.named_parameters():
@@ -90,32 +114,98 @@ class ParameterDictModel(nn.Module):
             param_reprs.append(param_repr)
         return f"{self.__class__.__name__}({', '.join(param_reprs)})"
 
-    def __getitem__(self, key: str):
-        if not has_attr(self, key.split(".")):
+    def __iter__(self) -> Iterator[str]:
+        """
+        Iterate over the model's parameters.
+
+        Yields:
+            Tuples of (parameter name, parameter tensor).
+        """
+        yield from self.keys()
+
+    def __getitem__(
+        self, key: str
+    ) -> Union[nn.Parameter, torch.Tensor, "ParameterDictModel"]:
+        """
+        Retrieve a parameter or nested submodule by key.
+
+        Args:
+            key: Parameter name, which can contain dots for nested access.
+
+        Returns:
+            The parameter, tensor, or nested ParameterDictModel at the specified key.
+
+        Raises:
+            KeyError: If the key is not found in the model.
+        """
+        assert isinstance(
+            key, str
+        ), f"Key must be a string, but got {type(key)}: {key}."
+        if not has_nested_attr(self, key.split(".")):
             raise KeyError(f"Key {key} not found in {self}")
-        key = key.split(".")
+        key_parts = key.split(".")
         obj = self
-        for k in key:
+        for k in key_parts:
             obj = getattr(obj, k)
         return obj
 
-    def __setitem__(self, key: str, value: nn.Parameter):
-        if not has_attr(self, key.split(".")):
-            _set_attr(self, key.split("."), value, check_parent=True)
+    def __setitem__(self, key: str, value: Union[nn.Parameter, torch.Tensor]) -> None:
+        """
+        Set a parameter at the specified key, creating nested structure if needed.
+
+        Args:
+            key: Parameter name, which can contain dots for nested assignment.
+            value: Parameter or tensor to assign.
+        """
+        if not has_nested_attr(self, key.split(".")):
+            set_nested_attr(self, key.split("."), value, check_parent=True)
         else:
-            _set_attr(self, key.split("."), value, check_parent=False)
+            set_nested_attr(self, key.split("."), value, check_parent=False)
+
+    def __contains__(self, key: str) -> bool:
+        """
+        Check if a parameter key exists in the model.
 
-    def __contains__(self, key: str):
-        return has_attr(self, key.split("."))
+        Args:
+            key: Parameter name, which can contain dots for nested checking.
+
+        Returns:
+            True if the key exists, False otherwise.
+        """
+        return has_nested_attr(self, key.split("."))
 
     def keys(self):
-        return [name for name, _ in self.named_parameters()]
+        """
+        Return a list of all parameter names in the model.
+
+        Returns:
+            List of parameter names (including nested names with dots).
+        """
+        return self.state_dict().keys()
+
+    def items(self):
+        """
+        Return a list of (name, parameter) tuples.
+
+        Returns:
+            List of tuples containing parameter names and their corresponding tensors.
+        """
+        yield from self.state_dict().items()
 
-    def items(self) -> List[Tuple[str, nn.Parameter]]:
-        return [(name, self[name]) for name in self.keys()]
+    def values(self):
+        """
+        Return a list of all parameter values in the model.
 
-    def values(self) -> List[nn.Parameter]:
-        return [self[name] for name in self.keys()]
+        Returns:
+            List of parameter tensors.
+        """
+        yield from self.state_dict().values()
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """
+        Return the number of parameters in the model.
+
+        Returns:
+            The total number of parameters.
+        """
         return len(self.keys())