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

fusion_bench/__init__.py

@@ -41,6 +41,8 @@ _import_structure = {
         "CausalLMBackbonePool",
         "CausalLMPool",
         "CLIPVisionModelPool",
+        "ConvNextForImageClassificationPool",
+        "Dinov2ForImageClassificationPool",
         "GPT2ForSequenceClassificationPool",
         "HuggingFaceGPT2ClassificationPool",
         "NYUv2ModelPool",
@@ -107,6 +109,8 @@ if TYPE_CHECKING:
         CausalLMBackbonePool,
         CausalLMPool,
         CLIPVisionModelPool,
+        ConvNextForImageClassificationPool,
+        Dinov2ForImageClassificationPool,
         GPT2ForSequenceClassificationPool,
         HuggingFaceGPT2ClassificationPool,
         NYUv2ModelPool,

fusion_bench/dataset/clip_dataset.py

@@ -62,6 +62,7 @@ class CLIPDataset(torch.utils.data.Dataset):
         if self.processor is not None:
             if isinstance(self.processor, (ProcessorMixin, BaseImageProcessor)):
                 # Apply the processor to the image to get the input tensor
+                image = image.convert("RGB")  # ensure image is in RGB format
                 inputs = self.processor(images=[image], return_tensors="pt")[
                     "pixel_values"
                 ][0]

fusion_bench/method/__init__.py

@@ -55,6 +55,8 @@ _import_structure = {
         "GPT2LayerWiseAdaMergingAlgorithm",
         "LayerWiseAdaMergingForLlamaSFT",
         "FlanT5LayerWiseAdaMergingAlgorithm",
+        "ResNetLayerWiseAdamerging",
+        "ResNetTaskWiseAdamerging",
     ],
     "pwe_moe": [
         "PWEMoELinearScalarizationForCLIP",

fusion_bench/method/adamerging/__init__.py

@@ -1,6 +1,29 @@
 # flake8: noqa F401
-from .clip_layer_wise_adamerging import CLIPLayerWiseAdaMergingAlgorithm
-from .clip_task_wise_adamerging import CLIPTaskWiseAdaMergingAlgorithm
-from .flan_t5_layer_wise_adamerging import FlanT5LayerWiseAdaMergingAlgorithm
-from .gpt2_layer_wise_adamerging import GPT2LayerWiseAdaMergingAlgorithm
-from .llama_adamerging import LayerWiseAdaMergingForLlamaSFT
+import sys
+from typing import TYPE_CHECKING
+
+from fusion_bench.utils.lazy_imports import LazyImporter
+
+_import_structure = {
+    "clip_layer_wise_adamerging": ["CLIPLayerWiseAdaMergingAlgorithm"],
+    "clip_task_wise_adamerging": ["CLIPTaskWiseAdaMergingAlgorithm"],
+    "flan_t5_layer_wise_adamerging": ["FlanT5LayerWiseAdaMergingAlgorithm"],
+    "gpt2_layer_wise_adamerging": ["GPT2LayerWiseAdaMergingAlgorithm"],
+    "llama_adamerging": ["LayerWiseAdaMergingForLlamaSFT"],
+    "resnet_adamerging": ["ResNetLayerWiseAdamerging", "ResNetTaskWiseAdamerging"],
+}
+
+if TYPE_CHECKING:
+    from .clip_layer_wise_adamerging import CLIPLayerWiseAdaMergingAlgorithm
+    from .clip_task_wise_adamerging import CLIPTaskWiseAdaMergingAlgorithm
+    from .flan_t5_layer_wise_adamerging import FlanT5LayerWiseAdaMergingAlgorithm
+    from .gpt2_layer_wise_adamerging import GPT2LayerWiseAdaMergingAlgorithm
+    from .llama_adamerging import LayerWiseAdaMergingForLlamaSFT
+    from .resnet_adamerging import ResNetLayerWiseAdamerging, ResNetTaskWiseAdamerging
+
+else:
+    sys.modules[__name__] = LazyImporter(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+    )

fusion_bench/method/adamerging/resnet_adamerging.py

@@ -0,0 +1,279 @@
+import os
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Dict, Iterator, Optional, Union, override
+
+import torch
+from omegaconf import DictConfig
+from torch import nn
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from fusion_bench import (
+    BaseAlgorithm,
+    LightningFabricMixin,
+    auto_register_config,
+    get_rankzero_logger,
+    instantiate,
+)
+from fusion_bench.constants import RuntimeConstants
+from fusion_bench.dataset import CLIPDataset
+from fusion_bench.modelpool import ResNetForImageClassificationPool
+from fusion_bench.models.wrappers.layer_wise_fusion import LayerWiseMergedModel
+from fusion_bench.models.wrappers.task_wise_fusion import TaskWiseMergedModel
+from fusion_bench.utils import load_tensor_from_file
+from fusion_bench.utils.data import InfiniteDataLoader
+
+from .entropy_loss import entropy_loss
+from .utils import construct_layer_wise_merged_model, construct_task_wise_merged_model
+
+if TYPE_CHECKING:
+    from transformers import ResNetForImageClassification, ResNetModel
+
+log = get_rankzero_logger(__name__)
+
+
+@auto_register_config
+class _ResNetAdaMergingBase(
+    ABC,
+    LightningFabricMixin,
+    BaseAlgorithm,
+):
+    classification_heads: Dict[str, nn.Module]
+    shuffled_test_loader_iters: Dict[str, Iterator]
+
+    def __init__(
+        self,
+        max_steps: int,
+        optimizer: DictConfig,
+        lr_scheduler: DictConfig,
+        dataloader_kwargs: DictConfig,
+        init_values: Optional[float],
+        clamp_weights: bool = False,
+        tie_weights: bool = True,
+        strict: bool = False,
+        resume_weights_path: Union[str, None] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        if RuntimeConstants.debug:
+            log.info("Debug mode is on, setting max_steps to 10")
+            self.max_steps = 10
+
+    @override
+    def run(self, modelpool: ResNetForImageClassificationPool):
+        self.modelpool = modelpool
+
+        # setup models
+        wrapped_model = self.setup_wrapped_model(modelpool)
+
+        # if max_steps <= 0, skip training and return the merged model directly
+        # this can be used to evaluate the merging weights loaded from `resume_weights_path`
+        if self.max_steps <= 0:
+            # skip_training
+            return wrapped_model.merge_and_unload()
+
+        # setup dataloaders
+        self.setup_dataloaders()
+
+        # configure optimizer and lr_scheduler
+        optimizer = instantiate(self.optimizer, params=[wrapped_model.merge_weight])
+        if self.lr_scheduler is not None:
+            lr_scheduler = instantiate(self.lr_scheduler, optimizer=optimizer)
+        else:
+            lr_scheduler = None
+
+        wrapped_model, optimizer = self.fabric.setup(wrapped_model, optimizer)
+        wrapped_model = self.test_time_adaptation(
+            wrapped_model, optimizer, lr_scheduler
+        )
+
+        # save merging weights
+        if self.log_dir is not None:
+            self.fabric.save(
+                os.path.join(self.log_dir, "checkpoints", "merge_weight.ckpt"),
+                {"merge_weight": wrapped_model.merge_weight},
+            )
+
+        merged_model = wrapped_model.merge_and_unload()
+        if self.log_dir is not None:
+            modelpool.save_model(
+                merged_model,
+                os.path.join(self.log_dir, "checkpoints", "merged_model"),
+                algorithm_config=self.config,
+                description="Merged ResNet model using AdaMerging (E Yang, 2023).",
+            )
+
+        return merged_model
+
+    def test_time_adaptation(
+        self,
+        wrapped_model: TaskWiseMergedModel,
+        optimizer: torch.optim.Optimizer,
+        lr_scheduler: Optional[torch.optim.lr_scheduler.LRScheduler],
+    ):
+        model_names = self.modelpool.model_names
+        wrapped_model.train()
+        wrapped_model.merge_weights()
+
+        for step_idx in tqdm(
+            range(self.max_steps),
+            disable=not self.fabric.is_global_zero,
+            dynamic_ncols=True,
+        ):
+            metrics = {"tta/total_loss": 0.0}
+            for task in model_names:
+                batch = next(self.get_shuffled_test_loader_iter(task))
+                logits = self.compute_logits(wrapped_model, batch[0], task)
+                loss = entropy_loss(logits)
+                metrics[f"tta/{task}_loss"] = loss.item()
+                metrics["tta/total_loss"] += loss.item()
+                self.fabric.backward(loss, retain_graph=True)
+
+            optimizer.step()
+            optimizer.zero_grad()
+            wrapped_model.merge_weights()  # merge weights for the next step
+            if lr_scheduler is not None:
+                lr_scheduler.step()
+
+            self.fabric.log_dict(metrics=metrics, step=step_idx)
+
+        return wrapped_model
+
+    def compute_logits(
+        self, module: Union["ResNetModel", nn.Module], images: torch.Tensor, task: str
+    ) -> torch.Tensor:
+        if self.modelpool.type == "transformers":
+            outputs = module(images, return_dict=True)
+            pooled_output = outputs.pooler_output
+            logits = self.classification_heads[task](pooled_output)
+            return logits
+        else:
+            raise NotImplementedError(
+                f"Model type {self.modelpool.type} is not supported."
+            )
+
+    def setup_dataloaders(self):
+        dataloader_kwargs = dict(self.dataloader_kwargs)
+        dataloader_kwargs["shuffle"] = True  # ensure shuffling for TTA
+        processor = self.modelpool.load_processor()
+        for task in self.modelpool.test_dataset_names:
+            test_dataset = self.modelpool.load_test_dataset(task)
+            test_dataset = CLIPDataset(test_dataset, processor=processor)
+            test_loader = DataLoader(test_dataset, **dataloader_kwargs)
+            self.shuffled_test_loader_iters[task] = iter(
+                InfiniteDataLoader(test_loader)
+            )
+
+    def get_shuffled_test_loader_iter(self, task: str):
+        return self.shuffled_test_loader_iters[task]
+
+    @abstractmethod
+    def setup_wrapped_model(
+        self, modelpool: ResNetForImageClassificationPool
+    ) -> Union[TaskWiseMergedModel, LayerWiseMergedModel]:
+        """
+        Setup the wrapped merged model.
+
+        Args:
+            modelpool (ResNetForImageClassificationPool): The model pool containing pretrained and finetuned models.
+
+        Returns:
+            Union[TaskWiseMergedModel, LayerWiseMergedModel] : The wrapped merged model.
+        """
+        pass
+
+
+class ResNetTaskWiseAdamerging(_ResNetAdaMergingBase):
+    @torch.no_grad()
+    def setup_wrapped_model(self, modelpool: ResNetForImageClassificationPool):
+        pretrained_model = modelpool.load_pretrained_model()
+        finetuned_models = dict(modelpool.named_models())
+
+        if modelpool.type == "transformers":
+            pretrained_model: "ResNetForImageClassification"
+            finetuned_models: Dict[str, "ResNetForImageClassification"]
+            for model_name in finetuned_models:
+                self.classification_heads[model_name] = finetuned_models[
+                    model_name
+                ].classifier
+                # fix the classification head during merging and move to device
+                self.classification_heads[model_name].requires_grad_(False)
+            pretrained_backbone: "ResNetModel" = pretrained_model.resnet
+            finetuned_backbones = [
+                finetuned_models[model_name].resnet for model_name in finetuned_models
+            ]
+        else:
+            raise NotImplementedError(f"Model type {modelpool.type} is not supported.")
+
+        wrapped_model = construct_task_wise_merged_model(
+            pretrained_model=pretrained_backbone,
+            finetuned_models=finetuned_backbones,
+            clamp_weights=self.clamp_weights,
+            tie_weights=self.tie_weights,
+            strict=self.strict,
+        )
+
+        if self.init_values is not None:
+            log.info(f"Initializing merging weights to {self.init_values}")
+            wrapped_model.merge_weight.data.fill_(self.init_values)
+
+        # load merging weights if provided
+        if self.resume_weights_path is not None:
+            merging_weights = load_tensor_from_file(
+                self.resume_weights_path, device="cpu"
+            )
+            log.info(f"Loaded merging weights from {self.resume_weights_path}")
+            assert merging_weights.shape == wrapped_model.merge_weight.shape, (
+                f"Merging weights shape {merging_weights.shape} does not match "
+                f"model's merge_weight shape {wrapped_model.merge_weight.shape}."
+            )
+            wrapped_model.merge_weight.data = merging_weights
+        return wrapped_model
+
+
+class ResNetLayerWiseAdamerging(_ResNetAdaMergingBase):
+    @torch.no_grad()
+    def setup_wrapped_model(self, modelpool: ResNetForImageClassificationPool):
+        pretrained_model = modelpool.load_pretrained_model()
+        finetuned_models = dict(modelpool.named_models())
+
+        if modelpool.type == "transformers":
+            pretrained_model: "ResNetForImageClassification"
+            finetuned_models: Dict[str, "ResNetForImageClassification"]
+            for model_name in finetuned_models:
+                self.classification_heads[model_name] = finetuned_models[
+                    model_name
+                ].classifier
+                # fix the classification head during merging and move to device
+                self.classification_heads[model_name].requires_grad_(False)
+            pretrained_backbone: "ResNetModel" = pretrained_model.resnet
+            finetuned_backbones = [
+                finetuned_models[model_name].resnet for model_name in finetuned_models
+            ]
+        else:
+            raise NotImplementedError(f"Model type {modelpool.type} is not supported.")
+
+        wrapped_model = construct_layer_wise_merged_model(
+            pretrained_model=pretrained_backbone,
+            finetuned_models=finetuned_backbones,
+            clamp_weights=self.clamp_weights,
+            tie_weights=self.tie_weights,
+            strict=self.strict,
+        )
+
+        if self.init_values is not None:
+            log.info(f"Initializing merging weights to {self.init_values}")
+            wrapped_model.merge_weight.data.fill_(self.init_values)
+
+        # load merging weights if provided
+        if self.resume_weights_path is not None:
+            merging_weights = load_tensor_from_file(
+                self.resume_weights_path, device="cpu"
+            )
+            log.info(f"Loaded merging weights from {self.resume_weights_path}")
+            assert merging_weights.shape == wrapped_model.merge_weight.shape, (
+                f"Merging weights shape {merging_weights.shape} does not match "
+                f"model's merge_weight shape {wrapped_model.merge_weight.shape}."
+            )
+            wrapped_model.merge_weight.data = merging_weights
+        return wrapped_model

fusion_bench/method/adamerging/task_wise_adamerging.py

@@ -18,21 +18,9 @@ from fusion_bench.models.wrappers.task_wise_fusion import (
     get_task_wise_weights,
 )
 
-log = logging.getLogger(__name__)
-
-
-def entropy_loss(logits: Tensor) -> Tensor:
-    """
-    Compute the entropy loss of a set of logits.
+from .entropy_loss import entropy_loss
 
-    Args:
-        logits (Tensor): The logits to compute the entropy loss of.
-
-    Returns:
-        Tensor: The entropy loss of the logits.
-    """
-    probs = torch.softmax(logits, dim=-1)
-    return -torch.sum(probs * torch.log(probs + 1e-8), dim=-1).mean()
+log = logging.getLogger(__name__)
 
 
 class TaskWiseAdaMergingAlgorithm(ModelFusionAlgorithm):

fusion_bench/method/adamerging/utils.py

@@ -1,4 +1,9 @@
+from typing import List
+
 import torch
+import torch.nn as nn
+
+from fusion_bench.utils.type import TorchModelType
 
 
 def get_memory_usage(desc):
@@ -13,3 +18,56 @@ def get_memory_usage(desc):
     return (
         f"{desc}\nAllocated Memory: {allocated:.2f} MB\nCached Memory: {cached:.2f} MB"
     )
+
+
+@torch.no_grad()
+def construct_task_wise_merged_model(
+    pretrained_model: TorchModelType,
+    finetuned_models: List[TorchModelType],
+    clamp_weights: bool = False,
+    tie_weights: bool = True,
+    strict: bool = False,
+):
+    from fusion_bench.models.wrappers.task_wise_fusion import (
+        TaskWiseMergedModel,
+        get_task_wise_weights,
+    )
+
+    merging_weights = get_task_wise_weights(num_models=len(finetuned_models))
+    module = TaskWiseMergedModel(
+        task_wise_weight=merging_weights,
+        pretrained_model=pretrained_model,
+        finetuned_models=finetuned_models,
+        clamp_weights=clamp_weights,
+        tie_weights=tie_weights,
+        strict=strict,
+    )
+    return module
+
+
+@torch.no_grad()
+def construct_layer_wise_merged_model(
+    pretrained_model: TorchModelType,
+    finetuned_models: List[TorchModelType],
+    clamp_weights: bool = False,
+    tie_weights: bool = True,
+    strict: bool = False,
+):
+    from fusion_bench.models.wrappers.layer_wise_fusion import (
+        LayerWiseMergedModel,
+        get_layer_wise_weights,
+    )
+
+    merging_weights = get_layer_wise_weights(
+        num_models=len(finetuned_models),
+        num_layers=len([p for p in pretrained_model.parameters() if p.requires_grad]),
+    )
+    module = LayerWiseMergedModel(
+        layer_wise_weight=merging_weights,
+        pretrained_model=pretrained_model,
+        finetuned_models=finetuned_models,
+        clamp_weights=clamp_weights,
+        tie_weights=tie_weights,
+        strict=strict,
+    )
+    return module