fusion-bench 0.2.23__py3-none-any.whl → 0.2.25__py3-none-any.whl

fusion_bench/mixins/serialization.py

@@ -6,6 +6,7 @@ from inspect import Parameter, _ParameterKind
 from pathlib import Path
 from typing import Dict, Mapping, Optional, Union
 
+from bidict import MutableBidict, bidict
 from omegaconf import DictConfig, OmegaConf
 
 from fusion_bench.constants import FUSION_BENCH_VERSION
@@ -15,22 +16,29 @@ from fusion_bench.utils.instantiate_utils import set_print_function_call
 log = logging.getLogger(__name__)
 
 __all__ = [
-    "YAMLSerializationMixin",
     "auto_register_config",
+    "YAMLSerializationMixin",
     "BaseYAMLSerializable",
 ]
 
 
-def _get_attr_name(config_mapping: Mapping[str, str], param_name):
-    for attr_name, p in config_mapping.items():
-        if p == param_name:
-            return attr_name
-    else:
-        raise ValueError(f"Parameter {param_name} not found in config mapping.")
+def _set_attr(self, param_name: str, value):
+    """
+    Set an attribute on the object using the parameter name from config mapping.
+
+    This function looks up the corresponding attribute name for the given parameter
+    name using the object's _config_mapping, then sets that attribute to the
+    specified value. It also logs the operation for debugging purposes.
 
+    Args:
+        self: The object instance to set the attribute on.
+        param_name (str): The parameter name (config key) to map to an attribute.
+        value: The value to assign to the attribute.
 
-def _set_attr(self, param_name: str, value):
-    attr_name = _get_attr_name(self._config_mapping, param_name)
+    Raises:
+        ValueError: If the parameter name is not found in the config mapping.
+    """
+    attr_name = self._config_mapping.inverse[param_name]
     log.debug(f"set {attr_name} to {value}. Parameter name: {param_name}")
     setattr(self, attr_name, value)
 
@@ -59,37 +67,16 @@ def auto_register_config(cls):
             functionality and modified __init__ behavior.
 
     Behavior:
-        - **Parameter Registration**: All non-variadic parameters (excluding *args, **kwargs)
+        - **Parameter Registration**: All non-variadic parameters (excluding ``*args``, ``**kwargs``)
           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
         - **Attribute Setting**: All parameters become instance attributes accessible via dot notation
 
-    Example:
-        ```python
-        @auto_register_config
-        class MyAlgorithm(BaseYAMLSerializable):
-            def __init__(self, learning_rate: float = 0.001, batch_size: int = 32, model_name: str = "default", **kwargs):
-                super().__init__(**kwargs)
-
-        # All instantiation methods work automatically:
-        algo1 = MyAlgorithm(0.01, 64)  # positional args
-        algo2 = MyAlgorithm(learning_rate=0.01, model_name="bert")  # keyword args
-        algo3 = MyAlgorithm(0.01, batch_size=128, model_name="gpt")  # mixed args
-
-        # Attributes are automatically set and can be serialized:
-        print(algo1.learning_rate)  # 0.01
-        print(algo1.batch_size)     # 64
-        print(algo1.model_name)     # "default" (from default value)
-
-        config = algo1.config
-        # DictConfig({'_target_': 'MyAlgorithm', 'learning_rate': 0.01, 'batch_size': 64, 'model_name': 'default'})
-        ```
-
     Note:
         - The decorator wraps the original __init__ method while preserving its signature for IDE support
-        - Parameters with *args or **kwargs signatures are ignored during registration
+        - Parameters with ``*args`` or ``**kwargs`` signatures are ignored during registration
         - The attributes are auto-registered, then the original __init__ method is called,
         - Type hints, method name, and other metadata are preserved using functools.wraps
         - This decorator is designed to work seamlessly with the YAML serialization system
@@ -103,7 +90,10 @@ def auto_register_config(cls):
 
     # Auto-register parameters in _config_mapping
     if not "_config_mapping" in cls.__dict__:
-        cls._config_mapping = deepcopy(getattr(cls, "_config_mapping", {}))
+        cls._config_mapping = deepcopy(getattr(cls, "_config_mapping", bidict()))
+    if not isinstance(cls._config_mapping, bidict):
+        cls._config_mapping = bidict(cls._config_mapping)
+
     registered_parameters = tuple(cls._config_mapping.values())
 
     for param_name in list(sig.parameters.keys())[1:]:  # Skip 'self'
@@ -116,6 +106,7 @@ def auto_register_config(cls):
         ) and (param_name not in registered_parameters):
             cls._config_mapping[param_name] = param_name
 
+    @wraps(original_init)
     def __init__(self, *args, **kwargs):
         log.debug(f"set attributes for {self.__class__.__name__} in {cls.__name__}")
         # auto-register the attributes based on the signature
@@ -162,33 +153,10 @@ def auto_register_config(cls):
 
 class YAMLSerializationMixin:
     _config_key: Optional[str] = None
-    _config_mapping: Dict[str, str] = {}
+    _config_mapping: MutableBidict[str, str] = bidict()
     R"""
     `_config_mapping` is a dictionary mapping the attribute names of the class to the config option names. This is used to convert the class to a DictConfig.
 
-    For example, if an algorithm class is defined as follows:
-    
-    ```python
-    class SomeModelFusionAlgorithm(BaseModelFusionAlgorithm):
-        hyper_parameter_1 = None
-        hyper_parameter_2 = None
-
-        _config_mapping = BaseModelFusionAlgorithm._config_mapping | {
-            "hyper_parameter_1" : "hyper_param_1",
-            "hyper_parameter_2" : "hyper_param_2",
-        }
-        def __init__(self, hyper_param_1: int, hyper_param_2: int):
-            self.hyper_parameter_1 = hyper_param_1
-            self.hyper_parameter_2 = hyper_param_2
-            super().__init__()
-    ```
-
-    The model pool will be converted to a DictConfig as follows:
-        
-    ```python
-    algorithm = SomeModelFusionAlgorithm(hyper_param_1=1, hyper_param_2=2)
-    ```
-
     >>> algorithm.config
         DictCOnfig({'_target_': 'SomeModelFusionAlgorithm', 'hyper_param_1': 1, 'hyper_param_2': 2})
 
@@ -207,17 +175,6 @@ class YAMLSerializationMixin:
         This property converts the model pool instance into a dictionary
         configuration, which can be used for serialization or other purposes.
 
-        Example:
-
-        ```python
-        model = SomeModelFusionAlgorithm(hyper_param_1=1, hyper_param_2=2)
-        config = model.config
-        print(config)
-        # DictConfig({'_target_': 'SomeModelFusionAlgorithm', 'hyper_param_1': 1, 'hyper_param_2': 2})
-        ```
-
-        This is useful for serializing the object to a YAML file or for debugging.
-
         Returns:
             DictConfig: The configuration of the model pool.
         """
@@ -282,16 +239,6 @@ class YAMLSerializationMixin:
                 serialization. This is how the attribute will appear in YAML output.
             value: The value to assign to the attribute.
 
-        Example:
-            ```python
-            model = BaseYAMLSerializable()
-            model.set_option("learning_rate", "lr", 0.001)
-
-            # This sets model.learning_rate = 0.001
-            # and maps it to "lr" in the config output
-            config = model.config
-            # config will contain: {"lr": 0.001, ...}
-            ```
         """
         setattr(self, attr_name, value)
         self._config_mapping[attr_name] = param_name

fusion_bench/mixins/simple_profiler.py

@@ -9,27 +9,33 @@ __all__ = ["SimpleProfilerMixin"]
 
 class SimpleProfilerMixin:
     """
-    A mixin class that provides simple profiling capabilities.
+    A mixin class that provides simple profiling capabilities using Lightning's SimpleProfiler.
 
-    This mixin allows for easy profiling of code blocks using a context manager.
-    It also provides methods to start and stop profiling actions, and to print
-    a summary of the profiling results.
+    This mixin allows for easy profiling of code blocks using a context manager or manual
+    start/stop methods. It measures the execution time of named actions and provides
+    a summary of the profiling results. Unlike statistical profilers, this provides
+    precise timing measurements for specific code blocks.
 
