update docstrings

t-vi · t-vi · commit d20f1c0c1d17 · 2021-02-19T12:59:47.000+01:00
diff --git a/torchdrift/reducers/reducer.py b/torchdrift/reducers/reducer.py
@@ -1,10 +1,22 @@
 import torch
 
 class Reducer(torch.nn.Module):
-    """Base class for reducers"""
+    """Base class for reducers.
+
+This is a `torch.nn.Module` with an additional `fit` method.
+The usual forward is for testing after fitting."""
 
     def fit(self, x: torch.Tensor) -> torch.Tensor:
+        """Fits the reducer to reference data `x` and returns the reduced
+data.
+
+Override this in your reducer implementation.
+"""
         raise NotImplementedError("Override fit in subclass")
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Reduces the input `x` (in testing) and returns the reduced data.
+
+Override this in your reducer implementation.
+"""
         raise NotImplementedError("Override forward in subclass")
diff --git a/torchdrift/utils/fit.py b/torchdrift/utils/fit.py
@@ -10,10 +10,20 @@ 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.
+
+The dataloader `dl` should provide the reference distribution. Optionally you can limit the number of batches sampled from the dataloader with `num_batches`.
+
+The `feature extractor` can be any module be anything that does not need to be fit.
+
+The reducers and detectors should be passed (in the order they should be applied, one takes the output from the previous) as a list. A single detector or reducer can also be passed.
+
+If you provide a `device`, data is moved there before running through the
+feature extractor, otherwise the functions try to infer the device from the `feature_extractor`.
 """
     if not isinstance(reducers_detectors, typing.Iterable):
         reducers_detectors = [reducers_detectors]
