fusion-bench 0.2.9__py3-none-any.whl

fusion_bench/compat/modelpool/huggingface_clip_vision.py

@@ -0,0 +1,178 @@
+import functools
+import logging
+from typing import Optional
+
+from omegaconf import DictConfig, open_dict
+from transformers import CLIPProcessor, CLIPVisionModel
+from typing_extensions import override
+
+from fusion_bench.dataset import CLIPDataset, load_dataset_from_config
+from fusion_bench.utils import timeit_context
+
+from .base_pool import ModelPool
+
+log = logging.getLogger(__name__)
+
+
+class HuggingFaceClipVisionPool(ModelPool):
+    """
+    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.
+    """
+
+    def __init__(self, modelpool_config: DictConfig):
+        super().__init__(modelpool_config)
+
+        self._clip_processor = None
+
+    @property
+    def clip_processor(self):
+        """
+        Returns the CLIP processor. If it's not already initialized, it initializes it using the path of the pretrained model.
+        """
+        if self._clip_processor is None:
+            if "_pretrained_" in self._model_names:
+                self._clip_processor = CLIPProcessor.from_pretrained(
+                    self.get_model_config("_pretrained_")["path"]
+                )
+            else:
+                log.warning(
+                    "No pretrained model found in the model pool. Returning the first model."
+                )
+                self._clip_processor = CLIPProcessor.from_pretrained(
+                    self.get_model_config(self.model_names[0])["path"]
+                )
+        return self._clip_processor
+
+    @override
+    def load_model(self, model_config: str | DictConfig) -> CLIPVisionModel:
+        """
+        Load a CLIP Vision model from the given configuration.
+
+        Args:
+            model_config (str | DictConfig): The configuration for the model to load.
+
+        Returns:
+            CLIPVisionModel: The loaded CLIP Vision model.
+        """
+        if isinstance(model_config, str):
+            model_config = self.get_model_config(model_config)
+
+        with timeit_context(
+            f"Loading CLIP vision model: '{model_config.name}' from '{model_config.path}'."
+        ):
+            vision_model = CLIPVisionModel.from_pretrained(model_config.path)
+        return vision_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 get_tta_dataset_config(self, dataset: str):
+        """
+        Retrieve the configuration for a TTA (Test-Time Adaptation) dataset.
+
+        Args:
+            dataset (str): The name of the dataset for which to retrieve the configuration.
+
+        Returns:
+            DictConfig: The configuration dictionary for the specified dataset.
+
+        Raises:
+            ValueError: If the specified dataset is not found in the configuration.
+        """
+        for dataset_config in self.config.tta_datasets:
+            if dataset_config.name == dataset:
+                return dataset_config
+        raise ValueError(f"Dataset {dataset} not found in config")
+
+    def prepare_dataset_config(self, dataset_config: DictConfig):
+        """
+        Prepare the dataset configuration by setting the dataset type if it's not already set.
+
+        Args:
+            dataset_config (DictConfig): The configuration dictionary for the dataset.
+
+        Returns:
+            DictConfig: The updated configuration dictionary for the dataset.
+        """
+        if not hasattr(dataset_config, "type"):
+            with open_dict(dataset_config):
+                dataset_config["type"] = self.config.dataset_type
+        return dataset_config
+
+    @functools.cache
+    def get_tta_test_dataset(
+        self, tta_dataset: str, clip_processor: Optional[CLIPProcessor] = None
+    ):
+        """
+        Load the test dataset for the task.
+        This method is cached, so the dataset is loaded only once.
+
+        Args:
+            tta_dataset (str): The name of the TTA dataset to load.
+            clip_processor (Optional[CLIPProcessor]): The CLIP processor to use for preprocessing the dataset. If None, the default processor is used.
+
+        Returns:
+            CLIPDataset: The loaded and preprocessed TTA test dataset.
+        """
+        if clip_processor is None:
+            # if clip_processor is not provided, try to load the clip_processor from pre-trained model
+            clip_processor = self.clip_processor
+        dataset_config = self.get_tta_dataset_config(tta_dataset)["dataset"]
+        dataset_config = self.prepare_dataset_config(dataset_config)
+        with timeit_context(f"Loading test dataset: {dataset_config.name}"):
+            dataset = load_dataset_from_config(dataset_config)
+        dataset = CLIPDataset(dataset, self.clip_processor)
+        return dataset
+
+    def get_train_dataset_config(self, model_name: str):
+        """
+        Retrieve the configuration for a specific training dataset.
+
+        Args:
+            model_name (str): The name of the model for which to retrieve the training dataset configuration.
+
+        Returns:
+            DictConfig: The configuration dictionary for the specified training dataset.
+
+        Raises:
+            ValueError: If the specified training dataset is not found in the configuration.
+        """
+        for dataset_config in self.config.train_datasets:
+            if dataset_config.name == model_name:
+                return dataset_config
+        raise ValueError(f"Dataset {model_name} not found in config")
+
+    def get_train_dataset(
+        self, model_name: str, clip_processor: Optional[CLIPProcessor] = None
+    ):
+        """
+        Load the training dataset for the specified model.
+
+        Args:
+            model_name (str): The name of the model for which to load the training dataset.
+            clip_processor (Optional[CLIPProcessor]): The CLIP processor to use for preprocessing the dataset. If None, the default processor is used.
+
+        Returns:
+            CLIPDataset: The loaded and preprocessed training dataset.
+        """
+        if clip_processor is None:
+            # if clip_processor is not provided, try to load the clip_processor from pre-trained model
+            clip_processor = self.clip_processor
+        dataset_config = self.get_train_dataset_config(model_name)["dataset"]
+        dataset_config = self.prepare_dataset_config(dataset_config)
+        with timeit_context(f"Loading train dataset: {dataset_config.name}"):
+            dataset = load_dataset_from_config(dataset_config)
+        dataset = CLIPDataset(dataset, self.clip_processor)
+        return dataset