-    Examples:
-
-    ```python
-    class MyClass(SimpleProfilerMixin):
-        def do_something(self):
-            with self.profile("work"):
-                # do some work here
-                ...
-            with self.profile("more work"):
-                # do more work here
-                ...
+    Note:
+        This mixin uses Lightning's SimpleProfiler which measures wall-clock time
+        for named actions. It's suitable for timing discrete operations rather than
+        detailed function-level profiling.
 
-            # print the profiling summary
-            self.print_profile_summary()
-    ```
+    Examples:
+        ```python
+        class MyClass(SimpleProfilerMixin):
+            def do_something(self):
+                with self.profile("data_loading"):
+                    # Load data here
+                    data = load_data()
+
+                with self.profile("model_training"):
+                    # Train model here
+                    model.train(data)
+
+                # Print the profiling summary
+                self.print_profile_summary("Training Profile")
+        ```
 
     Attributes:
         _profiler (SimpleProfiler): An instance of the SimpleProfiler class used for profiling.
@@ -38,7 +44,13 @@ class SimpleProfilerMixin:
     _profiler: SimpleProfiler = None
 
     @property
-    def profiler(self):
+    def profiler(self) -> SimpleProfiler:
+        """
+        Get the SimpleProfiler instance, creating it if necessary.
+
+        Returns:
+            SimpleProfiler: The profiler instance used for timing measurements.
+        """
         # Lazy initialization of the profiler instance
         if self._profiler is None:
             self._profiler = SimpleProfiler()
@@ -47,14 +59,24 @@ class SimpleProfilerMixin:
     @contextmanager
     def profile(self, action_name: str) -> Generator:
         """
-        Context manager for profiling a code block
+        Context manager for profiling a code block.
+
+        This context manager automatically starts profiling when entering the block
+        and stops profiling when exiting the block (even if an exception occurs).
+
+        Args:
+            action_name: A descriptive name for the action being profiled.
+                        This name will appear in the profiling summary.
+
+        Yields:
+            str: The action name that was provided.
 
         Example:
 
         ```python
-        with self.profile("work"):
-            # do some work here
-            ...
+        with self.profile("data_processing"):
+            # Process data here
+            result = process_large_dataset()
         ```
         """
         try:
@@ -64,18 +86,79 @@ class SimpleProfilerMixin:
             self.stop_profile(action_name)
 
     def start_profile(self, action_name: str):
+        """
+        Start profiling for a named action.
+
+        This method begins timing for the specified action. You must call
+        stop_profile() with the same action name to complete the measurement.
+
+        Args:
+            action_name: A descriptive name for the action being profiled.
+                        This name will appear in the profiling summary.
+
+        Example:
+            ```python
+            self.start_profile("model_inference")
+            result = model.predict(data)
+            self.stop_profile("model_inference")
+            ```
+        """
         self.profiler.start(action_name)
 
     def stop_profile(self, action_name: str):
+        """
+        Stop profiling for a named action.
+
+        This method ends timing for the specified action that was previously
+        started with start_profile().
+
+        Args:
+            action_name: The name of the action to stop profiling.
+                        Must match the name used in start_profile().
+
+        Example:
+            ```python
+            self.start_profile("data_loading")
+            data = load_data()
+            self.stop_profile("data_loading")
+            ```
+        """
         self.profiler.stop(action_name)
 
     @rank_zero_only
     def print_profile_summary(self, title: Optional[str] = None):
+        """
+        Print a summary of all profiled actions.
+
+        This method outputs a formatted summary showing the timing information
+        for all actions that have been profiled. The output includes action names
+        and their execution times.
+
+        Args:
+            title: Optional title to print before the profiling summary.
+                  If provided, this will be printed as a header.
+
+        Note:
+            This method is decorated with @rank_zero_only, meaning it will only
+            execute on the main process in distributed training scenarios.
+
+        Example:
+            ```python
+            # After profiling some actions
+            self.print_profile_summary("Training Performance Summary")
+            ```
+        """
         if title is not None:
             print(title)
         print(self.profiler.summary())
 
     def __del__(self):
+        """
+        Cleanup when the object is destroyed.
+
+        Ensures that the profiler instance is properly cleaned up to prevent
+        memory leaks when the mixin instance is garbage collected.
+        """
         if self._profiler is not None:
             del self._profiler
             self._profiler = None

fusion_bench/modelpool/__init__.py

