Add custom training fitness callbacks

docs/checkpoint_schema.md

@@ -391,6 +391,15 @@ distributed as training checkpoints.
 For release compatibility, readers accept legacy best-metric aliases such as
 `best_mAP50_95`, `best_mAP50`, `best_metric`, and `best_metric_name`.
 
+Custom-fitness training writes optional `fitness_source: "callback"` and
+`best_metric_key: "fitness/custom"` in its training and averaged checkpoints.
+`best_metric_value` and the legacy `best_mAP50_95` / `best_metric` aliases then
+contain the custom score, while `best_mAP50` remains the validation mAP50 at
+the selected epoch. The marker stores no callback code or state. Such files
+can be loaded for inference or as weights for a new run, but cannot resume
+training; neither can a default checkpoint resume with a custom scorer.
+Absent `fitness_source` preserves the existing checkpoint/resume contract.
+
 ## External Snapshot Exception
 
 The schema above governs LibreYOLO-authored `.pt` checkpoints. It does not

docs/training_loggers.md

@@ -49,6 +49,54 @@ reference but cannot be transported by value (for example, one that captures a
 write-only file handle) uses the guarded-script compatibility launcher and
 emits a warning; keep that call under the `__main__` guard.
 
+### Custom fitness
+
+For families using the shared `BaseTrainer`, including YOLO9 and RF-DETR,
+one callback object may define `fitness(metrics)`. Its finite real scalar
+return value selects `best.pt` and drives the existing `patience` setting;
+higher is better. Plain functions remain epoch observers. No CLI code loader
+is provided: this hook is a Python callable.
+
+```python
+class ValidationLossFitness:
+    def fitness(self, metrics):
+        return -metrics["metrics/loss"]
+
+model.train(
+    data="coco8.yaml", val_loss=True, patience=10,
+    callbacks=ValidationLossFitness(),
+)
+```
+
+The scorer receives a read-only copy of the scalar validation metrics, with
+the same names as `TrainEpochEvent.val_metrics`. Use the metrics your task
+actually reports; missing keys raise normally. It runs on rank zero before
+checkpoint selection and epoch observers, once for each validation result
+used for selection. Skipped validation does not call it. Precise-BN
+revalidation and final averaged-weight validation also use it. Keep the
+scorer deterministic and independent of invocation count.
+
+Events and checkpoints name the selected score `fitness/custom`; validation
+metrics keep their original values. `average_best` ranks snapshots by the
+same score. Ties keep the earlier best, and patience still counts epochs
+since improvement, including epochs without validation. Multiple scorers,
+non-numeric results, booleans, non-scalar tensors, NaN and infinity raise
+before updating best state or writing that epoch's checkpoint.
+
+The returned training result's legacy `best_mAP50_95` field also contains the
+selected score, as it already does for non-detection task metrics. Read raw
+accuracy values from `val_metrics`; `epoch_metrics` pairs `current_metric`
+and `best_metric` with their metric names.
+
+Custom-fitness training does not support `resume`. Callback code and state
+are not serialized, and the library cannot establish whether a new scorer
+is comparable with historical scores. Resuming with a scorer or from a
+custom-fitness checkpoint raises. To continue from its weights, load that
+checkpoint and start a new run with `resume=False` and a new output directory;
+this starts fresh optimizer, best-score and patience tracking. Ordinary
+training and ordinary resume are unchanged. VLM and VLA trainers do not
+support this scorer.
+
 ## Built-in loggers
 
 Built-in loggers are callback objects layered on the same universal hooks.
@@ -145,7 +193,8 @@ monitor overlays it like every other family.
 This option is off by default because target assignment adds work and memory
 to validation. It runs under `torch.no_grad()` with the evaluation/EMA model,
 and distributed training computes it locally on rank 0 without collectives.
-Best-checkpoint selection remains based on the configured accuracy metric.
+Best-checkpoint selection uses the task's default metric unless a custom
+fitness callback is supplied.
 Augmented validation, a task a family has not implemented it for, and
 inference-only (`g3`/`g4`) families all raise a clear configuration error.
 

libreyolo/models/fomo/trainer.py

