Trak wrapper

pyproject.toml

@@ -13,11 +13,13 @@ dependencies = [
     "lightning>=1.4.0",
     "torchmetrics>=1.4.0",
     "tqdm>=4.0.0",
+    "traker>=0.3.2"
 ]
 dynamic = ["version"]
 
 [tool.isort]
 profile = "black"
+extend_skip = ["__init__.py"]
 line_length = 79
 multi_line_output = 3
 include_trailing_comma = true
@@ -29,6 +31,10 @@ warn_unused_configs = true
 check_untyped_defs = true
 #ignore_errors = true # TODO: change this
 
+[[tool.mypy.overrides]]
+module = ["trak", "trak.projectors", "fast_jl"]
+ignore_missing_imports = true
+
 # Black formatting
 [tool.black]
 line-length = 127

quanda/explainers/__init__.py

@@ -1,11 +1,11 @@
+from quanda.explainers.base import BaseExplainer
 from quanda.explainers import utils, wrappers
 from quanda.explainers.aggregators import (
     AbsSumAggregator,
     BaseAggregator,
     SumAggregator,
     aggr_types,
 )
-from quanda.explainers.base import BaseExplainer
 from quanda.explainers.functional import ExplainFunc, ExplainFuncMini
 from quanda.explainers.random import RandomExplainer
 

quanda/explainers/wrappers/__init__.py

@@ -11,6 +11,8 @@
     captum_tracincp_self_influence,
 )
 
+from quanda.explainers.wrappers.trak_wrapper import TRAK, trak_explain, trak_self_influence
+
 __all__ = [
     "CaptumInfluence",
     "CaptumSimilarity",
@@ -22,4 +24,7 @@
     "CaptumTracInCP",
     "captum_tracincp_explain",
     "captum_tracincp_self_influence",
+    "TRAK",
+    "trak_explain",
+    "trak_self_influence",
 ]

quanda/explainers/wrappers/captum_influence.py

@@ -11,7 +11,7 @@
     ArnoldiInfluenceFunction,
 )
 