fusion_bench/compat/taskpool/__init__.py

@@ -0,0 +1,95 @@
+# flake8: noqa F401
+from omegaconf import DictConfig
+
+from fusion_bench.taskpool.dummy import DummyTaskPool
+
+from .base_pool import TaskPool
+
+
+class TaskPoolFactory:
+    """
+    Factory class to create and manage different task pools.
+    This is for v0.1.x versions, deprecated.
+    For implementing new task pool, use `fusion_bench.taskpool.BaseTaskPool` instead.
+
+    This class provides methods to create task pools based on a given configuration,
+    register new task pools, and list available task pools.
+    """
+
+    _taskpool_types = {
+        "dummy": DummyTaskPool,
+        "clip_vit_classification": ".clip_image_classification.CLIPImageClassificationTaskPool",
+        "FlanT5GLUETextGenerationTaskPool": ".flan_t5_glue_text_generation.FlanT5GLUETextGenerationTaskPool",
+        "NYUv2TaskPool": "fusion_bench.taskpool.nyuv2_taskpool.NYUv2TaskPool",
+    }
+
+    @staticmethod
+    def create_taskpool(taskpool_config: DictConfig):
+        """
+        Create an instance of a task pool based on the provided configuration.
+
+        Args:
+            taskpool_config (DictConfig): The configuration for the task pool. Must contain a 'type' attribute that specifies the type of the task pool.
+
+        Returns:
+            TaskPool: An instance of the specified task pool.
+
+        Raises:
+            ValueError: If 'type' attribute is not found in the configuration or does not match any known task pool types.
+        """
+        from fusion_bench.utils import import_object
+
+        taskpool_type = taskpool_config.get("type")
+        if taskpool_type is None:
+            raise ValueError("Task pool type not specified")
+
+        if taskpool_type not in TaskPoolFactory._taskpool_types:
+            raise ValueError(
+                f"Unknown task pool: {taskpool_type}, available task pools: {TaskPoolFactory._taskpool_types.keys()}. You can register a new task pool using `TaskPoolFactory.register_taskpool()` method."
+            )
+        taskpool_cls = TaskPoolFactory._taskpool_types[taskpool_type]
+        if isinstance(taskpool_cls, str):
+            if taskpool_cls.startswith("."):
+                taskpool_cls = f"fusion_bench.compat.taskpool.{taskpool_cls[1:]}"
+            taskpool_cls = import_object(taskpool_cls)
+        return taskpool_cls(taskpool_config)
+
+    @staticmethod
+    def register_taskpool(name: str, taskpool_cls):
+        """
+        Register a new task pool with the factory.
+
+        Args:
+            name (str): The name of the task pool.
+            taskpool_cls: The class of the task pool to register.
+        """
+        TaskPoolFactory._taskpool_types[name] = taskpool_cls
+
+    @classmethod
+    def available_taskpools(cls):
+        """
+        Get a list of available task pools.
+
+        Returns:
+            list: A list of available task pool names.
+        """
+        return list(cls._taskpool_types.keys())
+
+
+def load_taskpool_from_config(taskpool_config: DictConfig):
+    """
+    Loads a task pool based on the provided configuration.
+
+    The function checks the 'type' attribute of the configuration and returns an instance of the corresponding task pool.
+    If the 'type' attribute is not found or does not match any known task pool types, a ValueError is raised.
+
+    Args:
+        taskpool_config (DictConfig): The configuration for the task pool. Must contain a 'type' attribute that specifies the type of the task pool.
+
+    Returns:
+        An instance of the specified task pool.
+
+    Raises:
+        ValueError: If 'type' attribute is not found in the configuration or does not match any known task pool types.
+    """
+    return TaskPoolFactory.create_taskpool(taskpool_config)

