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

fusion_bench/method/simple_average.py

@@ -3,11 +3,16 @@ from copy import deepcopy
 from typing import Dict, List, Mapping, Optional, Union
 
 import torch
-from torch import nn
+from torch import Tensor, nn
 
 from fusion_bench.method.base_algorithm import BaseAlgorithm
 from fusion_bench.mixins import SimpleProfilerMixin, auto_register_config
 from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.models.utils import (
+    get_target_state_dict,
+    load_state_dict_into_target_modules,
+    validate_target_modules_equal,
+)
 from fusion_bench.utils import LazyStateDict
 from fusion_bench.utils.state_dict_arithmetic import (
     state_dict_add,
@@ -21,21 +26,22 @@ log = logging.getLogger(__name__)
 
 
 def simple_average(
-    modules: List[Union[nn.Module, StateDictType]],
-    base_module: Optional[nn.Module] = None,
+    modules: List[Union[nn.Module, StateDictType, Tensor]],
+    base_module: Optional[Union[nn.Module, StateDictType, Tensor]] = None,
 ):
     R"""
     Averages the parameters of a list of PyTorch modules or state dictionaries.
 
     This function takes a list of PyTorch modules or state dictionaries and returns a new module with the averaged parameters, or a new state dictionary with the averaged parameters.
 
+    If `_fusion_bench_target_modules` attribute is set on the modules, only the parameters of the specified target submodules will be averaged.
+
     Args:
-        modules (List[Union[nn.Module, StateDictType]]): A list of PyTorch modules or state dictionaries.
-        base_module (Optional[nn.Module]): A base module to use for the new module. If provided, the averaged parameters will be loaded into this module. If not provided, a new module will be created by copying the first module in the list.
+        modules (List[Union[nn.Module, StateDictType, Tensor]]): A list of PyTorch modules or state dictionaries.
+        base_module (Optional[Union[nn.Module, StateDictType, Tensor]]): A base module to use for the new module. If provided, the averaged parameters will be loaded into this module. If not provided, a new module will be created by copying the first module in the list.
 
     Returns:
-        module_or_state_dict (Union[nn.Module, StateDictType]): A new PyTorch module with the averaged parameters, or a new state dictionary with the averaged parameters.
-
+        module_or_state_dict (Union[nn.Module, StateDictType, Tensor]): A new PyTorch module with the averaged parameters, or a new state dictionary with the averaged parameters.
     Examples:
         >>> import torch.nn as nn
         >>> model1 = nn.Linear(10, 10)
@@ -47,23 +53,42 @@ def simple_average(
         >>> averaged_state_dict = simple_average([state_dict1, state_dict2])
     """
     assert len(modules) > 0, "modules must be a non-empty list"
+    validate_target_modules_equal(modules)
+
     if isinstance(modules[0], nn.Module):
         if base_module is None:
             new_module = deepcopy(modules[0])
         else:
             new_module = base_module
-        state_dict = state_dict_avg([module.state_dict() for module in modules])
-        new_module.load_state_dict(state_dict)
+        state_dict = state_dict_avg(
+            [get_target_state_dict(module) for module in modules]
+        )
+        load_state_dict_into_target_modules(new_module, state_dict)
         return new_module
     elif isinstance(modules[0], Mapping):
-        return state_dict_avg(modules)
+        # if the modules are state dicts
+        # compute the average state dict
+        avg_state_dict = state_dict_avg(modules)
+        # load into base_module if provided
+        if base_module is not None:
+            for k in avg_state_dict:
+                base_module[k] = avg_state_dict[k]
+            return base_module
+        else:
+            return avg_state_dict
+    elif isinstance(modules[0], Tensor):
+        mean_tensor = torch.stack(modules, dim=0).mean(dim=0)
+        if base_module is not None:
+            base_module.data = mean_tensor
+            return base_module
+        else:
+            return mean_tensor
+    else:
+        raise ValueError(f"Unsupported type: {type(modules[0])}")
 
 
 @auto_register_config
-class SimpleAverageAlgorithm(
-    SimpleProfilerMixin,
-    BaseAlgorithm,
-):
+class SimpleAverageAlgorithm(SimpleProfilerMixin, BaseAlgorithm):
     def __init__(self, show_pbar: bool = False, inplace: bool = True, **kwargs):
         """
         Args:
@@ -87,13 +112,20 @@ class SimpleAverageAlgorithm(
         Returns:
             The fused model obtained by simple averaging.
         """
-        if isinstance(modelpool, dict):
+        if not isinstance(modelpool, BaseModelPool):
             modelpool = BaseModelPool(modelpool)
 
         log.info(
             f"Fusing models using simple average on {len(modelpool.model_names)} models. "
             f"models: {modelpool.model_names}"
         )
+        if modelpool.has_instance_models and self.inplace:
+            log.warning(
+                "The model pool contains instance models, and inplace is set to True. "
+                "Therefore, the weights of the first model will be overwritten. "
+                "If this is desired behavior, this warning can be ignored."
+            )
+
         sd: Optional[StateDictType] = None
         forward_model = None
         merged_model_names = []
@@ -106,12 +138,12 @@ class SimpleAverageAlgorithm(
             with self.profile("merge weights"):
                 if sd is None:
                     # Initialize the state dictionary with the first model's state dictionary
-                    sd = model.state_dict()
+                    sd = get_target_state_dict(model)
                     forward_model = model if self.inplace else deepcopy(model)
                 else:
                     # Add the current model's state dictionary to the accumulated state dictionary
                     sd = state_dict_add(
-                        sd, model.state_dict(), show_pbar=self.show_pbar
+                        sd, get_target_state_dict(model), show_pbar=self.show_pbar
                     )
         with self.profile("merge weights"):
             # Divide the accumulated state dictionary by the number of models to get the average
@@ -124,11 +156,13 @@ class SimpleAverageAlgorithm(
             forward_model = deepcopy(forward_model.meta_module).to_empty(
                 device=forward_model._device
             )
-        result = forward_model.load_state_dict(sd, strict=False)
+
+        result = load_state_dict_into_target_modules(forward_model, sd, strict=False)
         if result.unexpected_keys:
             raise ValueError(f"Unexpected keys in state dict: {result.unexpected_keys}")
         if result.missing_keys:
             log.warning(f"Missing keys in state dict: {result.missing_keys}")
+
         # print profile report and log the merged models
         self.print_profile_summary()
         log.info(f"merged {len(merged_model_names)} models:")

fusion_bench/method/task_arithmetic/task_arithmetic.py

@@ -50,7 +50,7 @@ def task_arithmetic_merge(
         finetuned_models (List[nn.Module]): A list of fine-tuned models from which task vectors will be calculated.
         scaling_factor (float): A factor by which the task vectors will be scaled before merging.
         inplace (bool, optional): If True, the pre-trained model will be modified in place.
-                                  If False, a copy of the pre-trained model will be modified. Defaults to True.
+            If False, a copy of the pre-trained model will be modified. Defaults to True.
 
     Returns:
         nn.Module: The pre-trained model with the merged task vectors.

fusion_bench/mixins/lightning_fabric.py

@@ -1,6 +1,7 @@
 import functools
 import logging
 import os
+import sys
 from typing import TYPE_CHECKING, Any, List, Mapping, Optional, TypeVar
 
 import lightning as L
@@ -12,17 +13,32 @@ from omegaconf import DictConfig, OmegaConf
 
 from fusion_bench.constants import RuntimeConstants
 from fusion_bench.utils import import_object
+from fusion_bench.utils.hydra_utils import get_hydra_output_dir
 from fusion_bench.utils.instantiate_utils import instantiate
 
 if TYPE_CHECKING:
     import lightning.fabric.loggers.tensorboard
     from lightning.fabric.strategies import FSDPStrategy
+    from lightning.pytorch.loggers import MLFlowLogger
+    from mlflow.tracking.client import MlflowClient
 
 log = logging.getLogger(__name__)
 
 TensorOrModule = TypeVar("TensorOrModule", torch.Tensor, torch.nn.Module, Any)
 
 
+def _fabric_has_logger(fabric: L.Fabric) -> bool:
+    """
+    Check if the fabric has a logger.
+
+    Args:
+        fabric (L.Fabric): The Lightning Fabric instance.
+    Returns:
+        bool: True if the fabric has a logger, False otherwise.
+    """
+    return fabric._loggers is not None and len(fabric._loggers) > 0
+
+
 def get_policy(*args: str) -> set:
     """
     Get the policy from the provided list of policy names.
@@ -43,6 +59,21 @@ def get_size_based_auto_wrap_policy(*args, **kwargs):
     return policy
 
 
+def _is_mlflow_logger(fabric: L.Fabric) -> bool:
+    """
+    Check if the fabric's logger is an instance of MLFlowLogger.
+
+    Args:
+        fabric (L.Fabric): The Lightning Fabric instance.
+
+    Returns:
+        bool: True if the logger is an instance of MLFlowLogger, False otherwise.
+    """
+    if not _fabric_has_logger(fabric):
+        return False
+    return fabric.logger.__class__.__name__ == "MLFlowLogger"
+
+
 class LightningFabricMixin:
     """
     A mixin class for integrating Lightning Fabric into a project.
@@ -79,8 +110,8 @@ class LightningFabricMixin:
         """
         if self._fabric_instance is None:
             if config.get("fabric", None) is None:
-                log.warning("No fabric configuration found. use default settings.")
-                self._fabric_instance = L.Fabric()
+                log.warning("No fabric configuration found. use default settings. By default, use 1 device.")
+                self._fabric_instance = L.Fabric(devices=1)
             else:
                 self._fabric_instance = instantiate(config.fabric)
             if not _is_using_cli():  # if not using cli, launch the fabric
@@ -123,7 +154,10 @@ class LightningFabricMixin:
         Retrieves the log directory from the fabric's logger.
         """
         if self.fabric is not None and len(self.fabric._loggers) > 0:
-            log_dir = self.fabric.logger.log_dir
+            if hasattr(self.fabric.logger, "log_dir"):
+                log_dir = self.fabric.logger.log_dir
+            else:
+                log_dir = None
 
             # Special handling for SwanLabLogger to get the correct log directory
             if (
@@ -132,6 +166,20 @@ class LightningFabricMixin:
             ):
                 log_dir = self.fabric.logger.save_dir or self.fabric.logger._logdir
 
+            if (
+                log_dir is None
+                and self.fabric.logger.__class__.__name__ == "MLFlowLogger"
+            ):
+                log_dir = self.fabric.logger.save_dir
+                if log_dir is None:
+                    try:
+                        log_dir = self._program.config.path.log_dir
+                    except Exception:
+                        log.error(
+                            "Failed to get log_dir from program config for MLFlowLogger."
+                        )
+                        log_dir = "outputs"
+
             assert log_dir is not None, "log_dir should not be None"
             if self.fabric.is_global_zero and not os.path.exists(log_dir):
                 os.makedirs(log_dir, exist_ok=True)
@@ -246,3 +294,60 @@ class LightningFabricMixin:
         """
         for i, param_group in enumerate(optimizer.param_groups):
             self.fabric.log(name_template.format(i), param_group["lr"], step=step)
+
+    def log_artifact(self, local_path: str, artifact_path: str | None = None):
+        """
+        Logs a file as an artifact to the fabric's logger.
+
+        Args:
+            local_dir: The path to the directory to log as an artifact.
+            artifact_path: The directory within the logger's artifact storage to save the file.
+        """
+        if _is_mlflow_logger(self.fabric):
+            logger: "MLFlowLogger" = self.fabric.logger
+            experiment: "MlflowClient" = logger.experiment
+            experiment.log_artifact(
+                logger.run_id,
+                local_path=local_path,
+                artifact_path=(artifact_path),
+            )
+
+    def log_artifacts(self, local_dir: str, artifact_path: str | None = None):
+        """
+        Logs a directory as artifacts to the fabric's logger.
+
+        Args:
+            local_dir: The path to the directory to log as artifacts.
+            artifact_path: The directory within the logger's artifact storage to save the files.
+        """
+        if _is_mlflow_logger(self.fabric):
+            logger: "MLFlowLogger" = self.fabric.logger
+            experiment: "MlflowClient" = logger.experiment
+            experiment.log_artifacts(
+                logger.run_id,
+                local_dir=local_dir,
+                artifact_path=artifact_path,
+            )
+
+    def finalize(self):
+        """
+        Destructor to ensure proper cleanup of the Lightning Fabric instance.
+        """
+        if self._fabric_instance is None:
+            return
+
+        if _fabric_has_logger(self.fabric) and _is_mlflow_logger(self.fabric):
+            if sys.exc_info()[0] is None:
+                status = "success"
+            else:
+                status = "failed"
+            self.fabric.logger.finalize(status)
+
+        del self._fabric_instance
+        self._fabric_instance = None
+
+    def __del__(self):
+        """
+        Destructor to ensure proper cleanup of the Lightning Fabric instance.
+        """
+        self.finalize()

fusion_bench/mixins/serialization.py

@@ -68,7 +68,7 @@ def auto_register_config(cls):
 
     Behavior:
         - **Parameter Registration**: All non-variadic parameters (excluding ``*args``, ``**kwargs``)
-          from the __init__ method are automatically added to _config_mapping
+            from the __init__ method are automatically added to _config_mapping
         - **Positional Arguments**: Handled in order and mapped to corresponding parameter names
         - **Keyword Arguments**: Processed after positional arguments, overriding any conflicts
         - **Default Values**: Applied when parameters are not provided via arguments

fusion_bench/modelpool/base_pool.py

@@ -7,11 +7,12 @@ from omegaconf import DictConfig, OmegaConf, UnsupportedValueType
 from torch import nn
 from torch.utils.data import Dataset
 
-from fusion_bench import TorchModelType
+from fusion_bench import StateDictType, TorchModelType
 from fusion_bench.mixins import BaseYAMLSerializable, HydraConfigMixin
 from fusion_bench.utils import (
     ValidationError,
     instantiate,
+    state_dict_sub,
     timeit_context,
     validate_model_name,
 )
@@ -57,6 +58,10 @@ class BaseModelPool(
         **kwargs,
     ):
         if isinstance(models, List):
+            log.debug(
+                "Initializing BaseModelPool with a list of models. "
+                "Converting to a dictionary with integer string keys."
+            )
             models = {str(model_idx): model for model_idx, model in enumerate(models)}
 
         if isinstance(models, dict):
@@ -81,6 +86,22 @@ class BaseModelPool(
         self._test_datasets = test_datasets
         super().__init__(**kwargs)
 
+    @property
+    def has_instance_models(self) -> bool:
+        """
+        Check if the model pool contains any pre-instantiated models.
+
+        Attention:
+            Some algorithms may modify the models in-place if they are pre-instantiated.
+
+        Returns:
+            bool: True if there are pre-instantiated models, False otherwise.
+        """
+        for model_cfg in self._models.values():
+            if isinstance(model_cfg, nn.Module):
+                return True
+        return False
+
     @property
     def has_pretrained(self) -> bool:
         """
@@ -329,6 +350,21 @@ class BaseModelPool(
         for model_name in self.model_names:
             yield model_name, self.load_model(model_name)
 
+    def load_pretrained_model_and_task_vectors(
+        self,
+    ) -> Tuple[TorchModelType, List[StateDictType]]:
+        pretrained_model = self.load_pretrained_model()
+
+        task_vectors = []
+        for model_name in self.model_names:
+            finetuned_model = self.load_model(model_name)
+            task_vector = state_dict_sub(
+                finetuned_model.state_dict(), pretrained_model.state_dict()
+            )
+            task_vectors.append(task_vector)
+
+        return pretrained_model, task_vectors
+
     @property
     def has_train_dataset(self) -> bool:
         """

fusion_bench/modelpool/convnext_for_image_classification.py

@@ -98,7 +98,7 @@ class ConvNextForImageClassificationPool(BaseModelPool):
     - Load ConvNeXt models either from a pretrained checkpoint or from config.
     - Optionally adapt the classifier head to match dataset classnames.
     - Override `forward` to return logits for consistent interfaces within
-      FusionBench.
+        FusionBench.
 
     See `fusion_bench.modelpool.resnet_for_image_classification` for a closely
     related ResNet-based pool with analogous behavior.
@@ -161,6 +161,9 @@ class ConvNextForImageClassificationPool(BaseModelPool):
         ).logits
         model.original_forward = original_forward
 
+        # Mark ConvNeXt layers for FusionBench fusion
+        model._fusion_bench_target_modules = ["convnext"]
+
         return model
 
     @override
@@ -180,7 +183,7 @@ class ConvNextForImageClassificationPool(BaseModelPool):
         - The ConvNeXt model via `model.save_pretrained`.
         - The paired image processor via `AutoImageProcessor.save_pretrained`.
         - If `algorithm_config` is provided and on rank-zero, a README model card
-          documenting the FusionBench configuration.
+            documenting the FusionBench configuration.
         """
         model.save_pretrained(path)
         self.load_processor().save_pretrained(path)

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