-from quanda.explainers.base import BaseExplainer
+from quanda.explainers import BaseExplainer
 from quanda.explainers.utils import (
     explain_fn_from_explainer,
     self_influence_fn_from_explainer,

quanda/explainers/wrappers/trak_wrapper.py

@@ -0,0 +1,162 @@
+import warnings
+from importlib.util import find_spec
+from typing import Any, Iterable, List, Literal, Optional, Sized, Union
+
+import torch
+from trak import TRAKer
+from trak.projectors import BasicProjector, CudaProjector, NoOpProjector
+
+from quanda.explainers import BaseExplainer
+from quanda.explainers.utils import (
+    explain_fn_from_explainer,
+    self_influence_fn_from_explainer,
+)
+
+TRAKProjectorLiteral = Literal["cuda", "noop", "basic"]
+TRAKProjectionTypeLiteral = Literal["rademacher", "normal"]
+
+
+class TRAK(BaseExplainer):
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        train_dataset: torch.utils.data.Dataset,
+        model_id: str,
+        cache_dir: str,
+        projector: TRAKProjectorLiteral,
+        device: Union[str, torch.device] = "cpu",
+        proj_dim: int = 128,
+        proj_type: TRAKProjectionTypeLiteral = "normal",
+        seed: int = 42,
+        batch_size: int = 32,
+        params_ldr: Optional[Iterable] = None,
+    ):
+        super(TRAK, self).__init__(
+            model=model, train_dataset=train_dataset, model_id=model_id, cache_dir=cache_dir, device=device
+        )
+        self.dataset = train_dataset
+        self.proj_dim = proj_dim
+        self.batch_size = batch_size
+        self.cache_dir = cache_dir if cache_dir is not None else f"./trak_{model_id}_cache"
+
+        num_params_for_grad = 0
+        params_iter = params_ldr if params_ldr is not None else self.model.parameters()
+        for p in list(params_iter):
+            num_params_for_grad = num_params_for_grad + p.numel()
+        # Check if traker was installer with the ["cuda"] option
+        if projector == "cuda":
+            if find_spec("fast_jl"):
+                projector = "cuda"
+            else:
+                warnings.warn("Could not find cuda installation of TRAK. Defaulting to BasicProjector.")
+                projector = "basic"
+
+        projector_cls = {
+            "cuda": CudaProjector,
+            "basic": BasicProjector,
+            "noop": NoOpProjector,
+        }
+
+        projector_kwargs = {
+            "grad_dim": num_params_for_grad,
+            "proj_dim": proj_dim,
+            "proj_type": proj_type,
+            "seed": seed,
+            "device": device,
+        }
+        if projector == "cuda":
+            projector_kwargs["max_batch_size"] = self.batch_size
+        projector_obj = projector_cls[projector](**projector_kwargs)
+        self.traker = TRAKer(
+            model=model,
+            task="image_classification",
+            train_set_size=self.dataset_length,
+            projector=projector_obj,
+            proj_dim=proj_dim,
+            projector_seed=seed,
+            save_dir=self.cache_dir,
+            device=device,
+            use_half_precision=False,
+        )
+
+        # Train the TRAK explainer: featurize the training data
+        ld = torch.utils.data.DataLoader(self.dataset, batch_size=self.batch_size)
+        self.traker.load_checkpoint(self.model.state_dict(), model_id=0)
+        for i, (x, y) in enumerate(iter(ld)):
+            batch = x.to(self.device), y.to(self.device)
+            self.traker.featurize(batch=batch, inds=torch.tensor([i * self.batch_size + j for j in range(x.shape[0])]))
+        self.traker.finalize_features()
+        if projector == "basic":
+            # finalize_features frees memory so projector.proj_matrix needs to be reconstructed
+            self.traker.projector = projector_cls[projector](**projector_kwargs)
+
+    @property
+    def dataset_length(self) -> int:
+        """
+        By default, the Dataset class does not always have a __len__ method.
+        :return:
+        """
+        if isinstance(self.dataset, Sized):
+            return len(self.dataset)
+        dl = torch.utils.data.DataLoader(self.dataset, batch_size=1)
+        return len(dl)
+
+    def explain(self, test, targets):
+        test = test.to(self.device)
+        self.traker.start_scoring_checkpoint(
+            model_id=0, checkpoint=self.model.state_dict(), exp_name="test", num_targets=test.shape[0]
+        )
+        self.traker.score(batch=(test, targets), num_samples=test.shape[0])
+        explanations = torch.from_numpy(self.traker.finalize_scores(exp_name="test")).T.to(self.device)
+
+        # os.remove(os.path.join(self.cache_dir, "scores", "test.mmap"))
+        # os.removedirs(os.path.join(self.cache_dir, "scores"))
+
+        return explanations
+
+
+def trak_explain(
+    model: torch.nn.Module,
+    model_id: str,
+    cache_dir: Optional[str],
+    test_tensor: torch.Tensor,
+    train_dataset: torch.utils.data.Dataset,
+    device: Union[str, torch.device],
+    explanation_targets: Optional[Union[List[int], torch.Tensor]] = None,
+    **kwargs: Any,
+) -> torch.Tensor:
+    return explain_fn_from_explainer(
+        explainer_cls=TRAK,
+        model=model,
+        model_id=model_id,
+        cache_dir=cache_dir,
+        test_tensor=test_tensor,
+        targets=explanation_targets,
+        train_dataset=train_dataset,
+        device=device,
+        **kwargs,
+    )
+
+
+def trak_self_influence(
+    model: torch.nn.Module,
+    model_id: str,
+    cache_dir: Optional[str],
+    train_dataset: torch.utils.data.Dataset,
+    device: Union[str, torch.device],
+    batch_size: Optional[int] = 32,
+    **kwargs: Any,
+) -> torch.Tensor:
+    self_influence_kwargs = {
+        "batch_size": batch_size,
+    }
+    return self_influence_fn_from_explainer(
+        explainer_cls=TRAK,
+        model=model,
+        model_id=model_id,
+        cache_dir=cache_dir,
+        train_dataset=train_dataset,
+        device=device,
+        self_influence_kwargs=self_influence_kwargs,
+        **kwargs,
+    )

quanda/metrics/__init__.py

@@ -1,9 +1,9 @@
+from quanda.metrics.base import GlobalMetric, Metric
 from quanda.metrics import localization, randomization, unnamed
 from quanda.metrics.aggr_strategies import (
     GlobalAggrStrategy,
     GlobalSelfInfluenceStrategy,
 )
