Callbacks

substrax.callbacks holds the training-callback protocol and metric-driven early stopping, in two shapes:

• EarlyStopping for a loop that feeds it one validation value per epoch and asks should_stop;
• EarlyStoppingCallback for a trainer that drives TrainingCallback hooks and reads the metric from the epoch's logs.

EarlyStopping keeps its counters in a BestMetricTracker; EarlyStoppingCallback applies an EarlyStopping to the logged metric and adds the goal and divergence thresholds and the non-finite check. patience is at least 1: patience counts epochs without improvement, so an improving epoch never stops training.

from substrax.callbacks import EarlyStopping, EarlyStoppingCallback, EarlyStoppingConfig

stopper = EarlyStopping(patience=3, min_delta=1e-3)
for epoch in range(epochs):
    stopper.update(validate())
    if stopper.should_stop:
        break

callback = EarlyStoppingCallback(EarlyStoppingConfig(monitor="val_loss", patience=10))

Learning-rate decay on a plateau is not here: use optax.contrib.reduce_on_plateau.

Resuming a run

BestMetricTracker, EarlyStopping, EarlyStoppingCallback and CallbackList satisfy substrax.typing.Checkpointable: get_state() returns their bookkeeping as a dictionary a checkpoint store writes, and set_state() on objects built the same way continues where the saved ones stood, so a resumed run stops at the epoch the uninterrupted run would have. CallbackList keeps each stateful callback's state under its position, skips callbacks without state, and refuses a state saved from a list whose stateful callbacks stood elsewhere.

from substrax.callbacks import CallbackList
from substrax.checkpoint import OrbaxCheckpointStore

callbacks = CallbackList([EarlyStoppingCallback(EarlyStoppingConfig(patience=10))])
with OrbaxCheckpointStore(directory) as store:
    store.save(step, {"extensions": {"callbacks": callbacks.get_state()}})
    callbacks.set_state(store.restore(step).items["extensions"]["callbacks"])

substrax.callbacks

Training callbacks, best-metric tracking and early stopping.

BaseCallback

No-op implementation of every hook; subclasses override what they need.

on_train_begin(_trainer)

Called once before the first epoch.

on_train_end(_trainer)

Called once after the last epoch.

on_epoch_begin(_trainer, _epoch)

Called before each epoch.

on_epoch_end(_trainer, _epoch, _logs)

Called after each epoch with that epoch's metrics.

on_batch_begin(_trainer, _batch)

Called before each batch.

on_batch_end(_trainer, _batch, _logs)

Called after each batch with that batch's metrics.

on_validation_begin(_trainer)

Called before a validation pass.

on_validation_end(_trainer, _logs)

Called after a validation pass with its metrics.

CallbackList(callbacks=None)

Dispatches every hook to its callbacks in registration order.

Wrap an initial set of callbacks.

Parameters:

Name	Type	Description	Default
callbacks	Iterable[TrainingCallback] | None	The callbacks to start with; each receives every hook, in this order.	None

add(callback)

Append a callback.

remove(callback)

Remove a callback; raises ValueError if it was never added.

__len__()

The number of callbacks.

__iter__()

Iterate over the callbacks in dispatch order.

get_state()

Every stateful callback's state, for a checkpoint.

Returns:

Name	Type	Description
Each	CheckpointState	class:~substrax.typing.Checkpointable callback's state under its position
	CheckpointState	in dispatch order; callbacks without state are skipped.

set_state(state)

Take back a state :meth:get_state returned, into a list built the same way.

Parameters:

Name	Type	Description	Default
state	CheckpointState	Each stateful callback's saved state, by position.	required

Raises:

Type	Description
ValueError	If the stateful callbacks do not stand at the saved positions.

on_train_begin(trainer)

Dispatch on_train_begin.

on_train_end(trainer)

Dispatch on_train_end.

on_epoch_begin(trainer, epoch)

Dispatch on_epoch_begin.

on_epoch_end(trainer, epoch, logs)

Dispatch on_epoch_end.

on_batch_begin(trainer, batch)

Dispatch on_batch_begin.

on_batch_end(trainer, batch, logs)

Dispatch on_batch_end.

on_validation_begin(trainer)

Dispatch on_validation_begin.

on_validation_end(trainer, logs)

Dispatch on_validation_end.

TrainerLike

Bases: Protocol

What a callback may ask of the trainer that drives it.

model property

The model being trained.

TrainingCallback

Bases: Protocol

Lifecycle hooks a training loop fires; implement any subset via BaseCallback.

Parameters are positional-only so implementations may name them as they like.

