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

fusion_bench/__init__.py

@@ -5,46 +5,156 @@
 # ██║     ╚██████╔╝███████║██║╚██████╔╝██║ ╚████║      ██████╔╝███████╗██║ ╚████║╚██████╗██║  ██║
 # ╚═╝      ╚═════╝ ╚══════╝╚═╝ ╚═════╝ ╚═╝  ╚═══╝      ╚═════╝ ╚══════╝╚═╝  ╚═══╝ ╚═════╝╚═╝  ╚═╝
 # flake8: noqa: F401
-from . import (
-    constants,
-    dataset,
-    method,
-    metrics,
-    mixins,
-    modelpool,
-    models,
-    optim,
-    programs,
-    taskpool,
-    tasks,
-    utils,
-)
+import sys
+from typing import TYPE_CHECKING
+
+from fusion_bench.utils.lazy_imports import LazyImporter
+
+from . import constants, metrics, optim, tasks
 from .constants import RuntimeConstants
-from .method import BaseAlgorithm, BaseModelFusionAlgorithm
-from .mixins import auto_register_config
-from .modelpool import BaseModelPool
-from .models import (
-    create_default_model_card,
-    load_model_card_template,
-    save_pretrained_with_remote_code,
-    separate_io,
-)
-from .programs import BaseHydraProgram
-from .taskpool import BaseTaskPool
-from .utils import (
-    BoolStateDictType,
-    LazyStateDict,
-    StateDictType,
-    TorchModelType,
-    cache_with_joblib,
-    get_rankzero_logger,
-    import_object,
-    instantiate,
-    parse_dtype,
-    print_parameters,
-    seed_everything_by_time,
-    set_default_cache_dir,
-    set_print_function_call,
-    set_print_function_call_permeanent,
-    timeit_context,
-)
+from .method import _available_algorithms
+
+_extra_objects = {
+    "RuntimeConstants": RuntimeConstants,
+    "constants": constants,
+    "metrics": metrics,
+    "optim": optim,
+    "tasks": tasks,
+}
+_import_structure = {
+    "dataset": ["CLIPDataset"],
+    "method": _available_algorithms,
+    "mixins": [
+        "CLIPClassificationMixin",
+        "FabricTrainingMixin",
+        "HydraConfigMixin",
+        "LightningFabricMixin",
+        "OpenCLIPClassificationMixin",
+        "PyinstrumentProfilerMixin",
+        "SimpleProfilerMixin",
+        "YAMLSerializationMixin",
+        "auto_register_config",
+    ],
+    "modelpool": [
+        "AutoModelPool",
+        "BaseModelPool",
+        "CausalLMBackbonePool",
+        "CausalLMPool",
+        "CLIPVisionModelPool",
+        "GPT2ForSequenceClassificationPool",
+        "HuggingFaceGPT2ClassificationPool",
+        "NYUv2ModelPool",
+        "OpenCLIPVisionModelPool",
+        "PeftModelForSeq2SeqLMPool",
+        "ResNetForImageClassificationPool",
+        "Seq2SeqLMPool",
+        "SequenceClassificationModelPool",
+    ],
+    "models": [
+        "create_default_model_card",
+        "load_model_card_template",
+        "save_pretrained_with_remote_code",
+        "separate_load",
+        "separate_save",
+    ],
+    "programs": ["BaseHydraProgram", "FabricModelFusionProgram"],
+    "taskpool": [
+        "BaseTaskPool",
+        "CLIPVisionModelTaskPool",
+        "DummyTaskPool",
+        "GPT2TextClassificationTaskPool",
+        "LMEvalHarnessTaskPool",
+        "OpenCLIPVisionModelTaskPool",
+        "NYUv2TaskPool",
+    ],
+    "utils": [
+        "ArithmeticStateDict",
+        "BoolStateDictType",
+        "LazyStateDict",
+        "StateDictType",
+        "TorchModelType",
+        "cache_with_joblib",
+        "get_rankzero_logger",
+        "import_object",
+        "instantiate",
+        "parse_dtype",
+        "print_parameters",
+        "seed_everything_by_time",
+        "set_default_cache_dir",
+        "set_print_function_call",
+        "set_print_function_call_permeanent",
+        "timeit_context",
+    ],
+}
+
+if TYPE_CHECKING:
+    from .dataset import CLIPDataset
+    from .method import BaseAlgorithm, BaseModelFusionAlgorithm
+    from .mixins import (
+        CLIPClassificationMixin,
+        FabricTrainingMixin,
+        HydraConfigMixin,
+        LightningFabricMixin,
+        OpenCLIPClassificationMixin,
+        PyinstrumentProfilerMixin,
+        SimpleProfilerMixin,
+        YAMLSerializationMixin,
+        auto_register_config,
+    )
+    from .modelpool import (
+        AutoModelPool,
+        BaseModelPool,
+        CausalLMBackbonePool,
+        CausalLMPool,
+        CLIPVisionModelPool,
+        GPT2ForSequenceClassificationPool,
+        HuggingFaceGPT2ClassificationPool,
+        NYUv2ModelPool,
+        OpenCLIPVisionModelPool,
+        PeftModelForSeq2SeqLMPool,
+        ResNetForImageClassificationPool,
+        Seq2SeqLMPool,
+        SequenceClassificationModelPool,
+    )
+    from .models import (
+        create_default_model_card,
+        load_model_card_template,
+        save_pretrained_with_remote_code,
+        separate_load,
+        separate_save,
+    )
+    from .programs import BaseHydraProgram, FabricModelFusionProgram
+    from .taskpool import (
+        BaseTaskPool,
+        CLIPVisionModelTaskPool,
+        DummyTaskPool,
+        GPT2TextClassificationTaskPool,
+        LMEvalHarnessTaskPool,
+        NYUv2TaskPool,
+        OpenCLIPVisionModelTaskPool,
+    )
+    from .utils import (
+        ArithmeticStateDict,
+        BoolStateDictType,
+        LazyStateDict,
+        StateDictType,
+        TorchModelType,
+        cache_with_joblib,
+        get_rankzero_logger,
+        import_object,
+        instantiate,
+        parse_dtype,
+        print_parameters,
+        seed_everything_by_time,
+        set_default_cache_dir,
+        set_print_function_call,
+        set_print_function_call_permeanent,
+        timeit_context,
+    )
+else:
+    sys.modules[__name__] = LazyImporter(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        extra_objects=_extra_objects,
+    )

