fusion-bench 0.2.20__py3-none-any.whl → 0.2.21__py3-none-any.whl

fusion_bench/method/smile_upscaling/projected_energy.py

@@ -0,0 +1,145 @@
+import os
+from typing import Literal
+
+import pandas as pd
+import torch
+
+from fusion_bench import BaseAlgorithm, BaseModelPool, auto_register_config
+from fusion_bench.mixins import LightningFabricMixin, SimpleProfilerMixin
+
+from tqdm import tqdm
+
+
+class ProjectedEnergyAnalysis(
+    SimpleProfilerMixin,
+    LightningFabricMixin,
+    BaseAlgorithm,
+):
+    def on_run_start(self):
+        self.device = self.fabric.device
+
+    def run(self, modelpool: BaseModelPool):
+        with self.profile("model loading"):
+            base_model = modelpool.load_pretrained_model()
+
+        results = {
+            "model_name": [],
+            "module_index": [],
+            "module_name": [],
+            "projected_energy_I": [],
+            "projected_energy_II": [],
+            "projected_energy_II_III": [],
+        }
+        for model_name in tqdm(
+            modelpool.model_names,
+            "analyzing",
+            dynamic_ncols=True,
+        ):
+            with self.profile("model loading"):
+                finetuned_model = modelpool.load_model(model_name)
+
+            module_index = 0
+            for module_name, base_module in tqdm(
+                list(base_model.named_modules()),
+                "analyzing modules",
+                dynamic_ncols=True,
+            ):
+                if isinstance(base_module, torch.nn.Linear):
+                    with self.profile("weight analysis"):
+                        _result = self.analyze_weight(
+                            base_module.weight,
+                            finetuned_model.get_submodule(module_name).weight,
+                        )
+                    results["model_name"].append(model_name)
+                    results["module_index"].append(module_index)
+                    results["module_name"].append(module_name)
+                    for key, value in _result.items():
+                        results[key].append(value)
+
+                    module_index += 1
+
+        # save results as csv
+        results = pd.DataFrame(results)
+        results.to_csv(
+            os.path.join(self.log_dir, "projected_energy_analysis.csv"), index=True
+        )
+
+        self.print_profile_summary()
+        return None
+
+    @torch.no_grad()
+    def analyze_weight(self, w: torch.Tensor, w_ft: torch.Tensor, k: int = -1):
+        w = w.to(dtype=torch.float32, device=self.device)
+        w_ft = w_ft.to(dtype=torch.float32, device=self.device)
+        w_diff = w_ft - w
+
+        # Perform analysis on the weight tensor
+        u, s, vh = torch.linalg.svd(w, full_matrices=False)
+        v = vh.T
+        if k < 0:
+            # find the position where the sum of singular values is larger than 50% of the total sum
+            cumsum = s.cumsum(0)
+            k = (cumsum < cumsum[-1] * 0.5).sum().item() + 1
+
+        # subspace I
+        w_diff_proj = self._project_subspace_low(u=u, s=s, v=v, k=k, w=w, w_ft=w_ft)
+        projected_energy_I = (
+            torch.linalg.norm(w_diff_proj, ord="fro") ** 2
+            / torch.linalg.norm(w_diff, ord="fro") ** 2
+        )
+
+        # subspace II
+        w_diff_proj = self._project_subspace_high(u=u, s=s, v=v, k=k, w=w, w_ft=w_ft)
+        projected_energy_II = (
+            torch.linalg.norm(w_diff_proj, ord="fro") ** 2
+            / torch.linalg.norm(w_diff, ord="fro") ** 2
+        )
+
+        ## subspace II+III
+        u, s, vh = torch.linalg.svd(w, full_matrices=True)
+        v = vh.T
+        w_diff_proj = self._project_subspace_high(u=u, s=s, v=v, k=k, w=w, w_ft=w_ft)
+        projected_energy_II_III = (
+            torch.linalg.norm(w_diff_proj, ord="fro") ** 2
+            / torch.linalg.norm(w_diff, ord="fro") ** 2
+        )
+
+        return {
+            "projected_energy_I": projected_energy_I.item(),
+            "projected_energy_II": projected_energy_II.item(),
+            "projected_energy_II_III": projected_energy_II_III.item(),
+        }
+
+    def _project_subspace_low(
+        self,
+        u: torch.Tensor,
+        s: torch.Tensor,
+        v: torch.Tensor,
+        k: int,
+        w: torch.Tensor,
+        w_ft: torch.Tensor,
+    ):
+        u = u[:, :k]
+        s = s[:k]
+        v = v[:, :k]
+
+        w_diff = w_ft - w
+        w_diff_proj = torch.linalg.multi_dot((u, u.T, w_diff, v, v.T))
+        return w_diff_proj
+
+    def _project_subspace_high(
+        self,
+        u: torch.Tensor,
+        s: torch.Tensor,
+        v: torch.Tensor,
+        k: int,
+        w: torch.Tensor,
+        w_ft: torch.Tensor,
+    ):
+        u = u[:, k:]
+        s = s[k:]
+        v = v[:, k:]
+
+        w_diff = w_ft - w
+        w_diff_proj = torch.linalg.multi_dot((u, u.T, w_diff, v, v.T))
+        return w_diff_proj

