fusion-bench 0.2.9__py3-none-any.whl

fusion_bench/modelpool/base_pool.py

@@ -0,0 +1,268 @@
+import logging
+from copy import deepcopy
+from typing import Dict, List, Optional, Union
+
+import torch
+from omegaconf import DictConfig
+from torch import nn
+from torch.utils.data import Dataset
+
+from fusion_bench.mixins import BaseYAMLSerializableModel
+from fusion_bench.utils import instantiate, timeit_context
+
+__all__ = ["BaseModelPool"]
+
+log = logging.getLogger(__name__)
+
+
+class BaseModelPool(BaseYAMLSerializableModel):
+    """
+    A class for managing and interacting with a pool of models along with their associated datasets or other specifications. For example, a model pool may contain multiple models, each with its own training, validation, and testing datasets. As for the specifications, a vision model pool may contain image preprocessor, and a language model pool may contain a tokenizer.
+
+    Attributes:
+        _models (DictConfig): Configuration for all models in the pool.
+        _train_datasets (Optional[DictConfig]): Configuration for training datasets.
+        _val_datasets (Optional[DictConfig]): Configuration for validation datasets.
+        _test_datasets (Optional[DictConfig]): Configuration for testing datasets.
+        _usage_ (Optional[str]): Optional usage information.
+        _version_ (Optional[str]): Optional version information.
+    """
+
+    _program = None
+    _models: Union[DictConfig, Dict[str, nn.Module]]
+    _config_mapping = BaseYAMLSerializableModel._config_mapping | {
+        "_models": "models",
+        "_train_datasets": "train_datasets",
+        "_val_datasets": "val_datasets",
+        "_test_datasets": "test_datasets",
+    }
+
+    def __init__(
+        self,
+        models: Union[DictConfig, Dict[str, nn.Module], List[nn.Module]],
+        *,
+        train_datasets: Optional[DictConfig] = None,
+        val_datasets: Optional[DictConfig] = None,
+        test_datasets: Optional[DictConfig] = None,
+        **kwargs,
+    ):
+        if isinstance(models, List):
+            models = {str(model_idx): model for model_idx, model in enumerate(models)}
+        self._models = models
+        self._train_datasets = train_datasets
+        self._val_datasets = val_datasets
+        self._test_datasets = test_datasets
+        super().__init__(**kwargs)
+
+    @property
+    def has_pretrained(self):
+        """
+        Check if the model pool contains a pretrained model.
+
+        Returns:
+            bool: True if a pretrained model is available, False otherwise.
+        """
+        return "_pretrained_" in self._models
+
+    @property
+    def all_model_names(self) -> List[str]:
+        """
+        Get the names of all models in the pool, including special models.
+
+        Returns:
+            List[str]: A list of all model names.
+        """
+        return [name for name in self._models]
+
+    @property
+    def model_names(self) -> List[str]:
+        """
+        Get the names of regular models, excluding special models.
+
+        Returns:
+            List[str]: A list of regular model names.
+        """
+        return [name for name in self._models if not self.is_special_model(name)]
+
+    @property
+    def train_dataset_names(self) -> List[str]:
+        """
+        Get the names of training datasets.
+
+        Returns:
+            List[str]: A list of training dataset names.
+        """
+        return (
+            list(self._train_datasets.keys())
+            if self._train_datasets is not None
+            else []
+        )
+
+    @property
+    def val_dataset_names(self) -> List[str]:
+        """
+        Get the names of validation datasets.
+
+        Returns:
+            List[str]: A list of validation dataset names.
+        """
+        return list(self._val_datasets.keys()) if self._val_datasets is not None else []
+
+    @property
+    def test_dataset_names(self) -> List[str]:
+        """
+        Get the names of testing datasets.
+
+        Returns:
+            List[str]: A list of testing dataset names.
+        """
+        return (
+            list(self._test_datasets.keys()) if self._test_datasets is not None else []
+        )
+
+    def __len__(self):
+        return len(self.model_names)
+
+    @staticmethod
+    def is_special_model(model_name: str):
+        """
+        Determine if a model is special based on its name.
+
+        Args:
+            model_name (str): The name of the model.
+
+        Returns:
+            bool: True if the model name indicates a special model, False otherwise.
+        """
+        return model_name.startswith("_") and model_name.endswith("_")
+
+    def get_model_config(self, model_name: str, return_copy: bool = True) -> DictConfig:
+        """
+        Get the configuration for the specified model.
+
+        Args:
+            model_name (str): The name of the model.
+
+        Returns:
+            DictConfig: The configuration for the specified model.
+        """
+        model_config = self._models[model_name]
+        if return_copy:
+            model_config = deepcopy(model_config)
+        return model_config
+
+    def load_model(
+        self, model_name_or_config: Union[str, DictConfig], *args, **kwargs
+    ) -> nn.Module:
+        """
+        Load a model from the pool based on the provided configuration.
+
+        Args:
+            model (Union[str, DictConfig]): The model name or configuration.
+
+        Returns:
+            nn.Module: The instantiated 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]
+        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)}."
+            )
+        return model
+
+    def load_pretrained_model(self, *args, **kwargs):
+        assert (
+            self.has_pretrained
+        ), "No pretrained model available. Check `_pretrained_` is in the `models` key."
+        model = self.load_model("_pretrained_", *args, **kwargs)
+        return model
+
+    def load_pretrained_or_first_model(self, *args, **kwargs):
+        """
+        Load the pretrained model if available, otherwise load the first available model.
+
+        Returns:
+            nn.Module: The loaded model.
+        """
+        if self.has_pretrained:
+            model = self.load_model("_pretrained_", *args, **kwargs)
+        else:
+            model = self.load_model(self.model_names[0], *args, **kwargs)
+        return model
+
+    def models(self):
+        for model_name in self.model_names:
+            yield self.load_model(model_name)
+
+    def named_models(self):
+        for model_name in self.model_names:
+            yield model_name, self.load_model(model_name)
+
+    def load_train_dataset(self, dataset_name: str, *args, **kwargs) -> Dataset:
+        """
+        Load the training dataset for the specified model.
+
+        Args:
+            dataset_name (str): The name of the model.
+
+        Returns:
+            Dataset: The instantiated training dataset.
+        """
+        return instantiate(self._train_datasets[dataset_name], *args, **kwargs)
+
+    def train_datasets(self):
+        for dataset_name in self.train_dataset_names:
+            yield self.load_train_dataset(dataset_name)
+
+    def load_val_dataset(self, dataset_name: str, *args, **kwargs) -> Dataset:
+        """
+        Load the validation dataset for the specified model.
+
+        Args:
+            dataset_name (str): The name of the model.
+
+        Returns:
+            Dataset: The instantiated validation dataset.
+        """
+        return instantiate(self._val_datasets[dataset_name], *args, **kwargs)
+
+    def val_datasets(self):
+        for dataset_name in self.val_dataset_names:
+            yield self.load_val_dataset(dataset_name)
+
+    def load_test_dataset(self, dataset_name: str, *args, **kwargs) -> Dataset:
+        """
+        Load the testing dataset for the specified model.
+
+        Args:
+            dataset_name (str): The name of the model.
+
+        Returns:
+            Dataset: The instantiated testing dataset.
+        """
+        return instantiate(self._test_datasets[dataset_name], *args, **kwargs)
+
+    def test_datasets(self):
+        for dataset_name in self.test_dataset_names:
+            yield self.load_test_dataset(dataset_name)
+
+    def save_model(self, model: nn.Module, path: str):
+        """
+        Save the state dictionary of the model to the specified path.
+
+        Args:
+            model (nn.Module): The model whose state dictionary is to be saved.
+            path (str): The path where the state dictionary will be saved.
+        """
+        with timeit_context(f"Saving the state dict of model to {path}"):
+            torch.save(model.state_dict(), path)

fusion_bench/modelpool/causal_lm/__init__.py

@@ -0,0 +1,2 @@
+# flake8: noqa F401
+from .causal_lm import CausalLMBackbonePool, CausalLMPool, load_peft_causal_lm

fusion_bench/modelpool/causal_lm/causal_lm.py

@@ -0,0 +1,139 @@
+import logging
+import os
+from copy import deepcopy
+from typing import Any, Optional, TypeAlias, Union, cast  # noqa: F401
+
+import peft
+from omegaconf import DictConfig, flag_override
+from torch import nn
+from torch.nn.modules import Module
+from transformers import (
+    LlamaForCausalLM,
+    MistralForCausalLM,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+)
+from typing_extensions import override
+
+from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.utils import instantiate
+from fusion_bench.utils.dtype import parse_dtype
+
+log = logging.getLogger(__name__)
+
+CausalLM: TypeAlias = Union[LlamaForCausalLM, MistralForCausalLM, Any]
+
+
+class CausalLMPool(BaseModelPool):
+    _config_mapping = BaseModelPool._config_mapping | {
+        "_tokenizer": "tokenizer",
+        "_model_kwargs": "model_kwargs",
+    }
+
+    def __init__(
+        self,
+        models,
+        *,
+        tokenizer: Optional[DictConfig],
+        model_kwargs: Optional[DictConfig] = None,
+        **kwargs,
+    ):
+        super().__init__(models, **kwargs)
+        # process `model_kwargs`
+        self._tokenizer = tokenizer
+        self._model_kwargs = model_kwargs
+        if self._model_kwargs is None:
+            self._model_kwargs = DictConfig({})
+        with flag_override(self._model_kwargs, "allow_objects", True):
+            if hasattr(self._model_kwargs, "torch_dtype"):
+                self._model_kwargs.torch_dtype = parse_dtype(
+                    self._model_kwargs.torch_dtype
+                )
+
+    @override
+    def load_model(
+        self,
+        model_name_or_config: str | DictConfig,
+        *args,
+        **kwargs,
+    ) -> LlamaForCausalLM | MistralForCausalLM | nn.Module:
+        model_kwargs = deepcopy(self._model_kwargs)
+        model_kwargs.update(kwargs)
+        if isinstance(model_name_or_config, str):
+            log.info(f"Loading model: {model_name_or_config}", stacklevel=2)
+        return super().load_model(model_name_or_config, *args, **model_kwargs)
+
+    def load_tokenizer(self, *args, **kwargs) -> PreTrainedTokenizer:
+        assert self._tokenizer is not None, "Tokenizer is not defined in the config"
+        log.info("Loading tokenizer.", stacklevel=2)
+        tokenizer = instantiate(self._tokenizer, *args, **kwargs)
+        return tokenizer
+
+    @override
+    def save_model(
+        self,
+        model: PreTrainedModel,
+        path: str,
+        push_to_hub: bool = False,
+        model_dtype: Optional[str] = None,
+        save_tokenizer: bool = False,
+        tokenizer_kwargs=None,
+        **kwargs,
+    ):
+        """
+        Save the model to the specified path.
+
+        Args:
+            model (PreTrainedModel): The model to be saved.
+            path (str): The path where the model will be saved.
+            push_to_hub (bool, optional): Whether to push the model to the Hugging Face Hub. Defaults to False.
+            save_tokenizer (bool, optional): Whether to save the tokenizer along with the model. Defaults to False.
+            **kwargs: Additional keyword arguments passed to the `save_pretrained` method.
+        """
+        path = os.path.expanduser(path)
+        if save_tokenizer:
+            if tokenizer_kwargs is None:
+                tokenizer_kwargs = {}
+            # load the tokenizer
+            tokenizer = self.load_tokenizer(**tokenizer_kwargs)
+            tokenizer.save_pretrained(
+                path,
+                push_to_hub=push_to_hub,
+            )
+        if model_dtype is not None:
+            model.to(dtype=parse_dtype(model_dtype))
+        model.save_pretrained(
+            path,
+            push_to_hub=push_to_hub,
+            **kwargs,
+        )
+
+
+class CausalLMBackbonePool(CausalLMPool):
+    def load_model(
+        self, model_name_or_config: str | DictConfig, *args, **kwargs
+    ) -> Module:
+        model: Union[MistralForCausalLM, LlamaForCausalLM, Any] = super().load_model(
+            model_name_or_config, *args, **kwargs
+        )
+        return model.model.layers
+
+
+def load_peft_causal_lm(
+    base_model_path: str,
+    peft_model_path: str,
+    torch_dtype: str = "bfloat16",
+    is_trainable: bool = True,
+    merge_and_unload: bool = False,
+):
+    base_model = LlamaForCausalLM.from_pretrained(
+        base_model_path, torch_dtype=torch_dtype
+    )
+    model = peft.PeftModel.from_pretrained(
+        base_model,
+        peft_model_path,
+        is_trainable=is_trainable,
+    )
+    if merge_and_unload:
+        model = model.merge_and_unload()
+    return model

fusion_bench/modelpool/clip_vision/__init__.py

@@ -0,0 +1 @@
+from .modelpool import CLIPVisionModelPool

fusion_bench/modelpool/clip_vision/modelpool.py

@@ -0,0 +1,145 @@
+import logging
+from copy import deepcopy
+from typing import Optional, Union
+
+from datasets import load_dataset
+from omegaconf import DictConfig, open_dict
+from torch import nn
+from torch.utils.data import Dataset
+from transformers import CLIPModel, CLIPProcessor, CLIPVisionModel
+from typing_extensions import override
+
+from fusion_bench.utils import instantiate, timeit_context
+
+from ..base_pool import BaseModelPool
+
+log = logging.getLogger(__name__)
+
+
+class CLIPVisionModelPool(BaseModelPool):
+    """
+    A model pool for managing Hugging Face's CLIP Vision models.
+
+    This class extends the base `ModelPool` class and overrides its methods to handle
+    the specifics of the CLIP Vision models provided by the Hugging Face Transformers library.
+    """
+
+    _config_mapping = BaseModelPool._config_mapping | {"_processor": "processor"}
+
+    def __init__(
+        self,
+        models: DictConfig,
+        *,
+        processor: Optional[DictConfig] = None,
+        **kwargs,
+    ):
+        super().__init__(models, **kwargs)
+
+        self._processor = processor
+
+    def load_processor(self, *args, **kwargs) -> CLIPProcessor:
+        assert self._processor is not None, "Processor is not defined in the config"
+        if isinstance(self._processor, str):
+            log.info(f"Loading `transformers.CLIPProcessor`: {self._processor}")
+            processor = CLIPProcessor.from_pretrained(self._processor)
+        else:
+            processor = instantiate(self._processor, *args, **kwargs)
+        return processor
+
+    def load_clip_model(self, model_name: str, *args, **kwargs) -> CLIPModel:
+        model_config = self._models[model_name]
+
+        if isinstance(model_config, str):
+            log.info(f"Loading `transformers.CLIPModel`: {model_config}")
+            clip_model = CLIPModel.from_pretrained(model_config, *args, **kwargs)
+            return clip_model
+        else:
+            assert isinstance(
+                model_config, DictConfig
+            ), "Model config must be a DictConfig"
+            model_config = deepcopy(model_config)
+            with open_dict(model_config):
+                model_config._target_ = "transformers.CLIPModel.from_pretrained"
+            clip_model = instantiate(model_config, *args, **kwargs)
+            return clip_model
+
+    @override
+    def save_model(self, model: CLIPVisionModel, path: str):
+        """
+        Save a CLIP Vision model to the given path.
+
+        Args:
+            model (CLIPVisionModel): The model to save.
+            path (str): The path to save the model to.
+        """
+        with timeit_context(f'Saving clip vision model to "{path}"'):
+            model.save_pretrained(path)
+
+    def load_model(
+        self, model_name_or_config: Union[str, DictConfig], *args, **kwargs
+    ) -> CLIPVisionModel:
+        """
+        This method is used to load a CLIPVisionModel from the model pool.
+
+        Example configuration could be:
+
+        ```yaml
+        models:
+            cifar10: tanganke/clip-vit-base-patch32_cifar10
+            sun397: tanganke/clip-vit-base-patch32_sun397
+            stanford-cars: tanganke/clip-vit-base-patch32_stanford-cars
+        ```
+
+        Args:
+            model_name_or_config (Union[str, DictConfig]): The name of the model or the model configuration.
+
+        Returns:
+            CLIPVisionModel: The loaded CLIPVisionModel.
+        """
+        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):
+                log.info(f"Loading `transformers.CLIPVisionModel`: {model}")
+                return CLIPVisionModel.from_pretrained(model, *args, **kwargs)
+            if isinstance(model, nn.Module):
+                log.info(f"Returning existing model: {model}")
+                return model
+
+        # If the model is not a string, we use the default load_model method
+        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]
+        if isinstance(dataset_config, str):
+            log.info(
+                f"Loading train dataset using `datasets.load_dataset`: {dataset_config}"
+            )
+            dataset = load_dataset(dataset_config, split="train")
+        else:
+            dataset = super().load_train_dataset(dataset_name, *args, **kwargs)
+        return dataset
+
+    def load_val_dataset(self, dataset_name: str, *args, **kwargs):
+        dataset_config = self._val_datasets[dataset_name]
+        if isinstance(dataset_config, str):
+            log.info(
+                f"Loading validation dataset using `datasets.load_dataset`: {dataset_config}"
+            )
+            dataset = load_dataset(dataset_config, split="validation")
+        else:
+            dataset = super().load_val_dataset(dataset_name, *args, **kwargs)
+        return dataset
+
+    def load_test_dataset(self, dataset_name: str, *args, **kwargs):
+        dataset_config = self._test_datasets[dataset_name]
+        if isinstance(dataset_config, str):
+            log.info(
+                f"Loading test dataset using `datasets.load_dataset`: {dataset_config}"
+            )
+            dataset = load_dataset(dataset_config, split="test")
+        else:
+            dataset = super().load_test_dataset(dataset_name, *args, **kwargs)
+        return dataset

fusion_bench/modelpool/huggingface_automodel.py

@@ -0,0 +1,20 @@
+import logging
+
+from omegaconf import DictConfig
+from torch.nn.modules import Module
+from transformers import AutoModel
+
+from fusion_bench.compat.modelpool import ModelPool
+
+log = logging.getLogger(__name__)
+
+
+class AutoModelPool(ModelPool):
+    def load_model(self, model_config: str | DictConfig) -> Module:
+        if isinstance(model_config, str):
+            model_config = self.get_model_config(model_config)
+        else:
+            model_config = model_config
+
+        model = AutoModel.from_pretrained(model_config.path)
+        return model

fusion_bench/modelpool/huggingface_gpt2_classification.py

@@ -0,0 +1,63 @@
+import functools
+import logging
+from typing import Optional
+
+from omegaconf import DictConfig
+from transformers import GPT2ForSequenceClassification, GPT2Tokenizer
+
+from fusion_bench.dataset.gpt2_glue import TokenizedGLUE
+from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.utils import instantiate
+
+log = logging.getLogger(__name__)
+tokenizer: GPT2Tokenizer = None
+
+
+@functools.cache
+def load_gpt2_dataset(name: str, split: Optional[str] = None):
+    global tokenizer
+    dataset = TokenizedGLUE(tokenizer=tokenizer).load_dataset(name)
+    if split is not None:
+        dataset = dataset[split]
+    return dataset
+
+
+def load_gpt2_tokenizer(pretrained_model_name_or_path: str):
+    tokenizer = GPT2Tokenizer.from_pretrained(pretrained_model_name_or_path)
+    tokenizer.model_max_length = 512
+    if tokenizer.pad_token is None:
+        if tokenizer.unk_token is not None:
+            tokenizer.pad_token = tokenizer.unk_token
+        elif tokenizer.eos_token is not None:
+            tokenizer.pad_token = tokenizer.eos_token
+        else:
+            raise ValueError
+    return tokenizer
+
+
+class GPT2ForSequenceClassificationPool(BaseModelPool):
+    _config_mapping = BaseModelPool._config_mapping | {"_tokenizer": "tokenizer"}
+
+    def __init__(self, tokenizer: DictConfig, **kwargs):
+        self._tokenizer = tokenizer
+        super().__init__(**kwargs)
+        self.setup()
+
+    def setup(self):
+        global tokenizer
+        self.tokenizer = tokenizer = instantiate(self._tokenizer)
+
+    def load_classifier(
+        self, model_config: str | DictConfig
+    ) -> GPT2ForSequenceClassification:
+        if isinstance(model_config, str):
+            model_config = self.get_model_config(model_config, return_copy=True)
+        model_config._target_ = (
+            "transformers.GPT2ForSequenceClassification.from_pretrained"
+        )
+        model = instantiate(model_config)
+        return model
+
+
+# For compatibility
+HuggingFaceGPT2ClassificationPool = GPT2ForSequenceClassificationPool

fusion_bench/modelpool/nyuv2_modelpool.py

@@ -0,0 +1,40 @@
+import logging
+
+import torch
+from omegaconf import DictConfig
+from torch import nn
+
+from fusion_bench.compat.modelpool.base_pool import ModelPool
+from fusion_bench.dataset.nyuv2 import NYUv2
+from fusion_bench.models.nyuv2.aspp import DeepLabHead
+from fusion_bench.models.nyuv2.lightning_module import NYUv2Model
+from fusion_bench.models.nyuv2.resnet_dilated import ResnetDilated, resnet_dilated
+
+log = logging.getLogger(__name__)
+
+
+class NYUv2ModelPool(ModelPool):
+    def load_model(
+        self, model_config: str | DictConfig, encoder_only: bool = True
+    ) -> ResnetDilated | NYUv2Model:
+        if isinstance(model_config, str):
+            model_config = self.get_model_config(model_config)
+
+        encoder = resnet_dilated(model_config.encoder)
+        decoders = nn.ModuleDict(
+            {
+                task: DeepLabHead(2048, NYUv2.num_out_channels[task])
+                for task in model_config.decoders
+            }
+        )
+        model = NYUv2Model(encoder=encoder, decoders=decoders)
+        if model_config.get("ckpt_path", None) is not None:
+            ckpt = torch.load(model_config.ckpt_path, map_location="cpu")
+            if "state_dict" in ckpt:
+                ckpt = ckpt["state_dict"]
+            model.load_state_dict(ckpt, strict=False)
+
+        if encoder_only:
+            return model.encoder
+        else:
+            return model

fusion_bench/modelpool/seq2seq_lm/__init__.py

@@ -0,0 +1,2 @@
+# flake8: noqa F401
+from .modelpool import Seq2SeqLMPool