-from quanda.metrics.base import GlobalMetric, Metric
 
 __all__ = [
     "Metric",

quanda/metrics/localization/class_detection.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from quanda.metrics.base import Metric
+from quanda.metrics import Metric
 
 
 class ClassDetectionMetric(Metric):

quanda/metrics/localization/mislabeling_detection.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from quanda.metrics.base import GlobalMetric
+from quanda.metrics import GlobalMetric
 
 
 class MislabelingDetectionMetric(GlobalMetric):

quanda/metrics/randomization/model_randomization.py

@@ -3,7 +3,7 @@
 
 import torch
 
-from quanda.metrics.base import Metric
+from quanda.metrics import Metric
 from quanda.utils.common import get_parent_module_from_name
 from quanda.utils.functions import CorrelationFnLiterals, correlation_functions
 

quanda/metrics/unnamed/dataset_cleaning.py

@@ -4,7 +4,7 @@
 import lightning as L
 import torch
 
-from quanda.metrics.base import GlobalMetric
+from quanda.metrics import GlobalMetric
 from quanda.utils.common import class_accuracy
 from quanda.utils.training import BaseTrainer
 

quanda/metrics/unnamed/top_k_overlap.py

@@ -1,6 +1,6 @@
 import torch
 
-from quanda.metrics.base import Metric
+from quanda.metrics import Metric
 
 
 class TopKOverlapMetric(Metric):

quanda/toy_benchmarks/__init__.py

@@ -1,4 +1,4 @@
-from quanda.toy_benchmarks import localization, randomization, unnamed
 from quanda.toy_benchmarks.base import ToyBenchmark
+from quanda.toy_benchmarks import localization, randomization, unnamed
 
 __all__ = ["ToyBenchmark", "randomization", "localization", "unnamed"]

quanda/toy_benchmarks/localization/class_detection.py

@@ -4,7 +4,7 @@
 from tqdm import tqdm
 
 from quanda.metrics.localization import ClassDetectionMetric
-from quanda.toy_benchmarks.base import ToyBenchmark
+from quanda.toy_benchmarks import ToyBenchmark
 
 
 class ClassDetection(ToyBenchmark):

quanda/toy_benchmarks/localization/mislabeling_detection.py

@@ -8,7 +8,7 @@
 from quanda.metrics.localization.mislabeling_detection import (
     MislabelingDetectionMetric,
 )
-from quanda.toy_benchmarks.base import ToyBenchmark
+from quanda.toy_benchmarks import ToyBenchmark
 from quanda.utils.datasets.transformed.label_flipping import (
     LabelFlippingDataset,
 )

quanda/toy_benchmarks/localization/subclass_detection.py

@@ -6,7 +6,7 @@
 from tqdm import tqdm
 
 from quanda.metrics.localization.class_detection import ClassDetectionMetric
-from quanda.toy_benchmarks.base import ToyBenchmark
+from quanda.toy_benchmarks import ToyBenchmark
 from quanda.utils.datasets.transformed.label_grouping import (
     ClassToGroupLiterals,
     LabelGroupingDataset,

quanda/toy_benchmarks/randomization/model_randomization.py

@@ -6,7 +6,7 @@
 from quanda.metrics.randomization.model_randomization import (
     ModelRandomizationMetric,
 )
-from quanda.toy_benchmarks.base import ToyBenchmark
+from quanda.toy_benchmarks import ToyBenchmark
 from quanda.utils.functions import CorrelationFnLiterals
 
 

quanda/toy_benchmarks/unnamed/dataset_cleaning.py

@@ -6,7 +6,7 @@
 from tqdm import tqdm
 
 from quanda.metrics.unnamed.dataset_cleaning import DatasetCleaningMetric
-from quanda.toy_benchmarks.base import ToyBenchmark
+from quanda.toy_benchmarks import ToyBenchmark
 from quanda.utils.training.trainer import BaseTrainer
 
 

quanda/toy_benchmarks/unnamed/top_k_overlap.py

@@ -4,7 +4,7 @@
 from tqdm import tqdm
 
 from quanda.metrics.unnamed import TopKOverlapMetric
-from quanda.toy_benchmarks.base import ToyBenchmark
+from quanda.toy_benchmarks import ToyBenchmark
 
 
 class TopKOverlap(ToyBenchmark):

src/explainers/__init__.py

@@ -0,0 +1,18 @@
+from quanda.explainers.base import BaseExplainer
+from quanda.explainers import utils, wrappers
+from quanda.explainers.functional import ExplainFunc, ExplainFuncMini
+from quanda.explainers.random import RandomExplainer
+from quanda.explainers.aggregators import BaseAggregator, SumAggregator
+
+
+__all__ = [
+    "BaseExplainer",
+    "RandomExplainer",
+    "ExplainFunc",
+    "ExplainFuncMini",
+    "utils",
+    "wrappers",
+    "BaseAggregator",
+    "SumAggregator",
+    "AbsSumAggretor",
+]

src/metrics/__init__.py

@@ -0,0 +1,14 @@
+from quanda.metrics.base import GlobalMetric, Metric
+from quanda.metrics import localization, randomization, unnamed
+from quanda.metrics.aggr_strategies import GlobalAggrStrategy, GlobalSelfInfluenceStrategy
+
+
+__all__ = [
+    "Metric",
+    "GlobalMetric",
+    "GlobalAggrStrategy",
+    "GlobalSelfInfluenceStrategy",
+    "randomization",
+    "localization",
+    "unnamed",
+]

tests/conftest.py

@@ -201,10 +201,20 @@ def load_mnist_test_labels_1():
 
 
 @pytest.fixture
-def load_mnist_explanations_1():
+def load_mnist_explanations_similarity_1():
     return torch.load("tests/assets/mnist_test_suite_1/mnist_SimilarityInfluence_tda.pt")
 
 
+@pytest.fixture
+def load_mnist_explanations_trak_1():
+    return torch.load("tests/assets/mnist_test_suite_1/mnist_TRAK_tda.pt")
+
+
+@pytest.fixture
+def load_mnist_explanations_trak_si_1():
+    return torch.load("tests/assets/mnist_test_suite_1/mnist_TRAK_tda_si.pt")
+
+
 @pytest.fixture
 def load_mnist_dataset_explanations():
     return torch.rand((MINI_BATCH_SIZE, MINI_BATCH_SIZE))