detectors

The detectors are the central components of TorchDrift.

The inputs are expected to be 2d-Tensor s: batch x feature_dim.

class torchdrift.detectors.Detector(*, return_p_value: bool = False)

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.

Constructor Args:

return_p_value (bool): If set, forward returns a p-value (estimate) instead of the raw test score.

compute_p_value(inputs: torch.Tensor) → torch.Tensor

Performs a statistical test for drift and returns the p-value.

This method calls predict_shift_from_features under the hood, so you only need to override that when subclassing.

fit(x: torch.Tensor)

Record a sample as the reference distribution

forward(inputs: torch.Tensor, individual_samples: bool = False) → torch.Tensor

Performs a statistical test for drift and returns the score or, if return_p_value has been set in the constructor, the p-value.

This method calls predict_shift_from_features under the hood, so you only need to override that when subclassing.

predict_shift_from_features(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

class torchdrift.detectors.KSDriftDetector(*, return_p_value: bool = False)

Drift detector based on (multiple) Kolmogorov-Smirnov tests.

This detector uses the Kolmogorov-Smirnov test on the marginals of the features for each feature.

For scores, it returns the maximum score. p-values are computed with the Bonferroni correction of multiplying the p-value of the maximum score by the number of features/tests.

This is modelled after the KS drift detection in S. Rabanser et al: Failing Loudly: An Empirical Study of Methods for Detecting Dataset Shift (NeurIPS), 2019.

class torchdrift.detectors.KernelMMDDriftDetector(*, return_p_value=False, n_perm: int = 1000, kernel=<torchdrift.detectors.mmd.GaussianKernel object>)

Drift detector based on the kernel Maximum Mean Discrepancy (MMD) test.

This is modelled after the MMD drift detection in S. Rabanser et al: Failing Loudly: An Empirical Study of Methods for Detecting Dataset Shift (NeurIPS), 2019.

Note that our heuristic choice of the kernel bandwith is more closely aligned with that of the original MMD paper and code than S. Rabanser’s.

The default kernel is the unnormalized Gaussian (or Squared Exponential) kernel.

torchdrift.detectors.kernel_mmd(x, y, n_perm=1000, kernel=<torchdrift.detectors.mmd.GaussianKernel object>)

Implements the kernel MMD two-sample test.

It is modelled after the kernel MMD paper and code: A. Gretton et al.: A kernel two-sample test, JMLR 13 (2012) http://www.gatsby.ucl.ac.uk/~gretton/mmd/mmd.htm

The arguments x and y should be two-dimensional tensors. The first is the batch dimension (which may differ), the second the features (which must be the same on both x and y).

n_perm is number of bootstrap permutations to get p-value, pass None to not get p-value.

torchdrift.detectors.ks_p_value(n: int, m: int, d: float) → float

Computes the p-value for the two-sided two-sample KS test from the D-statistic.

This uses the stable recursion from T. Viehmann: Numerically more stable computation of the p-values for the two-sample Kolmogorov-Smirnov test.

torchdrift.detectors.ks_two_sample_multi_dim(x: torch.Tensor, y: torch.Tensor) → torch.Tensor

Computes the two-sample two-sided Kolmorogov-Smirnov statistic.

The inputs x and y are expected to be 2-dimensional tensors, the first dimensions being batch (potentially different) and the second features (necessarily the same).

We return a one-dimensional tensor of KS scores D, one per dimension.

Kernels for MMD

class torchdrift.detectors.mmd.Kernel

Base class for kernels

Unless otherwise noted, all kernels implementing lengthscale detection use the median of pairwise distances as the lengthscale.

class torchdrift.detectors.mmd.GaussianKernel(lengthscale=None)

Unnormalized gaussian kernel

\[k(|x-y|) = \exp(-|x-y|^2/(2\ell^2))\]

where \(\ell\) is the lengthscale (autodetected or given).

class torchdrift.detectors.mmd.ExpKernel(lengthscale=None)

Unnormalized exponential kernel

\[k(|x-y|) = \exp(-|x-y|/\ell)\]

where \(\ell\) is the lengthscale (autodetected or given).

class torchdrift.detectors.mmd.RationalQuadraticKernel(lengthscale=None, alpha=1.0)

Unnormalized rational quadratic kernel

\[k(|x-y|) = (1+|x-y|^2/(2 \alpha \ell^2))^{-\alpha}\]

where \(\ell\) is the lengthscale (autodetected or given).