WiP: Globalization from local explainer

src/metrics/localization/identical_class.py

@@ -15,11 +15,7 @@ def __init__(
         super().__init__(model, train_dataset, device, *args, **kwargs)
         self.scores = []
 
-    def update(
-        self,
-        test_labels: torch.Tensor,
-        explanations: torch.Tensor
-    ):
+    def update(self, test_labels: torch.Tensor, explanations: torch.Tensor):
         """
         Used to implement metric-specific logic.
         """

src/metrics/randomization/model_randomization.py

@@ -79,9 +79,7 @@ def update(
         corrs = self.correlation_measure(explanations, rand_explanations)
         self.results["rank_correlations"].append(corrs)
 
-    def compute(
-        self,
-    ):
+    def compute(self):
         return torch.cat(self.results["rank_correlations"]).mean()
 
     def reset(self):

src/metrics/unnamed/top_k_overlap.py

@@ -29,7 +29,7 @@ def compute(self, *args, **kwargs):
         return len(torch.unique(self.all_top_k_examples))
 
     def reset(self, *args, **kwargs):
-        self.all_top_k_examples = []
+        self.all_top_k_examples = torch.empty(0, self.top_k)
 
     def load_state_dict(self, state_dict: dict, *args, **kwargs):
         self.all_top_k_examples = state_dict["all_top_k_examples"]

src/utils/aggregators.py

@@ -0,0 +1,24 @@
+from abc import ABC
+
+import torch
+
+
+class ExplanationsAggregator(ABC):
+    def __init__(self, training_size: int, *args, **kwargs):
+        self.scores = torch.zeros(training_size)
+
+    def update(self, explanations: torch.Tensor):
+        raise NotImplementedError
+
+    def get_global_ranking(self) -> torch.Tensor:
+        return self.scores.argsort()
+
+
+class SumAggregator(ExplanationsAggregator):
+    def update(self, explanations: torch.Tensor) -> torch.Tensor:
+        self.scores += explanations.sum(dim=0)
+
+
+class AbsSumAggregator(ExplanationsAggregator):
+    def update(self, explanations: torch.Tensor) -> torch.Tensor:
+        self.scores += explanations.abs().sum(dim=0)

src/utils/common.py

@@ -1,8 +1,12 @@
 import functools
 from functools import reduce
-from typing import Any, Callable, Mapping
+from typing import Any, Callable, Mapping, Optional
 
 import torch
+import torch.utils
+import torch.utils.data
+
+from utils.explain_wrapper import SelfInfluenceFunction
 
 
 def _get_module_from_name(model: torch.nn.Module, layer_name: str) -> Any:
@@ -22,3 +26,18 @@ def make_func(func: Callable, func_kwargs: Mapping[str, ...] | None, **kwargs) -
         func_kwargs = kwargs
 
     return functools.partial(func, **func_kwargs)
+
+
+def get_self_influence_ranking(
+    model: torch.nn.Module,
+    model_id: str,
+    cache_dir: Optional[str],
+    training_data: torch.utils.data.Dataset,
+    self_influence_fn: SelfInfluenceFunction,
+    self_influence_fn_kwargs: Optional[dict] = None,
+) -> torch.Tensor:
+    size = len(training_data)
+    self_inf = torch.zeros((size,))
+    for i, (x, y) in enumerate(training_data):
+        self_inf[i] = self_influence_fn(model, model_id, cache_dir, training_data, i, **self_influence_fn_kwargs)
+    return self_inf.argsort()

src/utils/explain_wrapper.py

@@ -16,7 +16,19 @@ def __call__(
         test_tensor: torch.Tensor,
         method: str,
     ) -> torch.Tensor:
-        ...
+        pass
+
+
+class SelfInfluenceFunction(Protocol):
+    def __call__(
+        self,
+        model: torch.nn.Module,
+        model_id: str,
+        cache_dir: Optional[str],
+        train_dataset: torch.utils.data.Dataset,
+        id: int,
+    ) -> torch.Tensor:
+        pass
 
 
 def explain(