fusion_bench/method/smile_upscaling/smile_qwen2_upscaling.py

@@ -16,10 +16,16 @@ from transformers.models.qwen2.modeling_qwen2 import Qwen2DecoderLayer
 
 from fusion_bench import BaseAlgorithm, BaseModelPool
 from fusion_bench.compat.modelpool import to_modelpool
-from fusion_bench.mixins import SimpleProfilerMixin
+from fusion_bench.mixins import SimpleProfilerMixin, auto_register_config
+from fusion_bench.modelpool import CausalLMPool
+from fusion_bench.models.hf_utils import (
+    generate_complete_readme,
+    save_pretrained_with_remote_code,
+)
 from fusion_bench.models.modeling_smile_qwen2 import (
     SmileQwen2Config,
     SmileQwen2ForCausalLM,
+    SmileQwen2Model,
 )
 from fusion_bench.models.modeling_smile_qwen2.modeling_smile_qwen2 import (
     SmileQwen2DecoderLayer,
@@ -34,6 +40,7 @@ from fusion_bench.utils.parameters import print_parameters
 log = logging.getLogger(__name__)
 
 
+@auto_register_config
 class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
     R"""
     SmileQwen2UpscalingAlgorithm is a model fusion algorithm designed to upscale
@@ -49,15 +56,7 @@ class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
             Merges the pretrained model with the fine-tuned models to create an upscaled model.
     """
 
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "device": "device",
-        "accelerator": "accelerator",
-        "model_path": "model_path",
-        "model_dtype": "model_dtype",
-        "num_experts_per_tok": "num_experts_per_tok",
-        "rank_of_router": "rank_of_router",
-        "rank_of_expert": "rank_of_expert",
-    }
+    modelpool: CausalLMPool
 
     def __init__(
         self,
@@ -68,20 +67,13 @@ class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
         num_experts_per_tok,
         rank_of_router,
         rank_of_expert,
+        save_with_remote_code: bool = True,
         **kwargs,
     ):
-        self.device = device
-        self.accelerator = accelerator
-        self.model_path = model_path
-        self.model_dtype = model_dtype
-        # SmileMoE parameters, except `num_local_experts` which is set later according to the number of finetuned models
-        self.num_experts_per_tok = num_experts_per_tok
-        self.rank_of_router = rank_of_router
-        self.rank_of_expert = rank_of_expert
         super().__init__(**kwargs)
 
     @torch.no_grad()
