fusion-bench 0.2.12__py3-none-any.whl → 0.2.13__py3-none-any.whl

fusion_bench/taskpool/clip_vision/clip_rankone_moe_taskpool.py

@@ -12,36 +12,7 @@ from fusion_bench.models.hf_clip import HFCLIPClassifier
 from fusion_bench.models.rankone_moe import RankOneMoE
 
 from .taskpool import CLIPVisionModelTaskPool
-
-
-class LayerWiseRoutingWeightSaver:
-    def __init__(self, save_path: Path, max_num: Optional[int] = None):
-        self.save_path = save_path
-        self.max_num = max_num
-        self.routing_weights = []
-
-    def __call__(self, module, input: Tuple[Tensor], output: Tensor):
-        assert isinstance(output, Tensor), "Output is expected to be a Tensor"
-        # (batch_size, num_tokens, num_experts)
-        routing_weights = output.detach().cpu()
-        if self.max_num is not None and self.max_num > 0:
-            if len(self.routing_weights) > self.max_num:
-                return
-            elif routing_weights.size(0) + len(self.routing_weights) > self.max_num:
-                self.routing_weights.append(
-                    routing_weights[: self.max_num - len(self.routing_weights)]
-                )
-            else:
-                self.routing_weights.append(routing_weights)
-        else:
-            self.routing_weights.append(routing_weights)
-
-    def save_routing_weights(self):
-        routing_weights = torch.cat(self.routing_weights, dim=0)
-        if self.save_path is not None:
-            self.save_path.parent.mkdir(parents=True, exist_ok=True)
-            print(f"Saving routing weights to {self.save_path}")
-            torch.save(routing_weights, self.save_path)
+from .utils.routing_analysis_utils import LayerWiseRoutingWeightSaver
 
 
 class RankoneMoECLIPVisionModelTaskPool(CLIPVisionModelTaskPool):
@@ -109,4 +80,5 @@ class RankoneMoECLIPVisionModelTaskPool(CLIPVisionModelTaskPool):
             # remove hooks for saving layer-wise routing weights
             for i, handle in self._layer_wise_routing_weights_save_hook_handles.items():
                 self._layer_wise_routing_weights_save_hooks[i].save_routing_weights()
+                self._layer_wise_routing_weights_save_hook_handles.pop(i)
                 handle.remove()

fusion_bench/taskpool/clip_vision/clip_smile_taskpool.py

