WiP: Datasets reworked

src/downstream_tasks/subclass_identification.py

@@ -7,7 +7,7 @@
 from src.explainers.functional import ExplainFunc
 from src.explainers.wrappers.captum_influence import captum_similarity_explain
 from src.metrics.localization.identical_class import IdenticalClass
-from src.utils.datasets.group_label_dataset import (
+from src.utils.datasets.transformed_datasets.label_grouping import (
     ClassToGroupLiterals,
     GroupLabelDataset,
 )

src/metrics/randomization/model_randomization.py

@@ -27,7 +27,7 @@ def __init__(
         seed: int = 42,
         model_id: str = "0",
         cache_dir: str = "./cache",
-        device: str = "cpu" if torch.cuda.is_available() else "cuda",
+        device: str = "cpu",
         *args,
         **kwargs,
     ):
@@ -101,15 +101,15 @@ def explain_update(
         corrs = self.corr_measure(explanations, rand_explanations)
         self.results["scores"].append(corrs)
 
-    def compute(self):
+    def compute(self) -> torch.Tensor:
         return torch.cat(self.results["scores"]).mean()
 
     def reset(self):
         self.results = {"scores": []}
         self.generator.manual_seed(self.seed)
         self.rand_model = self._randomize_model(self.model)
 
-    def state_dict(self):
+    def state_dict(self) -> Dict:
         state_dict = {
             "results_dict": self.results,
             "rnd_model": self.model.state_dict(),
@@ -129,7 +129,7 @@ def load_state_dict(self, state_dict: dict):
         # self.explain_fn = state_dict["explain_fn"]
         # self.generator.set_state(state_dict["generator_state"])
 
-    def _randomize_model(self, model: torch.nn.Module):
+    def _randomize_model(self, model: torch.nn.Module) -> torch.nn.Module:
         rand_model = copy.deepcopy(model)
         for name, param in list(rand_model.named_parameters()):
             random_param_tensor = torch.empty_like(param).normal_(generator=self.generator)

src/utils/datasets/transformed_datasets/base.py

@@ -0,0 +1,68 @@
+import random
+from typing import Any, Callable, Optional, Sized
+
+import torch
+from torch.utils.data.dataset import Dataset
+
+
+class TransformedDataset(Dataset):
+    @staticmethod
+    def _identity(x: Any) -> Any:
+        return x
+
+    def __init__(
+        self,
+        dataset: torch.utils.data.Dataset,
+        n_classes: int,
+        cache_path: str = "./cache",
+        cls_idx: Optional[int] = None,
+        # If isinstance(subset_idx,int): perturb this class with probability p,
+        # if isinstance(subset_idx,List[int]): perturb datapoints with these indices with probability p
+        p: float = 1.0,
+        seed: Optional[int] = None,
+        device: str = "cpu",
+        sample_fn: Optional[Callable] = None,
+        label_fn: Optional[Callable] = None,
+    ):
+        super().__init__()
+        self.dataset = dataset
+        self.n_classes = n_classes
+        self.cls_idx = cls_idx
+        self.cache_path = cache_path
+        self.p = p
+        if sample_fn is not None:
+            self.sample_fn = sample_fn
+        else:
+            self.sample_fn = TransformedDataset._identity
+        if label_fn is not None:
+            self.label_fn = label_fn
+        else:
+            self.label_fn = TransformedDataset._identity
+        self.seed = seed
+        if self.seed is not None:
+            random.seed(self.seed)
+        self.samples_to_perturb = []
+        for i in range(self.__len__()):
+            x, y = dataset[i]
+            perturb_sample = (self.cls_idx is None) or (y == self.cls_idx)
+            p_condition = (random.random() <= self.p) if self.p < 1.0 else True
+            perturb_sample = p_condition and perturb_sample
+            if perturb_sample:
+                self.samples_to_perturb.append(i)
+
+    def __len__(self):
+        if isinstance(self.dataset, Sized):
+            return len(self.dataset)
+        dl = torch.utils.data.DataLoader(self.dataset, batch_size=1)
+        return len(dl)
+
+    def __getitem__(self, index):
+        x, y = self.dataset[index]
+        xx = self.sample_fn(x)
+        yy = self.label_fn(y)
+
+        return xx, yy if index in self.samples_to_perturb else x, y
+
+    def _get_original_label(self, index):
+        _, y = self.dataset[index]
+        return y