@@ -18,6 +18,7 @@ _import_structure = {
         "GPT2ForSequenceClassificationPool",
     ],
     "seq_classification_lm": ["SequenceClassificationModelPool"],
+    "resnet_for_image_classification": ["ResNetForImageClassificationPool"],
 }
 
 
@@ -33,6 +34,7 @@ if TYPE_CHECKING:
     from .nyuv2_modelpool import NYUv2ModelPool
     from .openclip_vision import OpenCLIPVisionModelPool
     from .PeftModelForSeq2SeqLM import PeftModelForSeq2SeqLMPool
+    from .resnet_for_image_classification import ResNetForImageClassificationPool
     from .seq2seq_lm import Seq2SeqLMPool
     from .seq_classification_lm import SequenceClassificationModelPool
 

fusion_bench/modelpool/base_pool.py

@@ -180,26 +180,59 @@ class BaseModelPool(
 
         Args:
             model_name_or_config (Union[str, DictConfig]): The model name or configuration.
+                - If str: should be a key in self._models
+                - If DictConfig: should be a configuration dict for instantiation
+            *args: Additional positional arguments passed to model instantiation.
+            **kwargs: Additional keyword arguments passed to model instantiation.
 
         Returns:
-            nn.Module: The instantiated model.
+            nn.Module: The instantiated or retrieved model.
         """
         log.debug(f"Loading model: {model_name_or_config}", stacklevel=2)
-        if isinstance(self._models, DictConfig):
-            model_config = (
-                self._models[model_name_or_config]
-                if isinstance(model_name_or_config, str)
-                else model_name_or_config
-            )
-            model = instantiate(model_config, *args, **kwargs)
-        elif isinstance(self._models, Dict) and isinstance(model_name_or_config, str):
-            model = self._models[model_name_or_config]
+
+        if isinstance(model_name_or_config, str):
+            model_name = model_name_or_config
+            # Handle string model names - lookup in the model pool
+            if model_name not in self._models:
+                raise KeyError(
+                    f"Model '{model_name}' not found in model pool. "
+                    f"Available models: {list(self._models.keys())}"
+                )
+            model_config = self._models[model_name]
+
+            # Handle different types of model configurations
+            match model_config:
+                case dict() | DictConfig() as config:
+                    # Configuration that needs instantiation
+                    log.debug(f"Instantiating model '{model_name}' from configuration")
+                    return instantiate(config, *args, **kwargs)
+
+                case nn.Module() as model:
+                    # Pre-instantiated model - return directly
+                    log.debug(
+                        f"Returning pre-instantiated model '{model_name}' of type {type(model)}"
+                    )
+                    return model
+
+                case _:
+                    # Unsupported model configuration type
+                    raise ValueError(
+                        f"Unsupported model configuration type for '{model_name}': {type(model_config)}. "
+                        f"Expected nn.Module, dict, or DictConfig."
+                    )
+
+        elif isinstance(model_name_or_config, (dict, DictConfig)):
+            # Direct configuration - instantiate directly
+            log.debug("Instantiating model from direct DictConfig")
+            model_config = model_name_or_config
+            return instantiate(model_config, *args, **kwargs)
+
         else:
-            raise ValueError(
-                "The model pool configuration is not in the expected format."
-                f"We expected a DictConfig or Dict, but got {type(self._models)}."
+            # Unsupported input type
+            raise TypeError(
+                f"Unsupported input type: {type(model_name_or_config)}. "
+                f"Expected str or DictConfig."
             )
-        return model
 
     def load_pretrained_model(self, *args, **kwargs):
         assert (
@@ -229,6 +262,36 @@ class BaseModelPool(
         for model_name in self.model_names:
             yield model_name, self.load_model(model_name)
 
+    @property
+    def has_train_dataset(self) -> bool:
+        """
+        Check if the model pool contains training datasets.
+
+        Returns:
+            bool: True if training datasets are available, False otherwise.
+        """
+        return self._train_datasets is not None and len(self._train_datasets) > 0
+
+    @property
+    def has_val_dataset(self) -> bool:
+        """
+        Check if the model pool contains validation datasets.
+
+        Returns:
+            bool: True if validation datasets are available, False otherwise.
+        """
+        return self._val_datasets is not None and len(self._val_datasets) > 0
+
+    @property
+    def has_test_dataset(self) -> bool:
+        """
+        Check if the model pool contains testing datasets.
+
+        Returns:
+            bool: True if testing datasets are available, False otherwise.
+        """
+        return self._test_datasets is not None and len(self._test_datasets) > 0
+
     def load_train_dataset(self, dataset_name: str, *args, **kwargs) -> Dataset:
         """
         Load the training dataset for the specified model.

fusion_bench/modelpool/causal_lm/causal_lm.py

@@ -1,5 +1,5 @@
 """
-Online documentation for this module: https://tanganke.github.io/fusion_bench/modelpool/causal_lm
+Online documentation for this module: https://tanganke.github.io/fusion_bench/modelpool/llm
 """
 
 import logging
@@ -26,6 +26,7 @@ from fusion_bench import (
     instantiate,
     parse_dtype,
 )
+from fusion_bench.models.hf_utils import create_default_model_card
 from fusion_bench.utils.lazy_state_dict import LazyStateDict
 
 log = logging.getLogger(__name__)
@@ -271,13 +272,16 @@ class CausalLMPool(BaseModelPool):
         save_tokenizer: bool = False,
         tokenizer_kwargs=None,
         tokenizer: Optional[PreTrainedTokenizer] = None,
+        algorithm_config: Optional[DictConfig] = None,
+        description: Optional[str] = None,
+        base_model_in_modelcard: bool = True,
         **kwargs,
     ):
         """Save a model to the specified path with optional tokenizer and Hub upload.
 
         This method provides comprehensive model saving capabilities including
-        optional tokenizer saving, dtype conversion, and Hugging Face Hub upload.
-        The model is saved in the standard Hugging Face format.
+        optional tokenizer saving, dtype conversion, model card creation, and
+        Hugging Face Hub upload. The model is saved in the standard Hugging Face format.
 
         Args:
             model: The PreTrainedModel instance to be saved.
@@ -295,15 +299,13 @@ class CausalLMPool(BaseModelPool):
                 when save_tokenizer is True.
             tokenizer: Optional pre-loaded tokenizer instance. If provided, this
                 tokenizer will be saved regardless of the save_tokenizer flag.
+            algorithm_config: Optional DictConfig containing algorithm configuration.
+                If provided, a model card will be created with algorithm details.
+            description: Optional description for the model card. If not provided
+                and algorithm_config is given, a default description will be generated.
             **kwargs: Additional keyword arguments passed to the model's
                 save_pretrained method.
 
-        Side Effects:
-            - Creates model files in the specified directory
-            - Optionally creates tokenizer files in the same directory
-            - May convert the model to a different dtype
-            - May upload files to Hugging Face Hub
-
         Example:
             ```python
             >>> pool = CausalLMPool(models=..., tokenizer=...)
@@ -313,7 +315,9 @@ class CausalLMPool(BaseModelPool):
             ...     "/path/to/save",
             ...     save_tokenizer=True,
             ...     model_dtype="float16",
-            ...     push_to_hub=True
+            ...     push_to_hub=True,
+            ...     algorithm_config=algorithm_config,
+            ...     description="Custom merged model"
             ... )
             ```
         """