@@ -0,0 +1,102 @@
+from copy import deepcopy
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+from torch import Tensor
+from torch.utils.hooks import RemovableHandle
+from transformers import CLIPModel, CLIPProcessor, CLIPVisionModel
+from transformers.models.clip.modeling_clip import CLIPVisionTransformer
+
+from fusion_bench.method.smile_upscaling import SmileMoELinear
+from fusion_bench.models.hf_clip import HFCLIPClassifier
+
+from .taskpool import CLIPVisionModelTaskPool
+from .utils.routing_analysis_utils import LayerWiseRoutingWeightSaver
+
+
+class SmileCLIPVisionModelTaskPool(CLIPVisionModelTaskPool):
+
+    # hooks and handles for saving layer-wise routing weights
+    _layer_wise_routing_weights_save_hooks: Dict[Any, LayerWiseRoutingWeightSaver] = {}
+    _layer_wise_routing_weights_save_hook_handles: Dict[Any, RemovableHandle] = {}
+
+    def __init__(
+        self,
+        linear_module_names: Union[List[str], str],
+        layer_wise_routing_weights_save_path: Optional[str],
+        layer_wise_routing_weights_max_num: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Initialize the SMILECLIPVisionModelTaskPool.
+
+        Args:
+            linear_module_names (Union[List[str], str]): The names of the linear modules to save the layer-wise routing weights for.
+            layer_wise_routing_weights_save_path (Optional[str]): The path to save the layer-wise routing weights.
+            layer_wise_routing_weights_max_num (Optional[int]): The maximum number of layer-wise routing weights to save.
+        """
+        # linear module names
+        assert linear_module_names is not None, "linear_module_names must be provided"
+        self.linear_module_names = (
+            [linear_module_names]
+            if isinstance(linear_module_names, str)
+            else list(linear_module_names)
+        )
+        # save path for layer-wise routing weights
+        self._layer_wise_routing_weights_save_path = (
+            layer_wise_routing_weights_save_path
+        )
+        self.layer_wise_routing_weights_save_path = (
+            Path(layer_wise_routing_weights_save_path)
+            if layer_wise_routing_weights_save_path is not None
+            else None
+        )
+        self.layer_wise_routing_weights_max_num = layer_wise_routing_weights_max_num
+        super().__init__(**kwargs)
+
+    def on_task_evaluation_begin(self, classifier: HFCLIPClassifier, task_name: str):
+        super().on_task_evaluation_begin(classifier, task_name)
+        if self.layer_wise_routing_weights_save_path is not None:
+            # setup hooks for saving layer-wise routing weights
+            assert isinstance(
+                classifier.clip_model.vision_model,
+                (CLIPVisionTransformer, CLIPVisionModel),
+            ), "Vision model is expected to be a CLIPVisionTransformer"
+            vision_model = classifier.clip_model.vision_model
+            if isinstance(vision_model, CLIPVisionModel):
+                vision_model = vision_model.vision_model
+                # assign forward hooks for each layer
+
+            for i, layer in enumerate(vision_model.encoder.layers):
+                for linear_module_name in self.linear_module_names:
+                    linear_module = layer.get_submodule(linear_module_name)
+                    assert isinstance(
+                        linear_module,
+                        (SmileMoELinear),
+                    ), f"Linear module is expected to be a SmileMoELinear, but got {type(linear_module)}"
+                    # layer-wise routing weights
+                    hook = LayerWiseRoutingWeightSaver(
+                        self.layer_wise_routing_weights_save_path
+                        / task_name
+                        / f"layer_{i}_{linear_module_name}.pt",
+                        max_num=self.layer_wise_routing_weights_max_num,
+                    )
+                    self._layer_wise_routing_weights_save_hooks[
+                        (i, linear_module_name)
+                    ] = hook
+                    self._layer_wise_routing_weights_save_hook_handles[
+                        (i, linear_module_name)
+                    ] = linear_module.gate.register_forward_hook(hook)
+
+    def on_task_evaluation_end(self):
+        super().on_task_evaluation_end()
+        if self.layer_wise_routing_weights_save_path is not None:
+            # remove hooks for saving layer-wise routing weights
+            for (
+                key,
+                handle,
+            ) in self._layer_wise_routing_weights_save_hook_handles.items():
+                self._layer_wise_routing_weights_save_hooks[key].save_routing_weights()
+                self._layer_wise_routing_weights_save_hook_handles.pop(key)
+                handle.remove()

fusion_bench/taskpool/clip_vision/clip_sparse_wemoe_taskpool.py

@@ -15,36 +15,7 @@ from fusion_bench.models.sparse_we_moe import (
 )
 
 from .taskpool import CLIPVisionModelTaskPool
-
-
-class LayerWiseRoutingWeightSaver:
-    def __init__(self, save_path: Path, max_num: Optional[int] = None):
-        self.save_path = save_path
-        self.max_num = max_num
-        self.routing_weights = []
-
-    def __call__(self, module, input: Tuple[Tensor], output: Tensor):
-        assert isinstance(output, Tensor), "Output is expected to be a Tensor"
-        # (batch_size, num_tokens, num_experts)
-        routing_weights = output.detach().cpu()
-        if self.max_num is not None and self.max_num > 0:
-            if len(self.routing_weights) > self.max_num:
-                return
-            elif routing_weights.size(0) + len(self.routing_weights) > self.max_num:
-                self.routing_weights.append(
-                    routing_weights[: self.max_num - len(self.routing_weights)]
-                )
-            else:
-                self.routing_weights.append(routing_weights)
-        else:
-            self.routing_weights.append(routing_weights)
-
-    def save_routing_weights(self):
-        routing_weights = torch.cat(self.routing_weights, dim=0)
-        if self.save_path is not None:
-            self.save_path.parent.mkdir(parents=True, exist_ok=True)
-            print(f"Saving routing weights to {self.save_path}")
-            torch.save(routing_weights, self.save_path)
+from .utils.routing_analysis_utils import LayerWiseRoutingWeightSaver
 
 
 class SparseWEMoECLIPVisionModelTaskPool(CLIPVisionModelTaskPool):
@@ -117,4 +88,5 @@ class SparseWEMoECLIPVisionModelTaskPool(CLIPVisionModelTaskPool):
             # remove hooks for saving layer-wise routing weights
             for i, handle in self._layer_wise_routing_weights_save_hook_handles.items():
                 self._layer_wise_routing_weights_save_hooks[i].save_routing_weights()
+                self._layer_wise_routing_weights_save_hook_handles.pop(i)
                 handle.remove()

fusion_bench/taskpool/clip_vision/taskpool.py

@@ -32,8 +32,7 @@ from fusion_bench.mixins import LightningFabricMixin
 from fusion_bench.models.hf_clip import HFCLIPClassifier
 from fusion_bench.taskpool import BaseTaskPool
 from fusion_bench.tasks.clip_classification import get_classnames_and_templates
-from fusion_bench.utils import instantiate
-from fusion_bench.utils.parameters import count_parameters
+from fusion_bench.utils import count_parameters, instantiate
 
 if TYPE_CHECKING:
     from fusion_bench.models.surgery.surgerymodelwrapper import SurgeryModelWrapper

fusion_bench/taskpool/clip_vision/utils/routing_analysis_utils.py

@@ -0,0 +1,65 @@
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+import torch
+from torch import Tensor
+
+
+def _number_of_samples(routing_weights: List[Tensor]):
+    count = 0
+    for routing_weight in routing_weights:
+        count += routing_weight.size(0)
+    return count
+
+
+class LayerWiseRoutingWeightSaver:
+    """
+    A hook for saving layer-wise routing weights.
+    """
+
+    save_path: Path
+    "The path to save the layer-wise routing weights."
+    max_num: Optional[int]
+    "The maximum number of layer-wise routing weights to save. If None, all routing weights will be saved."
+    routing_weights: List[Tensor]
+    "The list of layer-wise routing weights."
+
+    def __init__(self, save_path: Path, max_num: Optional[int] = None):
+        """
+        Args:
+            save_path (Path): The path to save the layer-wise routing weights.
+            max_num (Optional[int]): The maximum number of layer-wise routing weights to save. If None, all routing weights will be saved.
+        """
+        self.save_path = save_path
+        self.max_num = max_num
+        self.routing_weights = []
+
+    def __call__(self, module, input: Tuple[Tensor], output: Tensor):
+        assert isinstance(output, Tensor), "Output is expected to be a Tensor"
+        # (batch_size, num_tokens, num_experts)
+        routing_weights = output.detach().cpu()
+        if self.max_num is not None and self.max_num > 0:
+            if _number_of_samples(self.routing_weights) > self.max_num:
+                return
+            elif (
+                routing_weights.size(0) + _number_of_samples(self.routing_weights)
+                > self.max_num
+            ):
+                self.routing_weights.append(
+                    routing_weights[
+                        : self.max_num - _number_of_samples(self.routing_weights)
+                    ]
+                )
+            else:
+                self.routing_weights.append(routing_weights)
+        else:
+            self.routing_weights.append(routing_weights)
+
+    def save_routing_weights(self):
+        routing_weights = torch.cat(self.routing_weights, dim=0)
+        if self.save_path is not None:
+            self.save_path.parent.mkdir(parents=True, exist_ok=True)
+            print(
+                f"Saving routing weights to {self.save_path}. Size: {routing_weights.size()}"
+            )
+            torch.save(routing_weights, self.save_path)

fusion_bench/taskpool/gpt2_text_classification.py

@@ -139,11 +139,40 @@ class GPT2TextClassificationTaskPool(BaseTaskPool, LightningFabricMixin):
         return dataloader
 
     @override
-    def evaluate(self, model: GPT2Model):
+    def evaluate(self, model: GPT2Model, name: str = None):
+        """Evaluate the model on the test datasets.
+
+        Args:
+            model (GPT2Model): The model to evaluate.
+            name (str, optional): The name of the model. Defaults to None. This is used to identify the model in the report.
+
+        Returns:
+            dict: A dictionary containing the evaluation results for each task.
+        """
         report = {}
+        if name is not None:
+            report["name"] = name
         for task_name in (pbar := tqdm(self._test_datasets, desc="Evaluating tasks")):
             pbar.set_description(f"Evaluating task {task_name}")
             dataloader = self.get_test_dataloader(task_name)
             result = self.evaluate_single_task(task_name, model, dataloader)
             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 isinstance(value, dict) and "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"Evaluation Result: {report}")
         return report

fusion_bench/taskpool/openclip_vision/__init__.py

@@ -0,0 +1 @@
+from .openclip_taskpool import OpenCLIPVisionModelTaskPool

fusion_bench/taskpool/openclip_vision/openclip_taskpool.py

@@ -0,0 +1,196 @@
+import itertools
+import json
+import logging
+import os
+from typing import TYPE_CHECKING, Callable, Dict, Optional, Union
+
+import lightning.fabric
+import open_clip
+import torch
+from omegaconf import DictConfig
+from torch.nn import functional as F
+from torch.utils.data import DataLoader, Dataset
+from torchmetrics import Accuracy, MeanMetric
+from torchmetrics.classification.accuracy import MulticlassAccuracy
+from tqdm.auto import tqdm
+
+from fusion_bench import BaseTaskPool
+from fusion_bench.dataset.clip_dataset import CLIPDataset
+from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.modelpool.openclip_vision.modelpool import load_classifier_head
+from fusion_bench.models.open_clip import (
+    ClassificationHead,
+    ImageClassifier,
+    ImageEncoder,
+)
+from fusion_bench.models.open_clip.variables_and_paths import OPENCLIP_CACHEDIR
+from fusion_bench.utils import count_parameters, instantiate
+
+if TYPE_CHECKING:
+    from fusion_bench.modelpool import OpenCLIPVisionModelPool
+    from fusion_bench.programs import FabricModelFusionProgram
+
+log = logging.getLogger(__name__)
+
+
+class OpenCLIPVisionModelTaskPool(
+    BaseTaskPool,
+    LightningFabricMixin,
+):
+    _is_setup = False
+
+    _program: "FabricModelFusionProgram"
+
+    processor: Optional[Callable] = None
+    test_datasets: Dict[str, CLIPDataset]
+
+    def __init__(
+        self,
+        test_datasets: Union[DictConfig, Dict[str, Dataset]],
+        classification_heads: Union[DictConfig, Dict[str, ClassificationHead]],
+        dataloader_kwargs: DictConfig,
+        model_name: Optional[str] = None,
+        fast_dev_run: bool = False,
+        **kwargs,
+    ):
+        self._test_datasets = test_datasets
+        self._classifier_heads = classification_heads
+        self._dataloader_kwargs = dataloader_kwargs
+        self._model_name = model_name
+        self.fast_dev_run = fast_dev_run
+        super().__init__(**kwargs)
+
+    def setup(self):
+        # setup the processor
+        if self._program is not None and self._program.modelpool is not None:
+            modelpool: "OpenCLIPVisionModelPool" = self._program.modelpool
+            self.processor = modelpool.test_processor
+        elif self._model_name is not None:
+            _, _, self.processor = open_clip.create_model_and_transforms(
+                self._model_name,
+                pretrained="openai",
+                cache_dir=OPENCLIP_CACHEDIR,
+            )
+        else:
+            raise ValueError("Modelpool or model_name is not set")
+
+        # setup the test datasets
+        self.test_datasets = {
+            name: instantiate(dataset) if isinstance(dataset, DictConfig) else dataset
+            for name, dataset in self._test_datasets.items()
+        }
+        self.test_datasets = {
+            name: CLIPDataset(dataset, self.processor)
+            for name, dataset in self.test_datasets.items()
+        }
+        self.test_dataloaders = {
+            name: self.fabric.setup_dataloaders(
+                DataLoader(dataset, **self._dataloader_kwargs)
+            )
+            for name, dataset in self.test_datasets.items()
+        }
+
+        # setup classifier heads
+        self.classifier_heads = {
+            name: load_classifier_head(head).to(self.fabric.device)
+            for name, head in self._classifier_heads.items()
+        }
+        self._is_setup = True
+
+    @torch.no_grad()
+    def _evaluate(
+        self,
+        classifier: ImageClassifier,
+        test_loader: DataLoader,
+        num_classes: int,
+        task_name: str,
+    ):
+        accuracy: MulticlassAccuracy = Accuracy(
+            task="multiclass", num_classes=num_classes
+        )
+        classifier.eval()
+        loss_metric = MeanMetric()
+        # if fast_dev_run is set, we only evaluate on a batch of the data
+        if self.fast_dev_run:
+            log.info("Running under fast_dev_run mode, evaluating on a single batch.")
+            test_loader = itertools.islice(test_loader, 1)
+        else:
+            test_loader = test_loader
+
+        pbar = tqdm(
+            test_loader,
+            desc=f"Evaluating {task_name}",
+            leave=False,
+            dynamic_ncols=True,
+        )
+        for batch in pbar:
+            inputs, targets = batch
+            logits = classifier(inputs)
+            loss = F.cross_entropy(logits, targets)
+            loss_metric.update(loss.detach().cpu())
+            acc = accuracy(logits.detach().cpu(), targets.detach().cpu())
+            pbar.set_postfix(
+                {
+                    "accuracy": accuracy.compute().item(),
+                    "loss": loss_metric.compute().item(),
+                }
+            )
+
+        acc = accuracy.compute().item()
+        loss = loss_metric.compute().item()
+        results = {"accuracy": acc, "loss": loss}
+        return results
+
+    def evaluate(self, model: ImageEncoder, **kwargs):
+        if not self._is_setup:
+            self.setup()
+
+        report = {}
+        # collect basic model information
+        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 not lightning.fabric.is_wrapped(model):
+            model = self.fabric.setup_module(model)
+
+        pbar = tqdm(
+            self.test_dataloaders.items(),
+            desc="Evaluating tasks",
+            total=len(self.test_dataloaders),
+        )
+        for task_name, test_dataloader in pbar:
+            classifier = ImageClassifier(model, self.classifier_heads[task_name])
+            num_classes = self.classifier_heads[task_name].weight.size(0)
+            result = self._evaluate(
+                classifier,
+                test_dataloader,
+                num_classes=num_classes,
+                task_name=task_name,
+            )
+            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"Evaluation Result: {report}")
+        if self.fabric.is_global_zero and len(self.fabric._loggers) > 0:
+            with open(os.path.join(self.log_dir, "report.json"), "w") as fp:
+                json.dump(report, fp)
+        return report

fusion_bench/utils/data.py

@@ -9,6 +9,18 @@ from torch.utils.data import DataLoader, Dataset
 
 
 class InfiniteDataLoader:
+    """
+    A wrapper class for DataLoader to create an infinite data loader.
+    This is useful in case we are only interested in the number of steps and not the number of epochs.
+
+    This class wraps a DataLoader and provides an iterator that resets
+    when the end of the dataset is reached, creating an infinite loop.
+
+    Attributes:
+        data_loader (DataLoader): The DataLoader to wrap.
+        data_iter (iterator): An iterator over the DataLoader.
+    """
+
     def __init__(self, data_loader: DataLoader):
         self.data_loader = data_loader
         self.data_iter = iter(data_loader)

fusion_bench/utils/devices.py

@@ -229,3 +229,17 @@ def cleanup_cuda():
     gc.collect()
     torch.cuda.empty_cache()
     torch.cuda.reset_peak_memory_stats()
+
+
+def print_memory_usage(print_fn=print):
+    """
+    Print the current GPU memory usage.
+
+    Returns:
+        str: A string containing the allocated and cached memory in MB.
+    """
+    allocated = torch.cuda.memory_allocated() / 1024**2  # 转换为 MB
+    cached = torch.cuda.memory_reserved() / 1024**2  # 转换为 MB
+    print_str = f"Allocated Memory: {allocated:.2f} MB\nCached Memory: {cached:.2f} MB"
+    print_fn(print_str)
+    return print_str

fusion_bench/utils/instantiate.py

@@ -2,6 +2,7 @@
 # Modified from Hydra
 import copy
 import functools
+from contextlib import contextmanager
 from enum import Enum
 from textwrap import dedent
 from typing import Any, Callable, Dict, List, Sequence, Tuple, Union
@@ -30,6 +31,17 @@ Function to be used for printing function calls.
 CATCH_EXCEPTION = True
 
 
+@contextmanager
+def set_print_function_call(value: bool):
+    global PRINT_FUNCTION_CALL
+    old_value = PRINT_FUNCTION_CALL
+    PRINT_FUNCTION_CALL = value
+    try:
+        yield
+    finally:
+        PRINT_FUNCTION_CALL = old_value
+
+
 def is_instantiable(config: Union[DictConfig, Any]) -> bool:
     if OmegaConf.is_dict(config):
         return "_target_" in config

fusion_bench/utils/misc.py

@@ -1,6 +1,6 @@
-from typing import Iterable
+from typing import Iterable, List
 
-__all__ = ["first", "has_length"]
+__all__ = ["first", "has_length", "join_list"]
 
 
 def first(iterable: Iterable):
@@ -16,3 +16,10 @@ def has_length(dataset):
     except TypeError:
         # TypeError: len() of unsized object
         return False
+
+
+def join_list(list_of_list: List[List]):
+    ans = []
+    for item in list_of_list:
+        ans.extend(item)
+    return ans

fusion_bench/utils/packages.py

@@ -82,3 +82,17 @@ def import_object(abs_obj_name: str):
     module_name, obj_name = abs_obj_name.rsplit(".", 1)
     module = importlib.import_module(module_name)
     return getattr(module, obj_name)
+
+
+def compare_versions(v1, v2):
+    """Compare two version strings.
+    Returns -1 if v1 < v2, 0 if v1 == v2, 1 if v1 > v2"""
+
+    v1 = version.parse(v1)
+    v2 = version.parse(v2)
+    if v1 < v2:
+        return -1
+    elif v1 > v2:
+        return 1
+    else:
+        return 0

fusion_bench/utils/parameters.py

@@ -252,7 +252,7 @@ def print_parameters(
 
 
 def check_parameters_all_equal(
-    list_of_param_names: List[Union[StateDictType, nn.Module, List[str]]]
+    list_of_param_names: List[Union[StateDictType, nn.Module, List[str]]],
 ) -> None:
     """
     Checks if all models have the same parameters.

fusion_bench/utils/tensorboard.py

@@ -1,5 +1,5 @@
 """
-functions deal with tensorboard logs.    
+functions deal with tensorboard logs.
 """
 
 from typing import Dict, Iterable, List

{fusion_bench-0.2.12.dist-info → fusion_bench-0.2.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fusion_bench
-Version: 0.2.12
+Version: 0.2.13
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
 License: MIT License