on_train_begin(trainer)

Called once before the first epoch.

on_train_end(trainer)

Called once after the last epoch.

on_epoch_begin(trainer, epoch)

Called before each epoch.

on_epoch_end(trainer, epoch, logs)

Called after each epoch with that epoch's metrics.

on_batch_begin(trainer, batch)

Called before each batch.

on_batch_end(trainer, batch, logs)

Called after each batch with that batch's metrics.

on_validation_begin(trainer)

Called before a validation pass.

on_validation_end(trainer, logs)

Called after a validation pass with its metrics.

EarlyStoppingCallback(config)

Bases: BaseCallback

Stop training when a monitored metric stops improving.

Initialize early stopping callback.

Parameters:

Name	Type	Description	Default
config	EarlyStoppingConfig	Early stopping configuration.	required

wait_count property

Epochs since the last improvement.

best_score property

The best monitored value so far, or None before the metric first appears.

stopped_epoch property

The epoch at which stopping was decided, if it has been.

should_stop property

Whether training should stop.

on_epoch_end(_trainer, epoch, logs)

Read the monitored metric, if present, and decide whether to stop.

Parameters:

Name	Type	Description	Default
_trainer	TrainerLike	The trainer instance (unused).	required
epoch	int	Current epoch number.	required
logs	dict[str, Any]	Dictionary of metrics from this epoch.	required

get_state()

The stopper's state and the stopping epoch, for a checkpoint.

Returns:

Type	Description
CheckpointState	stopper (see :meth:BestMetricTracker.get_state) and stopped_epoch
CheckpointState	(None while training continues).

set_state(state)

Take back a state :meth:get_state returned.

Parameters:

Name	Type	Description	Default
state	CheckpointState	The saved stopper state and stopping epoch.	required

EarlyStoppingConfig(*, monitor='val_loss', min_delta=0.0, patience=10, mode='min', check_finite=True, stopping_threshold=None, divergence_threshold=None) dataclass

Configuration for early stopping.

Attributes:

Name	Type	Description
monitor	str	Metric name to monitor (e.g., "val_loss", "accuracy").
min_delta	float	Minimum change to qualify as an improvement.
patience	int	Number of epochs with no improvement before stopping.
mode	Literal['min', 'max']	"min" if lower is better, "max" if higher is better.
check_finite	bool	If True, stop when metric becomes NaN or Inf.
stopping_threshold	float | None	Stop immediately when metric reaches this value.
divergence_threshold	float | None	Stop if metric exceeds this value (min mode only).

__post_init__()

Refuse a patience below one, which would stop on an improving epoch.

Raises:

Type	Description
ValueError	If patience is below one.

BestMetricTracker(*, mode, min_delta)

Shared best-so-far + stagnation bookkeeping for the plateau callbacks.

Tracks the best monitored value and the number of consecutive updates without a min_delta improvement, in either "min" or "max" mode.

Initialise the tracker.

Parameters:

Name	Type	Description	Default
mode	PlateauMode | str	"min" (lower is better) or "max" (higher is better).	required
min_delta	float	Minimum absolute change counted as an improvement.	required

Raises:

Type	Description
ValueError	If min_delta is negative.

mode property

The improvement direction.

best property

The best monitored value seen so far.

num_bad_epochs property

Consecutive updates without a min_delta improvement.

register(value)

Record value; return True if it improves on the best so far.

Updates the best value and resets the stagnation counter on an improvement; otherwise increments the stagnation counter.

get_state()

The best value and the stagnation count, for a checkpoint.

Returns:

Type	Description
CheckpointState	best (infinite before the first value) and num_bad_epochs.

set_state(state)

Take back a state :meth:get_state returned.

Parameters:

Name	Type	Description	Default
state	CheckpointState	The saved best value and stagnation count.	required

EarlyStopping(*, patience, min_delta=0.0, mode=PlateauMode.MIN)

Bases: BestMetricTracker

Signal to stop once a monitored metric stops improving.

Initialise the stopper.

Parameters:

Name	Type	Description	Default
patience	int	Epochs without a min_delta improvement before stopping.	required
min_delta	float	Minimum absolute change counted as an improvement.	0.0
mode	PlateauMode | str	"min" (lower is better) or "max" (higher is better).	MIN

Raises:

Type	Description
ValueError	If patience is not positive.

should_stop property

Whether the metric has stagnated for patience epochs.

update(value)

Record the latest metric; return True if it improved on the best.

PlateauMode

Bases: StrEnum

Whether a monitored metric improves by decreasing or increasing.