@@ -337,6 +341,24 @@ class CausalLMPool(BaseModelPool):
             **kwargs,
         )
 
+        # Create and save model card if algorithm_config is provided
+        if algorithm_config is not None:
+            if description is None:
+                description = "Model created using FusionBench."
+            model_card_str = create_default_model_card(
+                base_model=(
+                    self.get_model_path("_pretrained_")
+                    if base_model_in_modelcard and self.has_pretrained
+                    else None
+                ),
+                models=[self.get_model_path(m) for m in self.model_names],
+                description=description,
+                algorithm_config=algorithm_config,
+                modelpool_config=self.config,
+            )
+            with open(os.path.join(path, "README.md"), "w") as f:
+                f.write(model_card_str)
+
 
 class CausalLMBackbonePool(CausalLMPool):
     """A specialized model pool that loads only the transformer backbone layers.

fusion_bench/modelpool/clip_vision/modelpool.py

@@ -93,42 +93,79 @@ class CLIPVisionModelPool(BaseModelPool):
         self, model_name_or_config: Union[str, DictConfig], *args, **kwargs
     ) -> CLIPVisionModel:
         """
-        This method is used to load a CLIPVisionModel from the model pool.
+        Load a CLIPVisionModel from the model pool with support for various configuration formats.
 
-        Example configuration could be:
+        This method provides flexible model loading capabilities, handling different types of model
+        configurations including string paths, pre-instantiated models, and complex configurations.
 
