fusion-bench 0.2.24__py3-none-any.whl → 0.2.26__py3-none-any.whl

fusion_bench/mixins/__init__.py

@@ -11,6 +11,7 @@ _import_structure = {
     "hydra_config": ["HydraConfigMixin"],
     "lightning_fabric": ["LightningFabricMixin"],
     "openclip_classification": ["OpenCLIPClassificationMixin"],
+    "pyinstrument": ["PyinstrumentProfilerMixin"],
     "serialization": [
         "BaseYAMLSerializable",
         "YAMLSerializationMixin",
@@ -25,6 +26,7 @@ if TYPE_CHECKING:
     from .hydra_config import HydraConfigMixin
     from .lightning_fabric import LightningFabricMixin
     from .openclip_classification import OpenCLIPClassificationMixin
+    from .pyinstrument import PyinstrumentProfilerMixin
     from .serialization import (
         BaseYAMLSerializable,
         YAMLSerializationMixin,

fusion_bench/mixins/pyinstrument.py

@@ -0,0 +1,174 @@
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Generator, Optional, Union
+
+from lightning.fabric.utilities.rank_zero import rank_zero_only
+from pyinstrument import Profiler
+
+__all__ = ["PyinstrumentProfilerMixin"]
+
+
+class PyinstrumentProfilerMixin:
+    """
+    A mixin class that provides statistical profiling capabilities using pyinstrument.
+
+    This mixin allows for easy profiling of code blocks using a context manager.
+    It provides methods to start and stop profiling actions, save profiling results
+    to files, and print profiling summaries.
+
+    Note:
+        This mixin requires the `pyinstrument` package to be installed.
+        If not available, an ImportError will be raised when importing this module.
+
+    Examples:
+
+    ```python
+    class MyClass(PyinstrumentProfilerMixin):
+        def do_something(self):
+            with self.profile("work"):
+                # do some work here
+                ...
+
+            # save the profiling results
+            self.save_profile_report("profile_report.html")
+
+            # or print the summary
+            self.print_profile_summary()
+    ```
+
+    Attributes:
+        _profiler (Profiler): An instance of the pyinstrument Profiler class.
+    """
+
+    _profiler: Optional[Profiler] = None
+    _is_profiling: bool = False
+
+    @property
+    def profiler(self) -> Optional[Profiler]:
+        """Get the profiler instance, creating it if necessary."""
+        if self._profiler is None:
+            self._profiler = Profiler()
+        return self._profiler
+
+    @contextmanager
+    def profile(self, action_name: Optional[str] = None) -> Generator:
+        """
+        Context manager for profiling a code block.
+
+        Args:
+            action_name: Optional name for the profiling action (for logging purposes).
+
+        Example:
+
+        ```python
+        with self.profile("expensive_operation"):
+            # do some expensive work here
+            expensive_function()
+        ```
+        """
+        try:
+            self.start_profile(action_name)
+            yield action_name
+        finally:
+            self.stop_profile(action_name)
+
+    def start_profile(self, action_name: Optional[str] = None):
+        """
+        Start profiling.
+
+        Args:
+            action_name: Optional name for the profiling action.
+        """
+        if self._is_profiling:
+            return
+
+        self.profiler.start()
+        self._is_profiling = True
+        if action_name:
+            print(f"Started profiling: {action_name}")
+
+    def stop_profile(self, action_name: Optional[str] = None):
+        """
+        Stop profiling.
+
+        Args:
+            action_name: Optional name for the profiling action.
+        """
+        if not self._is_profiling:
+            return
+
+        self.profiler.stop()
+        self._is_profiling = False
+        if action_name:
+            print(f"Stopped profiling: {action_name}")
+
+    @rank_zero_only
+    def print_profile_summary(
+        self, title: Optional[str] = None, unicode: bool = True, color: bool = True
+    ):
+        """
+        Print a summary of the profiling results.
+
+        Args:
+            title: Optional title to print before the summary.
+            unicode: Whether to use unicode characters in the output.
+            color: Whether to use color in the output.
+        """
+        if self.profiler is None:
+            print("No profiling data available.")
+            return
+
+        if title is not None:
+            print(title)
+
+        print(self.profiler.output_text(unicode=unicode, color=color))
+
+    @rank_zero_only
+    def save_profile_report(
+        self,
+        output_path: Union[str, Path] = "profile_report.html",
+        format: str = "html",
+        title: Optional[str] = None,
+    ):
+        """
+        Save the profiling results to a file.
+
+        Args:
+            output_path: Path where to save the profiling report.
+            format: Output format ('html', or 'text').
+            title: Optional title for the report.
+        """
+        if self.profiler is None:
+            print("No profiling data available.")
+            return
+
+        output_path = Path(output_path)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        if format.lower() == "html":
+            content = self.profiler.output_html()
+        elif format.lower() == "text":
+            content = self.profiler.output_text(unicode=True, color=False)
+        else:
+            raise ValueError(f"Unsupported format: {format}. Use 'html', or 'text'.")
+
+        with open(output_path, "w", encoding="utf-8") as f:
+            f.write(content)
+
+        print(f"Profile report saved to: {output_path}")
+
+    def reset_profile(self):
+        """Reset the profiler to start fresh."""
+        if self._is_profiling:
+            self.stop_profile()
+
+        self._profiler = None
+
+    def __del__(self):
+        """Cleanup when the object is destroyed."""
+        if self._is_profiling:
+            self.stop_profile()
+
+        if self._profiler is not None:
+            del self._profiler
+            self._profiler = None

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/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]