fusion-bench 0.2.7__py3-none-any.whl → 0.2.8__py3-none-any.whl

fusion_bench/mixins/clip_classification.py

@@ -132,13 +132,13 @@ class CLIPClassificationMixin(LightningFabricMixin):
 
         # get cache directory
         if self.modelpool.has_pretrained:
-            model_name = self.modelpool.get_model_config(
-                "_pretrained_"
-            ).pretrained_model_name_or_path
+            model_name = self.modelpool.get_model_config("_pretrained_")
+            if not isinstance(model_name, str):
+                model_name = model_name.pretrained_model_name_or_path
         else:
-            model_name = self.modelpool.get_model_config(
-                self.modelpool.model_names[0]
-            ).pretrained_model_name_or_path
+            model_name = self.modelpool.get_model_config(self.modelpool.model_names[0])
+            if not isinstance(model_name, str):
+                model_name = model_name.pretrained_model_name_or_path
         cache_dir = os.path.join(
             self.zeroshot_weights_cache_dir,
             os.path.normpath(model_name.split("/")[-1]),

fusion_bench/mixins/lightning_fabric.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, List, Optional, TypeVar
 
 import lightning as L
 import torch
+from lightning.fabric.connector import _is_using_cli
 from lightning.fabric.loggers import TensorBoardLogger
 from lightning.fabric.utilities.rank_zero import rank_zero_only
 from omegaconf import DictConfig, OmegaConf
@@ -79,7 +80,8 @@ class LightningFabricMixin:
                 self._fabric_instance = L.Fabric()
             else:
                 self._fabric_instance = instantiate(config.fabric)
-            self._fabric_instance.launch()
+            if not _is_using_cli():  # if not using cli, launch the fabric
+                self._fabric_instance.launch()
             # Set the log directory in config if it is not already set
             if (
                 self.log_dir is not None

fusion_bench/modelpool/base_pool.py

@@ -147,7 +147,6 @@ class BaseModelPool(BaseYAMLSerializableModel):
             DictConfig: The configuration for the specified model.
         """
         model_config = self._models[model_name]
-        assert isinstance(model_config, DictConfig), "Model config must be a DictConfig"
         if return_copy:
             model_config = deepcopy(model_config)
         return model_config

fusion_bench/modelpool/clip_vision/modelpool.py

@@ -1,8 +1,11 @@
 import logging
 from copy import deepcopy
-from typing import Optional
+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
 
@@ -36,17 +39,29 @@ class CLIPVisionModelPool(BaseModelPool):
 
     def load_processor(self, *args, **kwargs) -> CLIPProcessor:
         assert self._processor is not None, "Processor is not defined in the config"
-        processor = instantiate(self._processor, *args, **kwargs)
+        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]
-        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
+
+        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):
@@ -59,3 +74,72 @@ class CLIPVisionModelPool(BaseModelPool):
         """
         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/models/surgery/__init__.py

@@ -0,0 +1 @@
+from .surgerymodelwrapper import SurgeryModelWrapper

fusion_bench/models/surgery/surgerymodelwrapper.py

@@ -1,5 +1,5 @@
 import math
-from typing import TYPE_CHECKING, List, Union, Callable, Generic
+from typing import TYPE_CHECKING, Callable, Generic, List, Union
 
 import torch
 from torch import nn
@@ -7,6 +7,7 @@ from transformers.models.clip.modeling_clip import (
     CLIPVisionModel,
     CLIPVisionTransformer,
 )
+
 from fusion_bench.utils.type import TorchModelType
 
 

fusion_bench/models/wrappers/layer_wise_fusion.py

@@ -16,7 +16,7 @@ import torch
 from torch import Tensor, nn
 from torch.func import functional_call
 
-from fusion_bench.utils.type import TorchModelType, StateDictType
+from fusion_bench.utils.type import StateDictType, TorchModelType
 
 __all__ = ["get_layer_wise_weights", "fuse_weights", "LayerWiseMergedModel"]
 

fusion_bench/models/wrappers/task_wise_fusion.py

@@ -22,7 +22,7 @@ import torch
 from torch import Tensor, nn
 from torch.func import functional_call
 
-from fusion_bench.utils.type import TorchModelType, StateDictType
+from fusion_bench.utils.type import StateDictType, TorchModelType
 
 log = logging.getLogger(__name__)
 

fusion_bench/programs/fabric_fusion_program.py

@@ -185,10 +185,13 @@ class FabricModelFusionProgram(
             report = taskpool.evaluate(merged_model)
             return report
         elif isinstance(merged_model, Dict):
-            model = merged_model.pop("model")
-            report: dict = taskpool.evaluate(model)
-            report.update(merged_model)
-            print(report)
+            report = {}
+            for key, item in merged_model.items():
+                if isinstance(item, nn.Module):
+                    report[key] = taskpool.evaluate(item)
+                else:
+                    # metadata
+                    report[key] = item
             return report
         elif isinstance(merged_model, Iterable):
             return [

fusion_bench/taskpool/llama/reward_model.py

@@ -11,10 +11,10 @@ import functools
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast
 
 import lightning as L
+import numpy as np
 import torch
 from omegaconf import DictConfig
 from torch.utils.data import Subset
-import numpy as np
 from tqdm.auto import tqdm
 
 from fusion_bench.dataset.llama.collate import bradley_terry_rm_collate

fusion_bench/tasks/clip_classification/__init__.py

@@ -58,11 +58,24 @@ class CLIPTemplateFactory:
             "templates": "templates",
         },
         "nateraw/rendered-sst2": ".rendered_sst2",
+        "rendered-sst2": ".rendered_sst2",
         "tanganke/stl10": ".stl10",
+        "stl10": ".stl10",
         "dpdl-benchmark/oxford_flowers102": ".flower102",
+        "oxford_flowers102": ".flower102",
         "timm/oxford-iiit-pet": ".oxford_iiit_pet",
+        "oxford-iiit-pet": ".oxford_iiit_pet",
         "imagenet": ".imagenet",
         "tiny-imagenet": ".tiny_imagenet",
+        "pcam": ".pcam",
+        "fer2013": ".fer2013",
+        "emnist_mnist": ".emnist_mnist",
+        "emnist_letters": ".emnist_letters",
+        "kmnist": ".kmnist",
+        "food101": ".food101",
+        "fashion_mnist": ".fashion_mnist",
+        "cub-200-2011": ".cub_200_2011",
+        "mango-leaf-disease": ".mango_leaf_disease",
     }
 
     @staticmethod
@@ -168,48 +181,3 @@ class CLIPTemplateFactory:
 
 def get_classnames_and_templates(dataset_name: str):
     return CLIPTemplateFactory.get_classnames_and_templates(dataset_name)
-
-
-def _load_hf_dataset(dataset_name: str):
-    """
-    Load a dataset from the Hugging Face datasets library based on the specified dataset name.
-
-    This function handles specific preprocessing steps for certain datasets to ensure consistency in dataset format.
-    For example, it renames columns, removes unnecessary columns, and specifies subsets for certain datasets.
-
-    Expected dataset format:
-        - The dataset should have an "image" column containing the image data.
-        - The dataset should have a "label" column containing the class labels.
-
-    Args:
-        dataset_name (str): The name of the dataset to load. Can be one of "svhn", "cifar10", "cifar100", "timm/oxford-iiit-pet", or any other dataset name supported by the Hugging Face datasets library. By default, the datasets have two columns: "image" and "label".
-
-    Returns:
-        A dataset object loaded from the Hugging Face datasets library, with any necessary preprocessing applied.
-    """
-    if dataset_name == "svhn":
-        return load_dataset(dataset_name, "cropped_digits")
-    elif dataset_name == "cifar10":
-        dataset = load_dataset(dataset_name)
-        dataset = dataset.rename_columns({"img": "image"})
-        return dataset
-    elif dataset_name == "cifar100":
-        dataset = load_dataset(dataset_name)
-        dataset = dataset.remove_columns(["coarse_label"]).rename_columns(
-            {"img": "image", "fine_label": "label"}
-        )
-        return dataset
-    elif dataset_name == "timm/oxford-iiit-pet":
-        dataset = load_dataset(dataset_name)
-        dataset = dataset.remove_columns(["image_id", "label_cat_dog"])
-        return dataset
-    else:
-        return load_dataset(dataset_name)
-
-
-def load_clip_dataset(dataset: str, processor):
-    hf_dataset = _load_hf_dataset(dataset)
-    return (
-        CLIPDataset(hf_dataset["train"], processor),
-        CLIPDataset(hf_dataset["test"], processor),
-    )

fusion_bench/tasks/clip_classification/clip_dataset.py

@@ -1,16 +1 @@
-import torch
-
-
-class CLIPDataset(torch.utils.data.Dataset):
-    def __init__(self, dataset, processor):
-        self.dataset = dataset
-        self.processor = processor
-
-    def __len__(self):
-        return len(self.dataset)
-
-    def __getitem__(self, idx):
-        item = self.dataset[idx]
-        image = item["image"]
-        inputs = self.processor(images=[image], return_tensors="pt")["pixel_values"][0]
-        return inputs, item["label"]
+from fusion_bench.dataset.clip_dataset import CLIPDataset

fusion_bench/tasks/clip_classification/cub_200_2011.py

@@ -0,0 +1,208 @@
+classname_mapping = {
+    "0": "Black_footed_Albatross",
+    "1": "Laysan_Albatross",
+    "2": "Sooty_Albatross",
+    "3": "Groove_billed_Ani",
+    "4": "Crested_Auklet",
+    "5": "Least_Auklet",
+    "6": "Parakeet_Auklet",
+    "7": "Rhinoceros_Auklet",
+    "8": "Brewer_Blackbird",
+    "9": "Red_winged_Blackbird",
+    "10": "Rusty_Blackbird",
+    "11": "Yellow_headed_Blackbird",
+    "12": "Bobolink",
+    "13": "Indigo_Bunting",
+    "14": "Lazuli_Bunting",
+    "15": "Painted_Bunting",
+    "16": "Cardinal",
+    "17": "Spotted_Catbird",
+    "18": "Gray_Catbird",
+    "19": "Yellow_breasted_Chat",
+    "20": "Eastern_Towhee",
+    "21": "Chuck_will_Widow",
+    "22": "Brandt_Cormorant",
+    "23": "Red_faced_Cormorant",
+    "24": "Pelagic_Cormorant",
+    "25": "Bronzed_Cowbird",
+    "26": "Shiny_Cowbird",
+    "27": "Brown_Creeper",
+    "28": "American_Crow",
+    "29": "Fish_Crow",
+    "30": "Black_billed_Cuckoo",
+    "31": "Mangrove_Cuckoo",
+    "32": "Yellow_billed_Cuckoo",
+    "33": "Gray_crowned_Rosy_Finch",
+    "34": "Purple_Finch",
+    "35": "Northern_Flicker",
+    "36": "Acadian_Flycatcher",
+    "37": "Great_Crested_Flycatcher",
+    "38": "Least_Flycatcher",
+    "39": "Olive_sided_Flycatcher",
+    "40": "Scissor_tailed_Flycatcher",
+    "41": "Vermilion_Flycatcher",
+    "42": "Yellow_bellied_Flycatcher",
+    "43": "Frigatebird",
+    "44": "Northern_Fulmar",
+    "45": "Gadwall",
+    "46": "American_Goldfinch",
+    "47": "European_Goldfinch",
+    "48": "Boat_tailed_Grackle",
+    "49": "Eared_Grebe",
+    "50": "Horned_Grebe",
+    "51": "Pied_billed_Grebe",
+    "52": "Western_Grebe",
+    "53": "Blue_Grosbeak",
+    "54": "Evening_Grosbeak",
+    "55": "Pine_Grosbeak",
+    "56": "Rose_breasted_Grosbeak",
+    "57": "Pigeon_Guillemot",
+    "58": "California_Gull",
+    "59": "Glaucous_winged_Gull",
+    "60": "Heermann_Gull",
+    "61": "Herring_Gull",
+    "62": "Ivory_Gull",
+    "63": "Ring_billed_Gull",
+    "64": "Slaty_backed_Gull",
+    "65": "Western_Gull",
+    "66": "Anna_Hummingbird",
+    "67": "Ruby_throated_Hummingbird",
+    "68": "Rufous_Hummingbird",
+    "69": "Green_Violetear",
+    "70": "Long_tailed_Jaeger",
+    "71": "Pomarine_Jaeger",
+    "72": "Blue_Jay",
+    "73": "Florida_Jay",
+    "74": "Green_Jay",
+    "75": "Dark_eyed_Junco",
+    "76": "Tropical_Kingbird",
+    "77": "Gray_Kingbird",
+    "78": "Belted_Kingfisher",
+    "79": "Green_Kingfisher",
+    "80": "Pied_Kingfisher",
+    "81": "Ringed_Kingfisher",
+    "82": "White_breasted_Kingfisher",
+    "83": "Red_legged_Kittiwake",
+    "84": "Horned_Lark",
+    "85": "Pacific_Loon",
+    "86": "Mallard",
+    "87": "Western_Meadowlark",
+    "88": "Hooded_Merganser",
+    "89": "Red_breasted_Merganser",
+    "90": "Mockingbird",
+    "91": "Nighthawk",
+    "92": "Clark_Nutcracker",
+    "93": "White_breasted_Nuthatch",
+    "94": "Baltimore_Oriole",
+    "95": "Hooded_Oriole",
+    "96": "Orchard_Oriole",
+    "97": "Scott_Oriole",
+    "98": "Ovenbird",
+    "99": "Brown_Pelican",
+    "100": "White_Pelican",
+    "101": "Western_Wood_Pewee",
+    "102": "Sayornis",
+    "103": "American_Pipit",
+    "104": "Whip_poor_Will",
+    "105": "Horned_Puffin",
+    "106": "Common_Raven",
+    "107": "White_necked_Raven",
+    "108": "American_Redstart",
+    "109": "Geococcyx",
+    "110": "Loggerhead_Shrike",
+    "111": "Great_Grey_Shrike",
+    "112": "Baird_Sparrow",
+    "113": "Black_throated_Sparrow",
+    "114": "Brewer_Sparrow",
+    "115": "Chipping_Sparrow",
+    "116": "Clay_colored_Sparrow",
+    "117": "House_Sparrow",
+    "118": "Field_Sparrow",
+    "119": "Fox_Sparrow",
+    "120": "Grasshopper_Sparrow",
+    "121": "Harris_Sparrow",
+    "122": "Henslow_Sparrow",
+    "123": "Le_Conte_Sparrow",
+    "124": "Lincoln_Sparrow",
+    "125": "Nelson_Sharp_tailed_Sparrow",
+    "126": "Savannah_Sparrow",
+    "127": "Seaside_Sparrow",
+    "128": "Song_Sparrow",
+    "129": "Tree_Sparrow",
+    "130": "Vesper_Sparrow",
+    "131": "White_crowned_Sparrow",
+    "132": "White_throated_Sparrow",
+    "133": "Cape_Glossy_Starling",
+    "134": "Bank_Swallow",
+    "135": "Barn_Swallow",
+    "136": "Cliff_Swallow",
+    "137": "Tree_Swallow",
+    "138": "Scarlet_Tanager",
+    "139": "Summer_Tanager",
+    "140": "Artic_Tern",
+    "141": "Black_Tern",
+    "142": "Caspian_Tern",
+    "143": "Common_Tern",
+    "144": "Elegant_Tern",
+    "145": "Forsters_Tern",
+    "146": "Least_Tern",
+    "147": "Green_tailed_Towhee",
+    "148": "Brown_Thrasher",
+    "149": "Sage_Thrasher",
+    "150": "Black_capped_Vireo",
+    "151": "Blue_headed_Vireo",
+    "152": "Philadelphia_Vireo",
+    "153": "Red_eyed_Vireo",
+    "154": "Warbling_Vireo",
+    "155": "White_eyed_Vireo",
+    "156": "Yellow_throated_Vireo",
+    "157": "Bay_breasted_Warbler",
+    "158": "Black_and_white_Warbler",
+    "159": "Black_throated_Blue_Warbler",
+    "160": "Blue_winged_Warbler",
+    "161": "Canada_Warbler",
+    "162": "Cape_May_Warbler",
+    "163": "Cerulean_Warbler",
+    "164": "Chestnut_sided_Warbler",
+    "165": "Golden_winged_Warbler",
+    "166": "Hooded_Warbler",
+    "167": "Kentucky_Warbler",
+    "168": "Magnolia_Warbler",
+    "169": "Mourning_Warbler",
+    "170": "Myrtle_Warbler",
+    "171": "Nashville_Warbler",
+    "172": "Orange_crowned_Warbler",
+    "173": "Palm_Warbler",
+    "174": "Pine_Warbler",
+    "175": "Prairie_Warbler",
+    "176": "Prothonotary_Warbler",
+    "177": "Swainson_Warbler",
+    "178": "Tennessee_Warbler",
+    "179": "Wilson_Warbler",
+    "180": "Worm_eating_Warbler",
+    "181": "Yellow_Warbler",
+    "182": "Northern_Waterthrush",
+    "183": "Louisiana_Waterthrush",
+    "184": "Bohemian_Waxwing",
+    "185": "Cedar_Waxwing",
+    "186": "American_Three_toed_Woodpecker",
+    "187": "Pileated_Woodpecker",
+    "188": "Red_bellied_Woodpecker",
+    "189": "Red_cockaded_Woodpecker",
+    "190": "Red_headed_Woodpecker",
+    "191": "Downy_Woodpecker",
+    "192": "Bewick_Wren",
+    "193": "Cactus_Wren",
+    "194": "Carolina_Wren",
+    "195": "House_Wren",
+    "196": "Marsh_Wren",
+    "197": "Rock_Wren",
+    "198": "Winter_Wren",
+    "199": "Common_Yellowthroat",
+}
+
+classnames = [classname_mapping[str(i)] for i in range(200)]
+templates = [
+    lambda c: f"a photo of a {c}.",
+    lambda c: f"a photo of the {c}.",
+]

fusion_bench/tasks/clip_classification/emnist_letters.py

@@ -0,0 +1,31 @@
+classnames_mapping = {
+    "0": "A",
+    "1": "B",
+    "2": "C",
+    "3": "D",
+    "4": "E",
+    "5": "F",
+    "6": "G",
+    "7": "H",
+    "8": "I",
+    "9": "J",
+    "10": "K",
+    "11": "L",
+    "12": "M",
+    "13": "N",
+    "14": "O",
+    "15": "P",
+    "16": "Q",
+    "17": "R",
+    "18": "S",
+    "19": "T",
+    "20": "U",
+    "21": "V",
+    "22": "W",
+    "23": "X",
+    "24": "Y",
+    "25": "Z",
+}
+
+classnames = [classnames_mapping[str(i)] for i in range(26)]
+templates = [lambda c: f'a photo of the digit character: "{c}".']

fusion_bench/tasks/clip_classification/emnist_mnist.py

@@ -0,0 +1,5 @@
+# https://huggingface.co/datasets/tanganke/emnist_mnist
+classnames = [str(i) for i in range(10)]
+templates = [
+    lambda c: f'a photo of the number: "{c}".',
+]

fusion_bench/tasks/clip_classification/fashion_mnist.py

@@ -0,0 +1,18 @@
+classname_mapping = {
+    "0": "T - shirt / top",
+    "1": "Trouser",
+    "2": "Pullover",
+    "3": "Dress",
+    "4": "Coat",
+    "5": "Sandal",
+    "6": "Shirt",
+    "7": "Sneaker",
+    "8": "Bag",
+    "9": "Ankle boot",
+}
+classnames = [classname_mapping[str(i)] for i in range(10)]
+
+templates = [
+    lambda c: f"a photo of a {c}.",
+    lambda c: f"a photo of the {c}.",
+]

fusion_bench/tasks/clip_classification/fer2013.py

@@ -0,0 +1,18 @@
+classnames = [
+    "angry",
+    "disgusted",
+    "fearful",
+    "happy",
+    "neutral",
+    "sad",
+    "surprised",
+]
+
+templates = [
+    lambda c: f"a photo of a {c} looking face.",
+    lambda c: f"a photo of a face showing the emotion: {c}.",
+    lambda c: f"a photo of a face looking {c}.",
+    lambda c: f"a face that looks {c}.",
+    lambda c: f"they look {c}.",
+    lambda c: f"look at how {c} they are.",
+]