fusion_bench/dataset/__init__.py

@@ -1,16 +1,20 @@
 # flake8: noqa F401
-from datasets import load_dataset
-from omegaconf import DictConfig, open_dict
+import sys
+from typing import TYPE_CHECKING
 
-from fusion_bench.utils import instantiate
+from omegaconf import DictConfig, open_dict
 
-from .clip_dataset import CLIPDataset
+from fusion_bench.utils.lazy_imports import LazyImporter
 
 
 def load_dataset_from_config(dataset_config: DictConfig):
     """
     Load the dataset from the configuration.
     """
+    from datasets import load_dataset
+
+    from fusion_bench.utils import instantiate
+
     assert hasattr(dataset_config, "type"), "Dataset type not specified"
     if dataset_config.type == "instantiate":
         return instantiate(dataset_config.object)
@@ -27,3 +31,22 @@ def load_dataset_from_config(dataset_config: DictConfig):
         return dataset
     else:
         raise ValueError(f"Unknown dataset type: {dataset_config.type}")
+
+
+_extra_objects = {
+    "load_dataset_from_config": load_dataset_from_config,
+}
+_import_structure = {
+    "clip_dataset": ["CLIPDataset"],
+}
+
+if TYPE_CHECKING:
+    from .clip_dataset import CLIPDataset
+
+else:
+    sys.modules[__name__] = LazyImporter(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        extra_objects=_extra_objects,
+    )

fusion_bench/dataset/clip_dataset.py

@@ -6,7 +6,7 @@ from typing import Optional, Tuple
 
 import torch
 from torch.utils.data import Dataset
-from transformers import CLIPProcessor, ProcessorMixin
+from transformers import BaseImageProcessor, CLIPProcessor, ProcessorMixin
 
 __all__ = ["CLIPDataset"]
 
@@ -60,7 +60,7 @@ class CLIPDataset(torch.utils.data.Dataset):
             raise ValueError("Each item should be a dictionary or a tuple of length 2")
         image = item["image"]
         if self.processor is not None:
-            if isinstance(self.processor, ProcessorMixin):
+            if isinstance(self.processor, (ProcessorMixin, BaseImageProcessor)):
                 # Apply the processor to the image to get the input tensor
                 inputs = self.processor(images=[image], return_tensors="pt")[
                     "pixel_values"

fusion_bench/method/__init__.py

@@ -2,6 +2,7 @@
 import sys
 from typing import TYPE_CHECKING
 
+from fusion_bench.utils import join_lists
 from fusion_bench.utils.lazy_imports import LazyImporter
 
 _import_structure = {
@@ -12,6 +13,8 @@ _import_structure = {
     "classification": [
         "ImageClassificationFineTuningForCLIP",
         "ContinualImageClassificationFineTuningForCLIP",
+        "ImageClassificationFineTuning",
+        "ImageClassificationFineTuning_Test",
     ],
     "lm_finetune": ["FullFinetuneSFT", "PeftFinetuneSFT", "BradleyTerryRewardModeling"],
     # analysis
@@ -131,7 +134,10 @@ _import_structure = {
         "ProgressivePruningForMixtral",
     ],
 }
-
+_available_algorithms = join_lists(list(_import_structure.values()))
+_extra_objects = {
+    "_available_algorithms": _available_algorithms,
+}
 
 if TYPE_CHECKING:
     from .ada_svd import AdaSVDMergingForCLIPVisionModel
@@ -141,6 +147,8 @@ if TYPE_CHECKING:
     from .bitdelta import BitDeltaAlgorithm
     from .classification import (
         ContinualImageClassificationFineTuningForCLIP,
+        ImageClassificationFineTuning,
+        ImageClassificationFineTuning_Test,
         ImageClassificationFineTuningForCLIP,
     )
     from .concrete_subspace import (
@@ -252,4 +260,5 @@ else:
         __name__,
         globals()["__file__"],
         _import_structure,
+        extra_objects=_extra_objects,
     )

fusion_bench/method/classification/__init__.py

@@ -1,3 +1,28 @@
 # flake8: noqa F401
-from .clip_finetune import ImageClassificationFineTuningForCLIP
-from .continual_clip_finetune import ContinualImageClassificationFineTuningForCLIP
+import sys
+from typing import TYPE_CHECKING
+
+from fusion_bench.utils.lazy_imports import LazyImporter
+
+_import_structure = {
+    "clip_finetune": ["ImageClassificationFineTuningForCLIP"],
+    "continual_clip_finetune": ["ContinualImageClassificationFineTuningForCLIP"],
+    "image_classification_finetune": [
+        "ImageClassificationFineTuning",
+        "ImageClassificationFineTuning_Test",
+    ],
+}
+
+if TYPE_CHECKING:
+    from .clip_finetune import ImageClassificationFineTuningForCLIP
+    from .continual_clip_finetune import ContinualImageClassificationFineTuningForCLIP
+    from .image_classification_finetune import (
+        ImageClassificationFineTuning,
+        ImageClassificationFineTuning_Test,
+    )
+else:
+    sys.modules[__name__] = LazyImporter(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+    )

fusion_bench/method/classification/image_classification_finetune.py

@@ -0,0 +1,214 @@
+import os
+from typing import Optional
+
+import lightning as L
+import lightning.pytorch.callbacks as pl_callbacks
+import torch
+from lightning.pytorch.loggers import TensorBoardLogger
+from lightning_utilities.core.rank_zero import rank_zero_only
+from lit_learn.lit_modules import ERM_LitModule
+from omegaconf import DictConfig
+from torch import nn
+from torch.utils.data import DataLoader
+from torchmetrics.classification import Accuracy
+
+from fusion_bench import (
+    BaseAlgorithm,
+    BaseModelPool,
+    RuntimeConstants,
+    auto_register_config,
+    get_rankzero_logger,
+    instantiate,
+)
+from fusion_bench.dataset import CLIPDataset
+from fusion_bench.modelpool import ResNetForImageClassificationPool
+from fusion_bench.tasks.clip_classification import get_num_classes
+
+log = get_rankzero_logger(__name__)
+
+
+@auto_register_config
+class ImageClassificationFineTuning(BaseAlgorithm):
+    def __init__(
+        self,
+        max_epochs: Optional[int],
+        max_steps: Optional[int],
+        label_smoothing: float,
+        optimizer: DictConfig,
+        lr_scheduler: DictConfig,
+        dataloader_kwargs: DictConfig,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        assert (max_epochs is None) or (
+            max_steps is None or max_steps < 0
+        ), "Only one of max_epochs or max_steps should be set."
+        self.training_interval = "epoch" if max_epochs is not None else "step"
+        if self.training_interval == "epoch":
+            self.max_steps = -1
+        log.info(f"Training interval: {self.training_interval}")
+        log.info(f"Max epochs: {max_epochs}, max steps: {max_steps}")
+
+    def run(self, modelpool: ResNetForImageClassificationPool):
+        # load model and dataset
+        model = modelpool.load_pretrained_or_first_model()
+        assert isinstance(model, nn.Module), "Loaded model is not a nn.Module."
+
+        assert (
+            len(modelpool.train_dataset_names) == 1
+        ), "Exactly one training dataset is required."
+        self.dataset_name = dataset_name = modelpool.train_dataset_names[0]
+        num_classes = get_num_classes(dataset_name)
+        train_dataset = modelpool.load_train_dataset(dataset_name)
+        train_dataset = CLIPDataset(
+            train_dataset, processor=modelpool.load_processor(stage="train")
+        )
+        train_loader = self.get_dataloader(train_dataset, stage="train")
+        if modelpool.has_val_dataset:
+            val_dataset = modelpool.load_val_dataset(dataset_name)
+            val_dataset = CLIPDataset(
+                val_dataset, processor=modelpool.load_processor(stage="val")
+            )
+            val_loader = self.get_dataloader(val_dataset, stage="val")
+
+        # configure optimizer
+        optimizer = instantiate(self.optimizer, params=model.parameters())
+        if self.lr_scheduler is not None:
+            lr_scheduler = instantiate(self.lr_scheduler, optimizer=optimizer)
+            optimizer = {
+                "optimizer": optimizer,
+                "lr_scheduler": {
+                    "scheduler": lr_scheduler,
+                    "interval": self.training_interval,
+                    "frequency": 1,
+                },
+            }
+        log.info(f"optimizer:\n{optimizer}")
+
+        lit_module = ERM_LitModule(
+            model,
+            optimizer,
+            objective=nn.CrossEntropyLoss(label_smoothing=self.label_smoothing),
+            metrics={
+                "acc@1": Accuracy(task="multiclass", num_classes=num_classes),
+                "acc@5": Accuracy(task="multiclass", num_classes=num_classes, top_k=5),
+            },
+        )
+
+        log_dir = (
+            self._program.path.log_dir
+            if self._program is not None
+            else "outputs/lightning_logs"
+        )
+        trainer = L.Trainer(
+            max_epochs=self.max_epochs,
+            max_steps=self.max_steps,
+            accelerator="auto",
+            devices="auto",
+            callbacks=[
+                pl_callbacks.LearningRateMonitor(logging_interval="step"),
+                pl_callbacks.DeviceStatsMonitor(),
+            ],
+            logger=TensorBoardLogger(
+                save_dir=log_dir,
+                name="",
+            ),
+            fast_dev_run=RuntimeConstants.debug,
+        )
+
+        trainer.fit(
+            lit_module, train_dataloaders=train_loader, val_dataloaders=val_loader
+        )
+        model = lit_module.model
+        if rank_zero_only.rank == 0:
+            log.info(f"Saving the final model to {log_dir}/raw_checkpoints/final")
+            modelpool.save_model(
+                model,
+                path=os.path.join(
+                    trainer.log_dir if trainer.log_dir is not None else log_dir,
+                    "raw_checkpoints",
+                    "final",
+                ),
+            )
+        return model
+
+    def get_dataloader(self, dataset, stage: str):
+        assert stage in ["train", "val", "test"], f"Invalid stage: {stage}"
+        dataloader_kwargs = dict(self.dataloader_kwargs)
+        if "shuffle" not in dataloader_kwargs:
+            dataloader_kwargs["shuffle"] = stage == "train"
+        return DataLoader(dataset, **dataloader_kwargs)
+
+
+@auto_register_config
+class ImageClassificationFineTuning_Test(BaseAlgorithm):
+    def __init__(self, checkpoint_path: str, dataloader_kwargs: DictConfig, **kwargs):
+        super().__init__(**kwargs)
+
+    def run(self, modelpool: BaseModelPool):
+        assert (
+            modelpool.has_val_dataset or modelpool.has_test_dataset
+        ), "No validation or test dataset found in the model pool."
+
+        # load model and dataset
+        model = modelpool.load_pretrained_or_first_model()
+        assert isinstance(model, nn.Module), "Loaded model is not a nn.Module."
+
+        if modelpool.has_test_dataset:
+            assert (
+                len(modelpool.test_dataset_names) == 1
+            ), "Exactly one test dataset is required."
+            self.dataset_name = dataset_name = modelpool.test_dataset_names[0]
+            dataset = modelpool.load_test_dataset(dataset_name)
+            dataset = CLIPDataset(
+                dataset, processor=modelpool.load_processor(stage="test")
+            )
+        else:
+            assert (
+                len(modelpool.val_dataset_names) == 1
+            ), "Exactly one validation dataset is required."
+            self.dataset_name = dataset_name = modelpool.val_dataset_names[0]
+            dataset = modelpool.load_val_dataset(dataset_name)
+            dataset = CLIPDataset(
+                dataset, processor=modelpool.load_processor(stage="test")
+            )
+        num_classes = get_num_classes(dataset_name)
+
+        test_loader = self.get_dataloader(dataset, stage="test")
+
+        if self.checkpoint_path is None:
+            lit_module = ERM_LitModule(
+                model,
+                metrics={
+                    "acc@1": Accuracy(task="multiclass", num_classes=num_classes),
+                    "acc@5": Accuracy(
+                        task="multiclass", num_classes=num_classes, top_k=5
+                    ),
+                },
+            )
+        else:
+            lit_module = ERM_LitModule.load_from_checkpoint(
+                checkpoint_path=self.checkpoint_path,
+                model=model,
+                metrics={
+                    "acc@1": Accuracy(task="multiclass", num_classes=num_classes),
+                    "acc@5": Accuracy(
+                        task="multiclass", num_classes=num_classes, top_k=5
+                    ),
+                },
+            )
+
+        trainer = L.Trainer(
+            devices=1, num_nodes=1, logger=False, fast_dev_run=RuntimeConstants.debug
+        )
+
+        test_metrics = trainer.test(lit_module, dataloaders=test_loader)
+        log.info(f"Test metrics: {test_metrics}")
+        return model
+
+    def get_dataloader(self, dataset, stage: str):
+        assert stage in ["train", "val", "test"], f"Invalid stage: {stage}"
+        dataloader_kwargs = dict(self.dataloader_kwargs)
+        if "shuffle" not in dataloader_kwargs:
+            dataloader_kwargs["shuffle"] = stage == "train"
+        return DataLoader(dataset, **dataloader_kwargs)

fusion_bench/method/opcm/opcm.py

@@ -87,6 +87,7 @@ class OPCMForCLIP(
         # get the average model
         with self.profile("loading model"):
             merged_model = modelpool.load_model(model_names[0])
+            assert merged_model is not None, "Failed to load the first model"
 
         if self.evaluate_on_every_step:
             with self.profile("evaluating model"):

fusion_bench/method/pwe_moe/module.py

@@ -13,8 +13,6 @@ import torch.func
 from torch import Tensor, nn
 from torch.nn import functional as F
 
-from fusion_bench.utils import join_list
-
 log = logging.getLogger(__name__)
 
 

fusion_bench/method/tall_mask/task_arithmetic.py

@@ -15,7 +15,7 @@ from fusion_bench.utils.state_dict_arithmetic import (
     state_dict_add,
     state_dict_binary_mask,
     state_dict_diff_abs,
-    state_dict_hadmard_product,
+    state_dict_hadamard_product,
     state_dict_mul,
     state_dict_sub,
     state_dict_sum,
@@ -111,7 +111,7 @@ class TallMaskTaskArithmeticAlgorithm(
 
         with self.profile("compress and retrieve"):
             for model_name in modelpool.model_names:
-                retrieved_task_vector = state_dict_hadmard_product(
+                retrieved_task_vector = state_dict_hadamard_product(
                     tall_masks[model_name], multi_task_vector
                 )
                 retrieved_state_dict = state_dict_add(

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,