fusion_bench/compat/taskpool/base_pool.py

@@ -0,0 +1,111 @@
+from typing import Union
+
+from omegaconf import DictConfig
+from tqdm.autonotebook import tqdm
+
+
+class TaskPool:
+    """
+    A class to manage a pool of tasks for evaluation.
+    This is the base class for version 0.1.x, deprecated.
+    Use `fusion_bench.taskpool.BaseTaskPool` instead.
+
+    Attributes:
+        config (DictConfig): The configuration for the task pool.
+        _all_task_names (List[str]): A list of all task names in the task pool.
+    """
+
+    _program = None
+
+    def __init__(self, taskpool_config: DictConfig):
+        """
+        Initialize the TaskPool with the given configuration.
+
+        Args:
+            taskpool_config (DictConfig): The configuration for the task pool.
+        """
+        super().__init__()
+        self.config = taskpool_config
+
+        # Check for duplicate task names
+        if self.config.get("tasks", None) is not None:
+            task_names = [task["name"] for task in self.config["tasks"]]
+            assert len(task_names) == len(
+                set(task_names)
+            ), "Duplicate task names found in the task pool"
+            self._all_task_names = task_names
+
+    def evaluate(self, model):
+        """
+        Evaluate the model on all tasks in the task pool, and return a report.
+
+        Take image classification as an example, the report will look like:
+
+        ```python
+        {
+            "mnist": {
+                "accuracy": 0.8,
+                "loss": 0.2,
+            },
+            <task_name>: {
+                <metric_name>: <metric_value>,
+                ...
+            },
+        }
+        ```
+
+        Args:
+            model: The model to evaluate.
+
+        Returns:
+            report (dict): A dictionary containing the results of the evaluation for each task.
+        """
+        report = {}
+        for task_name in tqdm(self.task_names, desc="Evaluating tasks"):
+            task = self.load_task(task_name)
+            result = task.evaluate(model)
+            report[task_name] = result
+        return report
+
+    @property
+    def task_names(self):
+        """
+        Return a list of all task names in the task pool.
+
+        Returns:
+            List[str]: A list of all task names.
+        """
+        return self._all_task_names
+
+    def get_task_config(self, task_name: str):
+        """
+        Retrieve the configuration for a specific task from the task pool.
+
+        Args:
+            task_name (str): The name of the task for which to retrieve the configuration.
+
+        Returns:
+            DictConfig: The configuration dictionary for the specified task.
+
+        Raises:
+            ValueError: If the specified task is not found in the task pool.
+        """
+        for task in self.config["tasks"]:
+            if task["name"] == task_name:
+                return task
+        raise ValueError(f"Task {task_name} not found in the task pool")
+
+    def load_task(self, task_name_or_config: Union[str, DictConfig]):
+        """
+        Load a task from the task pool.
+
+        Args:
+            task_name_or_config (Union[str, DictConfig]): The name or configuration of the task to load.
+
+        Returns:
+            Any: The loaded task.
+
+        Raises:
+            NotImplementedError: If the method is not implemented in the subclass.
+        """
+        raise NotImplementedError