-    def run(self, modelpool: BaseModelPool) -> SmileQwen2ForCausalLM:
+    def run(self, modelpool) -> SmileQwen2ForCausalLM:
         """
         Executes the upscaling process.
 
@@ -129,13 +121,29 @@ class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
             if os.path.dirname(config.model_path):
                 os.makedirs(os.path.dirname(config.model_path), exist_ok=True)
             log.info(f"Saving model to {config.model_path}")
-            pretrained_model_config = self.modelpool.get_model_config("_pretrained_")
-            pretrained_path = pretrained_model_config.get(
-                "path", pretrained_model_config["pretrained_model_name_or_path"]
-            )
-            tokenizer = AutoTokenizer.from_pretrained(pretrained_path)
+            tokenizer = self.modelpool.load_tokenizer()
             tokenizer.save_pretrained(config.model_path)
-            model.save_pretrained(config.model_path)
+            if not self.save_with_remote_code:
+                model.save_pretrained(config.model_path)
+            else:
+                save_pretrained_with_remote_code(
+                    model,
+                    auto_map={
+                        "AutoConfig": SmileQwen2Config,
+                        "AutoModel": SmileQwen2Model,
+                        "AutoModelForCausalLM": SmileQwen2ForCausalLM,
+                    },
+                    save_directory=config.model_path,
+                )
+
+            # save readme
+            complete_readme = generate_complete_readme(
+                algorithm=self,
+                modelpool=modelpool,
+                models=[modelpool.get_model_path(m) for m in modelpool.all_model_names],
+            )
+            with open(os.path.join(config.model_path, "README.md"), "w") as f:
+                f.write(complete_readme)
 
         return model
 
@@ -158,9 +166,12 @@ class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
 
         with init_empty_weights():
             pretrained_model_config = self.modelpool.get_model_config("_pretrained_")
-            pretrained_path = pretrained_model_config.get(
-                "path", pretrained_model_config["pretrained_model_name_or_path"]
-            )
+            if isinstance(pretrained_model_config, str):
+                pretrained_path = pretrained_model_config
+            else:
+                pretrained_path = pretrained_model_config.get(
+                    "path", pretrained_model_config["pretrained_model_name_or_path"]
+                )
             base_config = AutoConfig.from_pretrained(pretrained_path)
             model_config = SmileQwen2Config(
                 num_experts_per_tok=config.num_experts_per_tok,

fusion_bench/method/smile_upscaling/smile_upscaling.py

@@ -1,7 +1,7 @@
 import logging
 import os
 from copy import deepcopy
-from typing import Dict, List, Tuple  # noqa: F401
+from typing import Any, Dict, List, Tuple  # noqa: F401
 
 import torch
 import torch.nn.functional as F
@@ -21,6 +21,7 @@ from fusion_bench.models.smile_moe.linear_from_module import (
 )
 from fusion_bench.models.utils import get_attr, set_attr
 from fusion_bench.utils.parameters import print_parameters
+from fusion_bench.utils.devices import get_device
 
 log = logging.getLogger(__name__)
 
@@ -54,7 +55,7 @@ class SmileUpscalingAlgorithm(
         routing_use_diff: bool = True,
         average_experts: bool = False,
         model_path: str = None,
-        **kwargs,
+        **kwargs: Any,
     ):
         """
         Initialize the SmileUpscalingAlgorithm.