@@ -317,7 +317,6 @@ def _run_validation(
             raise
 
     def _checkpoint_extra_metadata(self) -> Dict[str, Any]:
-        return {
-            "task": "point",
-            "best_metric_key": "metrics/grid_F1",
-        }
+        # Direct trainer users may omit the model wrapper, whose task would
+        # otherwise supply this value to the shared checkpoint writer.
+        return {"task": "point"}

libreyolo/models/rfdetr/model.py

@@ -1236,7 +1236,9 @@ def train(
                 ``<RFDETRConfig.project>/<RFDETRConfig.name>`` when omitted;
                 ``project=`` / ``name=`` kwargs take precedence over this split.
             resume: Checkpoint path, or True to resume the loaded checkpoint.
-            callbacks: Optional training callback or iterable of callbacks.
+            callbacks: Optional callback or iterable. One object may define
+                fitness(metrics) to select best.pt and drive patience; custom
+                fitness requires a new run (resume=False).
             loggers: Optional built-in experiment loggers: a registered name,
                 a configured logger instance, or an iterable mixing both.
         """

libreyolo/models/vla/training/trainer.py

@@ -216,6 +216,8 @@ def __init__(
         self.callbacks = TrainCallbackList(callbacks)
         for logger_cb in resolve_loggers(loggers):
             self.callbacks.append(logger_cb)
+        if self.callbacks.fitness is not None:
+            raise NotImplementedError("VLA training does not support custom fitness callbacks")
         self.save_dir = self._resolve_save_dir()
 
     def _resolve_save_dir(self) -> Path:

libreyolo/models/vlm/training/trainer.py

@@ -111,6 +111,8 @@ def __init__(
         self.callbacks = TrainCallbackList(callbacks)
         for logger_cb in resolve_loggers(loggers):
             self.callbacks.append(logger_cb)
+        if self.callbacks.fitness is not None:
+            raise NotImplementedError("VLM training does not support custom fitness callbacks")
         if self.config.resume is True:
             # resume=True reads weights/last from this exact save_dir (see
             # _resolve_resume_dir); never let _resolve_save_dir() increment

libreyolo/models/yolo9/model.py

@@ -652,7 +652,9 @@ def train(
             pretrained: Optional training initialization weights. Use True to
                 load the matching LibreYOLO9 detect checkpoint for transfer
                 learning, or pass a checkpoint path/name.
-            callbacks: Optional training callback or iterable of callbacks.
+            callbacks: Optional callback or iterable. One object may define
+                fitness(metrics) to select best.pt and drive patience; custom
+                fitness requires a new run (resume=False).
             loggers: Optional built-in experiment loggers: a registered name,
                 a configured logger instance, or an iterable mixing both.
 

libreyolo/training/__init__.py

@@ -11,6 +11,7 @@
     TrainEndEvent as TrainEndEvent,
     TrainEpochEvent as TrainEpochEvent,
     TrainExceptionEvent as TrainExceptionEvent,
+    TrainFitnessCallback as TrainFitnessCallback,
     TrainStartEvent as TrainStartEvent,
 )
 from .config import (

libreyolo/training/callbacks.py

@@ -1,8 +1,9 @@
 """Public training callback types."""
 
+from collections.abc import Callable, Iterable, Mapping
 from dataclasses import dataclass, field
 from types import MappingProxyType
-from typing import Any, Callable, Iterable, Mapping, Protocol
+from typing import Any, Protocol
 
 
 @dataclass(frozen=True)
@@ -116,8 +117,22 @@ def on_train_exception(self, event: TrainExceptionEvent) -> None:
         """Handle a training exception."""
 
 
+class TrainFitnessCallback(Protocol):
+    """Optional scorer for best-checkpoint selection and patience.
+
+    ``metrics`` is the same read-only scalar mapping as
+    ``TrainEpochEvent.val_metrics``. Return a finite real scalar; higher is
+    better. Scoring runs on rank zero for each validation result used for
+    selection, before checkpoints and epoch callbacks. Callback code and
+    state are not saved, so custom-fitness training cannot be resumed.
+    """
+
+    def fitness(self, metrics: Mapping[str, float]) -> float:
+        """Score validation metrics without modifying them."""
+
+
 TrainEpochCallable = Callable[[TrainEpochEvent], None]
-TrainCallbackLike = TrainCallback | TrainEpochCallable
+TrainCallbackLike = TrainCallback | TrainFitnessCallback | TrainEpochCallable
 TrainCallbacks = TrainCallbackLike | Iterable[TrainCallbackLike] | None
 
 
@@ -131,6 +146,21 @@ def __init__(self, callbacks: TrainCallbacks = None):
             self._callbacks = [callbacks]
         else:
             self._callbacks = list(callbacks)
+        _ = self.fitness  # Validate scorers before any training work starts.
+
+    @property
+    def fitness(self) -> Callable[[Mapping[str, float]], float] | None:
+        """Return the sole optional fitness method, rejecting ambiguity."""
+        scorers = []
+        for callback in self._callbacks:
+            if hasattr(callback, "fitness"):
+                method = callback.fitness
+                if not callable(method):
+                    raise TypeError("Train callback attribute fitness must be callable")
+                scorers.append(method)
+        if len(scorers) > 1:
+            raise ValueError("At most one training callback may define fitness")
+        return scorers[0] if scorers else None
 
     def __bool__(self) -> bool:
         return bool(self._callbacks)
@@ -140,6 +170,11 @@ def __len__(self) -> int:
 
     def append(self, callback: TrainCallbackLike) -> None:
         self._callbacks.append(callback)
+        try:
+            _ = self.fitness
+        except (TypeError, ValueError):
+            self._callbacks.pop()
+            raise
 
     @staticmethod
     def _is_callback_object(callback) -> bool:
@@ -150,6 +185,7 @@ def _is_callback_object(callback) -> bool:
                 "on_train_epoch_end",
                 "on_train_end",
                 "on_train_exception",
+                "fitness",
             )
         )