fusion_bench/compat/taskpool/clip_image_classification.py

@@ -0,0 +1,210 @@
+import functools
+import json
+import logging
+import os
+from copy import deepcopy
+from functools import cached_property
+from typing import Callable, List, cast
+
+import lightning as L
+import torch
+from omegaconf import DictConfig, open_dict
+from torch.utils.data import DataLoader
+from tqdm.autonotebook import tqdm
+from transformers import CLIPModel, CLIPProcessor, CLIPVisionModel
+
+from fusion_bench.dataset import CLIPDataset, load_dataset_from_config
+from fusion_bench.models.hf_clip import HFCLIPClassifier
+from fusion_bench.tasks.classification import ClassificationTask
+from fusion_bench.tasks.clip_classification import get_classnames_and_templates
+from fusion_bench.utils.parameters import count_parameters
+
+from .base_pool import TaskPool
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+log = logging.getLogger(__name__)
+
+
+@functools.cache
+def load_dataset_from_config_cached(dataset_config: DictConfig):
+    return load_dataset_from_config(dataset_config)
+
+
+class CLIPImageClassificationTask(ClassificationTask):
+    """
+    This class is used to define the image classification task for CLIP models.
+    """
+
+    _fabric: L.Fabric = None
+    _clip_processor: CLIPProcessor = None
+    #
+    _taskpool: "CLIPImageClassificationTaskPool" = None
+
+    classnames: List[str] = []
+    templates: List[Callable[[str], str]] = []
+
+    def __init__(self, task_config: DictConfig):
+        self.config = task_config
+
+        self.classnames, self.templates = get_classnames_and_templates(
+            self.config["dataset"].name
+        )
+
+    @cached_property
+    def test_dataset(self):
+        """
+        Load the test dataset for the task.
+        This method is cached, so the dataset is loaded only once.
+        """
+        dataset_config = self.config["dataset"]
+        dataset_config = self._taskpool.prepare_dataset_config(dataset_config)
+        log.info(f"Loading test dataset: {dataset_config.name}")
+        dataset = load_dataset_from_config_cached(dataset_config)
+        dataset = CLIPDataset(dataset, self._clip_processor)
+        return dataset
+
+    @property
+    def num_classes(self):
+        return len(self.classnames)
+
+    @property
+    def test_loader(self):
+        loader = DataLoader(
+            self.test_dataset,
+            batch_size=self.config["batch_size"],
+            num_workers=self.config["num_workers"],
+            shuffle=False,
+        )
+        if self._fabric is not None:
+            loader = self._fabric.setup_dataloaders(loader)
+        return loader
+
+    def evaluate(self, clip_model: CLIPModel):
+        """
+        Evaluate the model on the image classification task.
+
+        Args:
+            clip_model (CLIPModel): The CLIP model to evaluate.
+
+        Returns:
+            dict: A dictionary containing the evaluation results.
+        """
+        classifier = HFCLIPClassifier(
+            clip_model=clip_model, processor=self._clip_processor
+        )
+        classifier.set_classification_task(self.classnames, self.templates)
+        if self._fabric is not None:
+            classifier = self._fabric.setup_module(deepcopy(classifier))
+        results = super().evaluate(classifier)
+        log.info(f"Results for task {self.config.name}: {results}")
+        return results
+
+
+class CLIPImageClassificationTaskPool(TaskPool):
+    _fabric: L.Fabric = None
+
+    # CLIP forward model and processor
+    _clip_model: CLIPModel = None
+    _clip_processor: CLIPProcessor = None
+
+    def __init__(self, taskpool_config: DictConfig):
+        super().__init__(taskpool_config)
+
+    def prepare_dataset_config(self, dataset_config: DictConfig):
+        if not hasattr(dataset_config, "type"):
+            with open_dict(dataset_config):
+                dataset_config["type"] = self.config.dataset_type
+        return dataset_config
+
+    def prepare_task_config(self, task_config: DictConfig):
+        # set default values for keys that are not present in per task configuration
+        for key in ["num_workers", "batch_size", "fast_dev_run"]:
+            if not hasattr(task_config, key):
+                with open_dict(task_config):
+                    task_config[key] = self.config[key]
+        return task_config
+
+    @property
+    def clip_model(self):
+        if self._clip_model is None:
+            self._clip_model = CLIPModel.from_pretrained(self.config["clip_model"])
+        return self._clip_model
+
+    @property
+    def clip_processor(self):
+        if self._clip_processor is None:
+            self._clip_processor = CLIPProcessor.from_pretrained(
+                self.config["clip_model"]
+            )
+        return self._clip_processor
+
+    def load_task(self, task_name_or_config: str | DictConfig):
+        if isinstance(task_name_or_config, str):
+            task_config = self.get_task_config(task_name_or_config)
+        else:
+            task_config = task_name_or_config
+        task_config = self.prepare_task_config(task_config)
+
+        # load the task from the configuration
+        task = CLIPImageClassificationTask(task_config)
+        task._fabric = self._fabric
+        task._taskpool = self
+        task._clip_processor = self.clip_processor
+
+        return task
+
+    def evaluate(self, model: CLIPVisionModel, name=None):
+        """
+        Evaluate the model on the image classification task.
+
+        Args:
+            model (CLIPVisionModel): The vision model to evaluate.
+
+        Returns:
+            dict: A dictionary containing the evaluation results for each task.
+        """
+        # if the fabric is not set, and we have a GPU, create a fabric instance
+        if self._fabric is None and torch.cuda.is_available():
+            self._fabric = L.Fabric(devices=1)
+            self._fabric.launch()
+
+        # CLIPVisionModel works the same with CLIPVisonTransformer, so we can use it directly
+        self.clip_model.vision_model = model
+        report = {}
+        training_params, all_params = count_parameters(model)
+        report["model_info"] = {
+            "trainable_params": training_params,
+            "all_params": all_params,
+            "trainable_percentage": training_params / all_params,
+        }
+        if name is not None:
+            report["model_info"]["name"] = name
+        for task_name in tqdm(self.task_names, desc="Evaluating tasks"):
+            task = self.load_task(task_name)
+            result = task.evaluate(self.clip_model)
+            report[task_name] = result
+
+        # calculate the average accuracy and loss
+        if "average" not in report:
+            report["average"] = {}
+            accuracies = [
+                value["accuracy"]
+                for key, value in report.items()
+                if "accuracy" in value
+            ]
+            if len(accuracies) > 0:
+                average_accuracy = sum(accuracies) / len(accuracies)
+                report["average"]["accuracy"] = average_accuracy
+            losses = [value["loss"] for key, value in report.items() if "loss" in value]
+            if len(losses) > 0:
+                average_loss = sum(losses) / len(losses)
+                report["average"]["loss"] = average_loss
+
+        log.info(f"Results for taskpool {self.config.name}: {report}")
+        if self._fabric.is_global_zero and len(self._fabric._loggers) > 0:
+            with open(
+                os.path.join(self._fabric.logger.log_dir, "report.json"), "w"
+            ) as fp:
+                json.dump(report, fp)
+        return report