Merge pull request #1 from alexshtf/early_stop_mode

Add mode and min_delta to early_stop

Author: alexshtf

README.md

@@ -120,12 +120,15 @@ events = pipe(
     epoch_stream(...),
     augment(validation_loss(...)),
     take(500),  # safety cap
-    early_stop(key="val_loss", patience=10),
+    early_stop(key="val_loss", patience=10, mode="min", min_delta=1e-4),
 )
 for event in events:
     print(event)
 ```
 
+`mode="min"` is the default. Use `mode="max"` for metrics such as accuracy, and set
+`min_delta` to ignore tiny noisy changes.
+
 ## Side effects
 Sometimes you want to log metrics (or write to an external system) without changing the stream. Use `tap(fn)`:
 ```python

docs/tutorial.md

@@ -241,7 +241,7 @@ events = pipe(
     epoch_stream((x_train, y_train), model, optimizer, loss_fn, batch_size=512, shuffle=True),
     augment(validation_loss((x_val, y_val), loss_fn)),
     take(500),
-    early_stop(key="val_loss", patience=10),
+    early_stop(key="val_loss", patience=10, mode="min", min_delta=1e-4),
 )
 
 history = list(events)  # finite now
@@ -250,9 +250,22 @@ print("stopped at epoch", history[-1]["step"])
 
 Notes:
 
-- `early_stop(..., key="val_loss", ...)` assumes “lower is better”.
+- `early_stop(..., mode="min")` (the default) treats lower values as better.
+- Use `mode="max"` for metrics where higher is better (for example `val_acc`).
+- `min_delta` is an absolute threshold for improvement, so tiny metric noise does not reset patience.
 - It yields events up to (and including) the epoch that triggers stopping.
 
+For an accuracy metric, switch to `mode="max"`:
+
+```python
+events = pipe(
+    epoch_stream(...),
+    augment(...),  # produce "val_acc" on each event
+    take(500),
+    early_stop(key="val_acc", patience=10, mode="max", min_delta=1e-3),
+)
+```
+
 ## 6) Become a hero: write your own augmenter
 
 An **augmenter** is a function `event -> dict` (or `None`) that adds keys to the event. This is great for:
@@ -462,7 +475,7 @@ events = pipe(
     ema("val_loss", alpha=0.2),
     print_every(10, keys=("train_loss", "val_loss", "val_loss_ema", "param_l2")),
     take(500),
-    early_stop(key="val_loss", patience=20),
+    early_stop(key="val_loss", patience=20, mode="min", min_delta=1e-4),
 )
 
 # Write the whole training history to disk (one JSON object per line).

src/fitstream/fit.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Iterable, Iterator, Sequence
-from typing import Any
+from typing import Any, Literal
 import time
 
 import torch
@@ -108,22 +108,43 @@ def stage(events: Iterable[dict[str, Any]]) -> Iterable[dict[str, Any]]:
 def early_stop(
     key: str,
     patience: int,
+    *,
+    mode: Literal["min", "max"] = "min",
+    min_delta: float = 0.0,
 ) -> Transform:
     """Yield events until the metric stops improving for `patience` steps.
 
+    Args:
+        key: Event key containing the monitored metric.
+        patience: Number of consecutive non-improving events tolerated before stopping.
+        mode: Improvement direction. ``"min"`` means lower is better, ``"max"`` means
+            higher is better.
+        min_delta: Minimum absolute change required to count as an improvement.
+
     Use as a pipe stage:
 
     - ``pipe(events, early_stop(key="val_loss", patience=10))``
     """
     if patience < 1:
         raise ValueError("patience must be >= 1.")
+    if mode not in {"min", "max"}:
+        raise ValueError("mode must be one of {'min', 'max'}.")
+    if min_delta < 0.0:
+        raise ValueError("min_delta must be >= 0.")
 
     def apply(stream: Iterable[dict[str, Any]]) -> Iterable[dict[str, Any]]:
-        best = float("inf")
+        best: float | None = None
         bad = 0
         for event in stream:
             value = float(event[key])
-            if value < best:
+            if best is None:
+                improved = True
+            elif mode == "min":
+                improved = value < (best - min_delta)
+            else:
+                improved = value > (best + min_delta)
+
+            if improved:
                 best = value
                 bad = 0
             else:

tests/test_early_stop.py

@@ -16,6 +16,49 @@ def test_early_stop_rejects_invalid_patience() -> None:
         early_stop(key="val_loss", patience=0)
 
 
+def test_early_stop_explicit_min_mode_matches_default() -> None:
+    events = [{"val_loss": loss} for loss in [5.0, 4.0, 4.0, 4.0, 3.0]]
+
+    result_default = list(early_stop(key="val_loss", patience=2)(events))
+    result_explicit = list(early_stop(key="val_loss", patience=2, mode="min", min_delta=0.0)(events))
+
+    assert result_explicit == result_default
+
+
+def test_early_stop_supports_max_mode() -> None:
+    events = [{"val_acc": value} for value in [0.50, 0.60, 0.60, 0.59, 0.70]]
+
+    result = list(early_stop(key="val_acc", patience=2, mode="max")(events))
+
+    assert [event["val_acc"] for event in result] == [0.50, 0.60, 0.60, 0.59]
+
+
+def test_early_stop_min_delta_for_min_mode() -> None:
+    events = [{"val_loss": value} for value in [1.00, 0.95, 0.92, 0.85]]
+
+    result = list(early_stop(key="val_loss", patience=2, min_delta=0.1)(events))
+
+    assert [event["val_loss"] for event in result] == [1.00, 0.95, 0.92]
+
+
+def test_early_stop_min_delta_for_max_mode() -> None:
+    events = [{"val_acc": value} for value in [0.50, 0.53, 0.56, 0.57, 0.58]]
+
+    result = list(early_stop(key="val_acc", patience=2, mode="max", min_delta=0.05)(events))
+
+    assert [event["val_acc"] for event in result] == [0.50, 0.53, 0.56, 0.57, 0.58]
+
+
+def test_early_stop_rejects_invalid_mode() -> None:
+    with pytest.raises(ValueError):
+        early_stop(key="val_loss", patience=2, mode="lower")  # type: ignore[arg-type]
+
+
+def test_early_stop_rejects_negative_min_delta() -> None:
+    with pytest.raises(ValueError):
+        early_stop(key="val_loss", patience=2, min_delta=-0.1)
+
+
 def test_early_stop_is_pipe_stage() -> None:
     events = [{"val_loss": loss} for loss in [5.0, 4.0, 4.0, 4.0, 3.0]]