@@ -91,7 +92,7 @@ class SmileUpscalingAlgorithm(
         print(f"=== Config for `{type(self).__name__}` ===")
 
     @torch.no_grad()
-    def run(self, modelpool: BaseModelPool):
+    def run(self, modelpool: BaseModelPool) -> nn.Module:
         """
         Executes the upscaling process.
 
@@ -142,7 +143,7 @@ class SmileUpscalingAlgorithm(
         pretrained_model: nn.Module,
         finetuned_models: List[nn.Module],
         in_place: bool = True,
-    ):
+    ) -> nn.Module:
         """
         Merges the pretrained model with the fine-tuned models to create an upscaled model.
 
@@ -180,7 +181,12 @@ class SmileUpscalingAlgorithm(
 
         name_list = name.split(".")
         module = get_attr(pretrained_model, name_list)
-        experts = [get_attr(m, name_list) for m in finetuned_models]
+        original_device = get_device(module)
+        module = module.to(self.device, non_blocking=True)
+        experts = [
+            get_attr(m, name_list).to(self.device, non_blocking=True)
+            for m in finetuned_models
+        ]
         try:
             moe_linear = SmileMoELinear(
                 module,
@@ -192,6 +198,7 @@ class SmileUpscalingAlgorithm(
                 full_matrices=self.full_matrices,
                 upscaling_accelerator=self.upscaling_accelerator,
             )
+            moe_linear = moe_linear.to(original_device, non_blocking=True)
         except ExpertNotTrainedError:
             print(f"skip {name} because the experts are not trained.")
             return

fusion_bench/method/tall_mask/task_arithmetic.py

@@ -9,7 +9,7 @@ from copy import deepcopy
 import torch
 
 from fusion_bench import BaseAlgorithm
-from fusion_bench.mixins import SimpleProfilerMixin
+from fusion_bench.mixins import SimpleProfilerMixin, auto_register_config
 from fusion_bench.modelpool import BaseModelPool
 from fusion_bench.utils.state_dict_arithmetic import (
     state_dict_add,
@@ -58,16 +58,11 @@ def generate_task_masks(
     return final_mask
 
 
+@auto_register_config
 class TallMaskTaskArithmeticAlgorithm(
-    BaseAlgorithm,
     SimpleProfilerMixin,
+    BaseAlgorithm,
 ):
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "tall_mask_lambda": "tall_mask_lambda",
-        "debug": "debug",
-        "verbose": "verbose",
-    }
-
     def __init__(
         self,
         tall_mask_lambda: float,
@@ -76,9 +71,6 @@ class TallMaskTaskArithmeticAlgorithm(
         **kwargs,
     ):
         super().__init__(**kwargs)
-        self.tall_mask_lambda = tall_mask_lambda
-        self.debug = debug
-        self.verbose = verbose
 
     @torch.no_grad()
     def run(self, modelpool: BaseModelPool):

fusion_bench/method/task_arithmetic/task_arithmetic.py

@@ -12,7 +12,7 @@ import torch
 from torch import nn
 
 from fusion_bench.method.base_algorithm import BaseAlgorithm
-from fusion_bench.mixins.simple_profiler import SimpleProfilerMixin
+from fusion_bench.mixins import SimpleProfilerMixin, auto_register_config
 from fusion_bench.modelpool import BaseModelPool
 from fusion_bench.utils.state_dict_arithmetic import (
     state_dict_add,
@@ -74,9 +74,10 @@ def task_arithmetic_merge(
     return pretrained_model
 
 
+@auto_register_config
 class TaskArithmeticAlgorithm(
-    BaseAlgorithm,
     SimpleProfilerMixin,
+    BaseAlgorithm,
 ):
     """
     Task Arithmetic Algorithm for model fusion.
@@ -89,22 +90,17 @@ class TaskArithmeticAlgorithm(
         scaling_factor (int): The factor by which the task vectors will be scaled before merging.
     """
 
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "scaling_factor": "scaling_factor"
-    }
-
-    def __init__(self, scaling_factor: int):
+    def __init__(self, scaling_factor: int, **kwargs):
         """
         Initializes the TaskArithmeticAlgorithm with the given scaling factor.
 
         Args:
             scaling_factor (int): The factor by which the task vectors will be scaled before merging.
         """
-        self.scaling_factor = scaling_factor
-        super().__init__()
+        super().__init__(**kwargs)
 
     @torch.no_grad()
-    def run(self, modelpool: Union[BaseModelPool, Dict[str, nn.Module]]):
+    def run(self, modelpool: Union[BaseModelPool, Dict[str, nn.Module]]) -> nn.Module:
         """
         Runs the Task Arithmetic Algorithm to fuse models in the given model pool.
 

fusion_bench/method/ties_merging/ties_merging.py

@@ -9,14 +9,14 @@ Overview of Ties-Merging:
 """
 
 import logging
-from typing import Dict, List, Literal, Mapping, Union  # noqa: F401
+from typing import Any, Dict, List, Literal, Mapping, Union  # noqa: F401
 
 import torch
 from torch import Tensor, nn
 
 from fusion_bench.compat.modelpool import to_modelpool
 from fusion_bench.method import BaseAlgorithm
-from fusion_bench.mixins import SimpleProfilerMixin
+from fusion_bench.mixins import SimpleProfilerMixin, auto_register_config
 from fusion_bench.modelpool import BaseModelPool
 from fusion_bench.utils.type import StateDictType
 
@@ -25,33 +25,22 @@ from .ties_merging_utils import state_dict_to_vector, ties_merging, vector_to_st
 log = logging.getLogger(__name__)
 
 
-class TiesMergingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
-    """
-    TiesMergingAlgorithm is a class for fusing multiple models using the TIES merging technique.
-
-    Attributes:
-        scaling_factor (float): The scaling factor to apply to the merged task vector.
-        threshold (float): The threshold for resetting values in the task vector.
-        remove_keys (List[str]): List of keys to remove from the state dictionary.
-        merge_func (Literal["sum", "mean", "max"]): The merge function to use for disjoint merging.
-    """
-
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "scaling_factor": "scaling_factor",
-        "threshold": "threshold",
-        "remove_keys": "remove_keys",
-        "merge_func": "merge_func",
-    }
-
+@auto_register_config
+class TiesMergingAlgorithm(
+    SimpleProfilerMixin,
+    BaseAlgorithm,
+):
     def __init__(
         self,
         scaling_factor: float,
         threshold: float,
         remove_keys: List[str],
         merge_func: Literal["sum", "mean", "max"],
-        **kwargs,
+        **kwargs: Any,
     ):
         """
+        TiesMergingAlgorithm is a class for fusing multiple models using the TIES merging technique.
+
         Initialize the TiesMergingAlgorithm with the given parameters.
 
         Args:
@@ -61,14 +50,12 @@ class TiesMergingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
             merge_func (Literal["sum", "mean", "max"]): The merge function to use for disjoint merging.
             **kwargs: Additional keyword arguments for the base class.
         """
-        self.scaling_factor = scaling_factor
-        self.threshold = threshold
-        self.remove_keys = remove_keys
-        self.merge_func = merge_func
         super().__init__(**kwargs)
 
     @torch.no_grad()
-    def run(self, modelpool: BaseModelPool | Dict[str, nn.Module], **kwargs):
+    def run(
+        self, modelpool: BaseModelPool | Dict[str, nn.Module], **kwargs: Any
+    ) -> nn.Module:
         """
         Run the TIES merging algorithm to fuse models in the model pool.
 

fusion_bench/method/we_moe/clip_we_moe.py

@@ -2,6 +2,7 @@ import functools
 import logging
 import os
 from copy import deepcopy
+from typing import Any, Iterator
 
 import torch
 from torch import Tensor
@@ -38,7 +39,7 @@ class CLIPWeightEnsemblingMoEAlgorithm(
 
     modelpool: CLIPVisionModelPool = None
 
-    def load_checkpoint(self, model, checkpoint):
+    def load_checkpoint(self, model: Any, checkpoint: Any):
         """
         Load the checkpoint file.
 
@@ -49,7 +50,7 @@ class CLIPWeightEnsemblingMoEAlgorithm(
         state = {"model": model}
         self._fabric.load(checkpoint, state)
 
-    def save_checkpoint(self, model, checkpoint):
+    def save_checkpoint(self, model: Any, checkpoint: Any):
         """
         Save the checkpoint file.
 
@@ -102,7 +103,7 @@ class CLIPWeightEnsemblingMoEAlgorithm(
         return moe_model
 
     @functools.cache
-    def get_shuffled_test_loader_iter(self, tta_dataset: str):
+    def get_shuffled_test_loader_iter(self, tta_dataset: str) -> Iterator:
         """
         Get an iterator for the shuffled test data loader.
 
@@ -131,7 +132,7 @@ class CLIPWeightEnsemblingMoEAlgorithm(
         """
         self.setup_zero_shot_classification_head()
 
-    def compute_logits(self, module, batch, task) -> Tensor:
+    def compute_logits(self, module: Any, batch: Any, task: Any) -> Tensor:
         """
         Compute the logits for the given batch and task.
 

fusion_bench/method/we_moe/we_moe.py

@@ -1,6 +1,6 @@
 import logging
 from abc import abstractmethod
-from typing import cast  # noqa: F401
+from typing import Any, cast  # noqa: F401
 
 import lightning as L
 import lightning.fabric.wrappers
@@ -70,7 +70,7 @@ class WeightEnsemblingMoEAlgorithm(
             assert "No CUDA device available."
 
     @abstractmethod
-    def load_checkpoint(self, model, checkpoint):
+    def load_checkpoint(self, model: Any, checkpoint: Any):
         """
         Load the checkpoint file.
 
@@ -81,7 +81,7 @@ class WeightEnsemblingMoEAlgorithm(
         pass
 
     @abstractmethod
-    def save_checkpoint(self, model, checkpoint):
+    def save_checkpoint(self, model: Any, checkpoint: Any):
         """
         Save the checkpoint file.
 
@@ -121,7 +121,7 @@ class WeightEnsemblingMoEAlgorithm(
         pass
 
     @abstractmethod
-    def compute_logits(self, module, batch, task) -> Tensor:
+    def compute_logits(self, module: Any, batch: Any, task: Any) -> Tensor:
         """
         Compute the logits for a given batch and task.
 
@@ -135,7 +135,7 @@ class WeightEnsemblingMoEAlgorithm(
         """
         pass
 
-    def test_time_adaptation(self, module: WeightEnsemblingMoE):
+    def test_time_adaptation(self, module: WeightEnsemblingMoE) -> WeightEnsemblingMoE:
         """
         Perform test-time adaptation for the given module.
 
@@ -208,7 +208,7 @@ class WeightEnsemblingMoEAlgorithm(
 
         return module
 
-    def run(self, modelpool: ModelPool):
+    def run(self, modelpool: ModelPool) -> WeightEnsemblingMoE:
         """
         Run the WeightEnsemblingMoEAlgorithm to fuse models using Weight Ensembling Mixture of Experts.
 

fusion_bench/method/weighted_average/llama.py

@@ -3,6 +3,7 @@ from typing import List, Mapping, Union  # noqa: F401
 
 import numpy as np
 import torch
+from transformers import PreTrainedModel
 from typing_extensions import override
 
 from fusion_bench.method import BaseAlgorithm
@@ -10,24 +11,17 @@ from fusion_bench.modelpool import CausalLMPool
 from fusion_bench.utils import timeit_context
 from fusion_bench.utils.state_dict_arithmetic import state_dict_add, state_dict_mul
 from fusion_bench.utils.type import StateDictType
+from fusion_bench.mixins import auto_register_config
 
 log = logging.getLogger(__name__)
 
 
+@auto_register_config
 class WeightedAverageForLLama(BaseAlgorithm):
     """
     A class to perform weighted averaging of LlaMa/Mistral models.
     """
 
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "normalize": "normalize",
-        "weights": "weights",
-        "backbone_only": "backbone_only",
-        "merged_model_save_path": "merged_model_save_path",
-        "save_tokenizer": "save_tokenizer",
-        "push_to_hub": "push_to_hub",
-    }
-
     def __init__(
         self,
         normalize: bool,
@@ -49,17 +43,11 @@ class WeightedAverageForLLama(BaseAlgorithm):
             save_tokenizer (bool): Whether to save the tokenizer.
             push_to_hub (bool): Whether to push the model to the hub.
         """
-        self.normalize = normalize
-        self.weights = weights
-        self.backbone_only = backbone_only
-        self.merged_model_save_path = merged_model_save_path
-        self.save_tokenizer = save_tokenizer
-        self.push_to_hub = push_to_hub
         super().__init__(**kwargs)
 
     @override
     @torch.no_grad()
-    def run(self, modelpool: CausalLMPool):
+    def run(self, modelpool: CausalLMPool) -> PreTrainedModel:
         """
         Executes the weighted averaging of models in the provided model pool.
 

fusion_bench/metrics/continual_learning/__init__.py

@@ -0,0 +1 @@
+from .backward_transfer import compute_backward_transfer

fusion_bench/metrics/continual_learning/backward_transfer.py

@@ -10,7 +10,7 @@ def compute_backward_transfer(
     Compute the backward transfer (BWT) of a model on a set of tasks.
 
     Equation:
-        BWT = \frac{1}{n} \sum_{k=1}^{n} (acc_{Ti}[k] - acc_{ii}[k])
+        $BWT = \frac{1}{n} \sum_{k=1}^{n} (acc_{T,i}[k] - acc_{i,i}[k])$
 
     Returns:
         float: The backward transfer of the model.

fusion_bench/metrics/nyuv2/__init__.py

@@ -1,10 +1,10 @@
 from .depth import DepthMetric
 from .noise import NoiseMetric
 from .normal import NormalMetric
-from .segmentation import SegmentationMertic
+from .segmentation import SegmentationMetric
 
 metric_classes = {
-    "segmentation": SegmentationMertic,
+    "segmentation": SegmentationMetric,
     "depth": DepthMetric,
     "normal": NormalMetric,
     "noise": NoiseMetric,

fusion_bench/metrics/nyuv2/segmentation.py

@@ -5,7 +5,7 @@ from torch import Tensor, nn
 from torchmetrics import Metric
 
 
-class SegmentationMertic(Metric):
+class SegmentationMetric(Metric):
     metric_names = ["mIoU", "pixAcc"]
 
     def __init__(self, num_classes=13):