+        Supported configuration formats:
+        1. String model paths (e.g., Hugging Face model IDs)
+        2. Pre-instantiated nn.Module objects
+        3. DictConfig objects for complex configurations
+
+        Example configuration:
         ```yaml
         models:
+            # Simple string paths to Hugging Face models
             cifar10: tanganke/clip-vit-base-patch32_cifar10
             sun397: tanganke/clip-vit-base-patch32_sun397
             stanford-cars: tanganke/clip-vit-base-patch32_stanford-cars
+
+            # Complex configuration with additional parameters
+            custom_model:
+                _target_: transformers.CLIPVisionModel.from_pretrained
+                pretrained_model_name_or_path: openai/clip-vit-base-patch32
+                torch_dtype: float16
         ```
 
         Args:
-            model_name_or_config (Union[str, DictConfig]): The name of the model or the model configuration.
+            model_name_or_config (Union[str, DictConfig]): Either a model name from the pool
+                or a configuration dictionary for instantiating the model.
+            *args: Additional positional arguments passed to model loading/instantiation.
+            **kwargs: Additional keyword arguments passed to model loading/instantiation.
 
         Returns:
-            CLIPVisionModel: The loaded CLIPVisionModel.
+            CLIPVisionModel: The loaded CLIPVisionModel instance.
         """
+        # Check if we have a string model name that exists in our model pool
         if (
             isinstance(model_name_or_config, str)
             and model_name_or_config in self._models
         ):
-            model = self._models[model_name_or_config]
-            if isinstance(model, str):
-                if rank_zero_only.rank == 0:
-                    log.info(f"Loading `transformers.CLIPVisionModel`: {model}")
-                repo_path = resolve_repo_path(
-                    model, repo_type="model", platform=self._platform
-                )
-                return CLIPVisionModel.from_pretrained(repo_path, *args, **kwargs)
-            if isinstance(model, nn.Module):
-                if rank_zero_only.rank == 0:
-                    log.info(f"Returning existing model: {model}")
-                return model
-        else:
-            # If the model is not a string, we use the default load_model method
-            return super().load_model(model_name_or_config, *args, **kwargs)
+            model_name = model_name_or_config
+
+            # handle different model configuration types
+            match self._models[model_name_or_config]:
+                case str() as model_path:
+                    # Handle string model paths (e.g., Hugging Face model IDs)
+                    if rank_zero_only.rank == 0:
+                        log.info(
+                            f"Loading model `{model_name}` of type `transformers.CLIPVisionModel` from {model_path}"
+                        )
+                    # Resolve the repository path (supports both HuggingFace and ModelScope)
+                    repo_path = resolve_repo_path(
+                        model_path, repo_type="model", platform=self._platform
+                    )
+                    # Load and return the CLIPVisionModel from the resolved path
+                    return CLIPVisionModel.from_pretrained(repo_path, *args, **kwargs)
+
+                case nn.Module() as model:
+                    # Handle pre-instantiated model objects
+                    if rank_zero_only.rank == 0:
+                        log.info(
+                            f"Returning existing model `{model_name}` of type {type(model)}"
+                        )
+                    return model
+
+                case _:
+                    # Handle other configuration types (e.g., DictConfig) via parent class
+                    # This fallback prevents returning None when the model config doesn't
+                    # match the expected string or nn.Module patterns
+                    return super().load_model(model_name_or_config, *args, **kwargs)
+
+        # If model_name_or_config is not a string in our pool, delegate to parent class
+        # This handles cases where model_name_or_config is a DictConfig directly
+        return super().load_model(model_name_or_config, *args, **kwargs)
 
     def load_train_dataset(self, dataset_name: str, *args, **kwargs):
         dataset_config = self._train_datasets[dataset_name]