Explainer base

Makefile

@@ -5,7 +5,7 @@ SHELL = /bin/bash
 .PHONY: style
 style:
 	black .
-	flake8 .
+	flake8 . --pytest-parametrize-names-type=csv
 	python -m isort .
 	rm -f .coverage
 	rm -f .coverage.*

pyproject.toml

@@ -56,6 +56,7 @@ dev = [ # Install wtih pip install .[dev] or pip install -e '.[dev]' in zsh
     "coverage>=7.2.3",
     "flake8>=6.0.0",
     "pytest<=7.4.4",
+    "flake8-pytest-style>=1.3.2",
     "pytest-cov>=4.0.0",
     "pytest-mock==3.10.0",
     "pre-commit>=3.2.0",

src/downstream_tasks/subclass_identification.py

@@ -3,12 +3,13 @@
 
 import torch
 
+from explainers.functional import ExplainFunc
+from explainers.wrappers.captum_influence import captum_similarity_explain
 from metrics.localization.identical_class import IdenticalClass
 from utils.datasets.group_label_dataset import (
     ClassToGroupLiterals,
     GroupLabelDataset,
 )
-from utils.explain_wrapper import explain
 from utils.training.trainer import BaseTrainer, Trainer
 
 
@@ -41,7 +42,7 @@ def evaluate(
         n_classes: int = 10,
         n_groups: int = 2,
         class_to_group: Union[ClassToGroupLiterals, Dict[int, int]] = "random",
-        explain_fn: Callable = explain,
+        explain_fn: ExplainFunc = captum_similarity_explain,
         explain_kwargs: Optional[dict] = None,
         trainer_kwargs: Optional[dict] = None,
         cache_dir: str = "./cache",
@@ -99,7 +100,8 @@ def evaluate(
                 cache_dir=os.path.join(cache_dir, run_id),
                 train_dataset=train_dataset,
                 test_tensor=input,
-                **explain_kwargs,
+                init_kwargs=explain_kwargs,
+                device=device,
             )
             metric.update(labels, explanations)
 

src/explainers/aggregators.py

@@ -0,0 +1,53 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+import torch
+
+
+class BaseAggregator(ABC):
+    def __init__(self):
+        self.scores: Optional[torch.Tensor] = None
+
+    @abstractmethod
+    def update(self, explanations: torch.Tensor):
+        raise NotImplementedError
+
+    def _validate_explanations(self, explanations: torch.Tensor):
+        if self.scores is None:
+            self.scores = torch.zeros(explanations.shape[1])
+
+        if explanations.shape[1] != self.scores.shape[0]:
+            raise ValueError(f"Explanations shape {explanations.shape} does not match the expected shape {self.scores.shape}")
+
+    def reset(self, *args, **kwargs):
+        """
+        Used to reset the aggregator state.
+        """
+        self.scores: torch.Tensor = None
+
+    def load_state_dict(self, state_dict: dict, *args, **kwargs):
+        """
+        Used to load the aggregator state.
+        """
+        self.scores = state_dict["scores"]
+
+    def state_dict(self, *args, **kwargs):
+        """
+        Used to return the metric state.
+        """
+        return {"scores": self.scores}
+
+    def compute(self) -> torch.Tensor:
+        return self.scores.argsort()
+
+
+class SumAggregator(BaseAggregator):
+    def update(self, explanations: torch.Tensor):
+        self._validate_explanations(explanations)
+        self.scores += explanations.sum(dim=0)
+
+
+class AbsSumAggregator(BaseAggregator):
+    def update(self, explanations: torch.Tensor):
+        self._validate_explanations(explanations)
+        self.scores += explanations.abs().sum(dim=0)

src/explainers/base.py

@@ -1,30 +1,49 @@
 from abc import ABC, abstractmethod
-from typing import Union
+from typing import Any, List, Optional, Union
 
 import torch
 
+from src.utils.common import cache_result
 
-class Explainer(ABC):
+
+class BaseExplainer(ABC):
     def __init__(
         self,
         model: torch.nn.Module,
-        dataset: torch.data.utils.Dataset,
+        model_id: str,
+        cache_dir: Optional[str],
+        train_dataset: torch.utils.data.Dataset,
         device: Union[str, torch.device],
+        **kwargs,
     ):
         self.model = model
+        self.model.to(device)
+
+        self.model_id = model_id
+        self.cache_dir = cache_dir
+        self.train_dataset = train_dataset
         self.device = torch.device(device) if isinstance(device, str) else device
-        self.images = dataset
-        self.samples = []
-        self.labels = []
-        dev = torch.device(device)
-        self.model.to(dev)
 
     @abstractmethod
-    def explain(self, x: torch.Tensor, explanation_targets: torch.Tensor) -> torch.Tensor:
-        pass
+    def explain(self, test: torch.Tensor, targets: Optional[Union[List[int], torch.Tensor]] = None, **kwargs: Any):
+        raise NotImplementedError
+
+    @cache_result
+    def self_influence(self, batch_size: Optional[int] = 32, **kwargs: Any) -> torch.Tensor:
+        """
+        Base class implements computing self influences by explaining the train dataset one by one
+
+        :param batch_size:
+        :param kwargs:
+        :return:
+        """
+
+        # Pre-allcate memory for influences, because torch.cat is slow
+        influences = torch.empty((len(self.train_dataset),), device=self.device)
+        ldr = torch.utils.data.DataLoader(self.train_dataset, shuffle=False, batch_size=batch_size)
 
-    def train(self) -> None:
-        pass
+        for i, (x, y) in zip(range(0, len(self.train_dataset), batch_size), ldr):
+            explanations = self.explain(test=x.to(self.device), **kwargs)
+            influences[i : i + batch_size] = explanations.diag(diagonal=i)
 
-    def save_coefs(self, dir_path: str) -> None:
-        pass
+        return influences

src/explainers/functional.py

@@ -0,0 +1,19 @@
+from typing import Dict, List, Optional, Protocol, Union
+
+import torch
+
+
+class ExplainFunc(Protocol):
+    def __call__(
+        self,
+        model: torch.nn.Module,
+        model_id: str,
+        cache_dir: Optional[str],
+        test_tensor: torch.Tensor,
+        explanation_targets: Optional[Union[List[int], torch.Tensor]],
+        train_dataset: torch.utils.data.Dataset,
+        explain_kwargs: Dict,
+        init_kwargs: Dict,
+        device: Union[str, torch.device],
+    ) -> torch.Tensor:
+        pass