fit reducers

t-vi · t-vi · commit a8d9171c6b47 · 2021-02-19T11:28:47.000+01:00
diff --git a/notebooks/drift_detection_on_images.ipynb b/notebooks/drift_detection_on_images.ipynb
diff --git a/torchdrift/detectors/__init__.py b/torchdrift/detectors/__init__.py
@@ -1,3 +1,3 @@
-from .detector import DriftDetector
+from .detector import Detector
 from .mmd import kernel_mmd, KernelMMDDriftDetector
 from .ks import ks_two_sample_multi_dim, KSDriftDetector, ks_p_value
diff --git a/torchdrift/detectors/detector.py b/torchdrift/detectors/detector.py
@@ -1,9 +1,6 @@
-from typing import Optional, Callable
-import tqdm
-
 import torch
 
-class DriftDetector(torch.nn.Module):
+class Detector(torch.nn.Module):
     """Detector class.
 
 The detector is is a `nn.Module` subclass that, after fitting, performs a drift test when called and returns a score or p-value.
@@ -16,28 +13,10 @@ def __init__(self, *, return_p_value: bool=False):
         self.register_buffer('base_outputs', None)
         self.return_p_value = return_p_value
 
-    def fit(
-        self,
-        ref_ds: torch.utils.data.Dataset,
-        feature_extractor: torch.nn.Module,
-        batch_size: int = 32,
-        num_batches: Optional[int] = None,
-    ):
-        """Train drift detector on reference distribution.
-"""
-        
-        feature_extractor.eval()  # careful about test time dropout
-        device = next(feature_extractor.parameters()).device
-        all_outputs = []
-        dl = torch.utils.data.DataLoader(ref_ds, batch_size=batch_size, shuffle=True)
-        nb = len(dl)
-        if num_batches is not None:
-            nb = min(nb, num_batches)
-        for i, (b, _) in tqdm.tqdm(zip(range(nb), dl), total=nb):
-            with torch.no_grad():
-                all_outputs.append(feature_extractor(b.to(device)))
-        all_outputs = torch.cat(all_outputs, dim=0)
-        self.base_outputs = all_outputs
+    def fit(self, x: torch.Tensor):
+        """Record a sample as the reference distribution"""
+        self.base_outputs = x.detach()
+        return x
 
     def predict_shift_from_features(self, base_outputs: torch.Tensor, outputs: torch.Tensor, compute_score: bool, compute_p_value: bool, individual_samples: bool = False) -> torch.Tensor:
         """stub to be overridden by subclasses"""
@@ -65,3 +44,4 @@ def forward(
         if self.return_p_value:
             return p_value
         return ood_score
+
diff --git a/torchdrift/detectors/ks.py b/torchdrift/detectors/ks.py
@@ -3,7 +3,7 @@
 import torch
 import numpy
 
-from . import DriftDetector
+from . import Detector
 
 try:
     import numba
@@ -64,7 +64,7 @@ def ks_two_sample_multi_dim(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
     ks_scores = sign.cumsum(0).abs().max(0).values
     return ks_scores
 
-class KSDriftDetector(DriftDetector):
+class KSDriftDetector(Detector):
     """Drift detector based on (multiple) Kolmogorov-Smirnov tests.
 
 This detector uses the Kolmogorov-Smirnov test on the marginals of the features
diff --git a/torchdrift/detectors/mmd.py b/torchdrift/detectors/mmd.py
@@ -2,7 +2,7 @@
 
 import torch
 
-from . import DriftDetector
+from . import Detector
 
 def kernel_mmd(x, y, n_perm=1000):
     """Implements the kernel MMD two-sample test.
@@ -59,7 +59,7 @@ def kernel_mmd(x, y, n_perm=1000):
     return mmd, p_val
 
 
-class KernelMMDDriftDetector(DriftDetector):
+class KernelMMDDriftDetector(Detector):
     """Drift detector based on the kernel Maximum Mean Discrepancy (MMD) test.
 
 This is modelled after the MMD drift detection in
diff --git a/torchdrift/reducers/__init__.py b/torchdrift/reducers/__init__.py
@@ -1 +1,2 @@
+from .reducer import Reducer
 from .pca import PCAReducer
diff --git a/torchdrift/reducers/pca.py b/torchdrift/reducers/pca.py
@@ -1,6 +1,7 @@
 import torch
+from . import Reducer
 
-class PCAReducer(torch.nn.Module):
+class PCAReducer(Reducer):
     """Reduce dimensions using PCA.
 
 This nn.Modue subclass reduces the dimensions of the inputs
@@ -13,14 +14,20 @@ def __init__(self, n_components:int = 2):
         super().__init__()
         self.n_components = n_components
 
-    def extra_repr(self):
+    def extra_repr(self) -> str:
         return f'n_components={self.n_components}'
 
-    def forward(self, x: torch.Tensor):
+    def fit(self, x: torch.Tensor) -> torch.Tensor:
         batch, feat = x.shape
         assert min(batch, feat) >= self.n_components
-        x = x - x.mean(1, keepdim=True)
+        self.mean = x.mean(0, keepdim=True)
+        x = x - self.mean
         u, s, v = x.svd()
-        comp = v[:, :self.n_components]
-        reduced = x @ comp
+        self.comp = v[:, :self.n_components]
+        reduced = x @ self.comp
+        return reduced
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x - self.mean
+        reduced = x @ self.comp
         return reduced
diff --git a/torchdrift/reducers/reducer.py b/torchdrift/reducers/reducer.py
@@ -0,0 +1,10 @@
+import torch
+
+class Reducer(torch.nn.Module):
+    """Base class for reducers"""
+
+    def fit(self, x: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError("Override fit in subclass")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError("Override forward in subclass")
diff --git a/torchdrift/utils/__init__.py b/torchdrift/utils/__init__.py
@@ -1 +1,2 @@
 from .experiments import DriftDetectionExperiment
+from .fit import fit
diff --git a/torchdrift/utils/experiments.py b/torchdrift/utils/experiments.py
@@ -1,6 +1,6 @@
 import torch
 import tqdm
-
+from .fit import fit
 
 class DriftDetectionExperiment:
     """An experimental setup to explore the ROC of drift detection setups
@@ -23,7 +23,7 @@ def __init__(self, drift_detector, feature_extractor, ood_ratio=1.0, sample_size
     # def extra_loss(self, ...):  add components to loss from training the detector
     def post_training(self, train_dataloader):
         "Called after training the main model, fits the drift detector."
-        self.drift_detector.fit(train_dataloader, self.feature_extractor)
+        fit(train_dataloader, self.feature_extractor, self.drift_detector)
 
     def evaluate(self, ind_datamodule, ood_datamodule, num_runs=50):
         """runs the experiment (`num_runs` inputs)
diff --git a/torchdrift/utils/fit.py b/torchdrift/utils/fit.py
@@ -0,0 +1,38 @@
+from typing import Optional, Callable, Union, Iterable, List
+import typing
+from ..reducers import Reducer
+from ..detectors import Detector
+
+import tqdm
+import torch
+
+def fit(
+        dl: torch.utils.data.DataLoader,
+        feature_extractor: torch.nn.Module,
+        reducers_detectors: Union[Reducer, Detector, List[Union[Reducer, Detector]]],
+        num_batches: Optional[int] = None,
+        device: Union[torch.device, str, None] = None
+    ):
+    """Train drift detector on reference distribution.
+"""
+    if not isinstance(reducers_detectors, typing.Iterable):
+        reducers_detectors = [reducers_detectors]
+    feature_extractor.eval()  # careful about test time dropout
+    if device is None:
+        device = next(feature_extractor.parameters()).device
+
+    all_outputs = []
+    # dl = torch.utils.data.DataLoader(ref_ds, batch_size=batch_size, shuffle=True)
+    nb = len(dl)
+    if num_batches is not None:
+        nb = min(nb, num_batches)
+    for i, (b, _) in tqdm.tqdm(zip(range(nb), dl), total=nb):
+        with torch.no_grad():
+            all_outputs.append(feature_extractor(b.to(device)))
+    all_outputs = torch.cat(all_outputs, dim=0)
+
+    for m in reducers_detectors:
+        if hasattr(m, 'fit'):
+            all_outputs = m.fit(all_outputs)
+        else:
+            all_outputs = m(all_outputs)

Original file line number	Diff line number	Diff line change
`@@ -1 +1,2 @@`
	`1`	`+from .reducer import Reducer`
`1`	`2`	`from .pca import PCAReducer`
Original file line number	Diff line number	Diff line change
`@@ -1 +1,2 @@`
`1`	`1`	`from .experiments import DriftDetectionExperiment`
	`2`	`+from .fit import fit`