Callbacks API

torch_batteries.callbacks

Reusable callbacks for workflow control and optimization.

Callback contract

• Callback — optional resumable-state base class for custom callbacks.

Training control

• EarlyStopping — stops training when a monitored metric stops improving.
• ModelCheckpoint — retains the best weights-only or full-state checkpoints.
• TerminateOnNonFinite — fails on selected NaN or infinite losses and metrics.
• ExperimentTrackingCallback — logs training lifecycle data through a tracker.

Optimization

• GradientAccumulation — groups batches into fewer optimizer steps.
• GradientClip — applies value- or norm-based gradient clipping.
• MixedPrecision — supplies autocast, scaled backward, and optimizer stepping.
• LearningRateScheduler — advances step-, epoch-, or validation-based schedulers.

Callback Source

Base class for callbacks with optional resumable state.

Decorator-only callback objects remain supported. Inheriting from this class enables a custom callback to participate in full training checkpoints.

state_dict() Source

Return state that should be stored in a full training checkpoint.

Stateless callbacks inherit the empty default. Stateful callbacks should return values supported by torch.save(..., weights_only=True).

load_state_dict(state_dict) Source

Restore state from a full training checkpoint.

Battery matches callback state strictly by callback type and configured order before invoking this method.

Parameters:

Name	Type	Description	Default
state_dict	dict[str, Any]	State previously returned by :meth:state_dict.	required

EarlyStopping Source

Bases: Callback

Early stops the training if selected metric doesn't improve after a given patience.

Parameters:

Name	Type	Description	Default
phase	MonitorPhase | None	"train" or "validation". "val" is deprecated.	None
metric	str | None	The name of the metric to monitor	None
min_delta	float	Minimum change in the monitored metric to qualify as an improvement	0.0
patience	int	Number of epochs with no improvement after which training will be stopped	5
mode	Literal['min', 'max']	One of 'min' or 'max'. In 'min' mode, training will stop when the monitored metric stops decreasing. In 'max' mode, it will stop when the metric stops increasing	'min'
restore_best_weights	bool	If True, restore model weights from the epoch with the best value of the monitored metric	False
stage	MonitorPhase | None	Deprecated keyword alias for phase.	None

best_score property

Get the best score observed so far.

best_weights property

Get the best model weights observed so far.

state_dict() Source

Return early-stopping state for resumable checkpoints.

load_state_dict(state_dict) Source

Restore early-stopping state from a checkpoint.

Parameters:

Name	Type	Description	Default
state_dict	dict[str, Any]	State returned by :meth:state_dict.	required

run_on_train_start(context) Source

Initialize early stopping parameters at the start of training.

Parameters:

Name	Type	Description	Default
context	EventContext	Event context. A resumed run preserves restored callback state.	required

run_on_epoch_end(context) Source

Check early stopping after training epoch ends.

Parameters:

Name	Type	Description	Default
context	EventContext	Event context containing training metrics.	required

run_on_validation_end(context) Source

Check early stopping after validation ends.

Parameters:

Name	Type	Description	Default
context	EventContext	Event context containing validation metrics.	required

run_on_train_end(context) Source

Restore best model weights after training ends if configured.

Parameters:

Name	Type	Description	Default
context	EventContext	Event context containing the model.	required

ExperimentTrackingCallback Source

Bases: Callback

Callback for automatic experiment tracking during training.

Integrates an ExperimentTracker with Battery training loop, automatically logging configuration, metrics, and summary.

This callback hooks into the event system to log: - Configuration at training start - Training metrics after each step - Validation metrics after validation - Summary statistics at training end

Example:

from torch_batteries.tracking import WandbTracker, Run

# Create tracker and configure run
tracker = WandbTracker(project="your-wandb-project")
run = Run(config={"lr": 0.001, "patience": 5})

# Create callback
callback = ExperimentTrackingCallback(
    tracker=tracker,
    run=run,
)

# Use with Battery
battery = Battery(model, optimizer=optimizer, callbacks=[callback])
battery.fit(train_loader, val_loader, epochs=10)

tracker property

Get the experiment tracker.

run property

Get the run configuration.

global_step property

Current global step.

current_epoch property

Current epoch.

log_every_n_steps property

How often to log metrics (in steps).

__init__(tracker, run=None, log_every_n_steps=1) Source

Initialize the experiment tracking callback.

Parameters:

Name	Type	Description	Default
tracker	ExperimentTracker	The experiment tracker instance	required
run	Run | None	Optional run configuration	None
log_every_n_steps	int	Positive interval between training metric log calls.	1

state_dict() Source

Return experiment progress counters.

load_state_dict(state_dict) Source

Restore experiment progress counters.

Parameters:

Name	Type	Description	Default
state_dict	dict[str, Any]	State returned by :meth:state_dict.	required

on_train_start(_) Source

Initialize tracker and log configuration.

Parameters:

Name	Type	Description	Default
_	EventContext	Event context, unused by this handler.	required

on_epoch_start(ctx) Source

Update current epoch.

Parameters:

Name	Type	Description	Default
ctx	EventContext	Event context	required

on_train_step_end(ctx) Source

Log training metrics after each step.

Parameters:

Name	Type	Description	Default
ctx	EventContext	Event context	required

on_validation_epoch_end(ctx) Source

Log validation metrics.

Parameters:

Name	Type	Description	Default
ctx	EventContext	Event context	required

on_train_end(ctx) Source

Log summary and finish tracking.

Parameters:

Name	Type	Description	Default
ctx	EventContext	Event context	required

on_exception(_) Source

Finish an initialized tracker with failure status.

Parameters:

Name	Type	Description	Default
_	EventContext	Exception context, unused by this handler.	required

GradientAccumulation Source

Bases: Callback

Accumulate gradients across multiple batches before optimizer steps.

Parameters:

Name	Type	Description	Default
steps	int	Maximum number of batches in one accumulation group.	required

steps property

Number of batches accumulated per optimizer step.

optimizer_step_idx property

Number of optimizer steps completed by this control.

reset() Source

Reset optimizer-step progress for a new training run.

is_group_start(batch_idx) Source

Return whether this batch begins an accumulation group.

Parameters:

Name	Type	Description	Default
batch_idx	int	Zero-based batch index in the current epoch.	required

is_group_end(batch_idx, total_batches) Source

Return whether this batch completes an accumulation group.

Parameters:

Name	Type	Description	Default
batch_idx	int	Zero-based batch index in the current epoch.	required
total_batches	int	Number of batches in the current epoch.	required

group_size(batch_idx, total_batches) Source

Return the actual size of the current accumulation group.

Parameters:

Name	Type	Description	Default
batch_idx	int	Zero-based batch index in the current epoch.	required
total_batches	int	Number of batches in the current epoch.	required

record_optimizer_step() Source

Record and return a completed optimizer-step index.

on_train_start(context) Source

Reset progress when training is not resuming from a checkpoint.

Parameters:

Name	Type	Description	Default
context	EventContext	Training-start event context containing the resume flag.	required

configure_train_step(context) Source

Return zeroing, scaling, and step decisions for the current batch.

Parameters:

Name	Type	Description	Default
context	EventContext	Train-step context with the batch index and batch count.	required

on_optimizer_step_end(context) Source

Synchronize callback state with Battery's completed-step counter.

Parameters:

Name	Type	Description	Default
context	EventContext	Optimizer event context containing optimizer_step_idx.	required

state_dict() Source

Return resumable accumulation state.

load_state_dict(state_dict) Source

Restore accumulation state.

Parameters:

Name	Type	Description	Default
state_dict	dict[str, Any]	State returned by :meth:state_dict.	required

GradientClip Source

Bases: Callback

Clip gradients immediately before optimizer steps.

norm scales all gradients proportionally when their combined norm exceeds value. value clamps each gradient element independently.

Parameters:

Name	Type	Description	Default
value	float	Non-negative clipping threshold.	required
algorithm	ClipAlgorithm	"norm" for global norm scaling or "value" for element-wise clamping.	'norm'

value property

Configured clipping threshold.

algorithm property

Configured clipping algorithm.

apply(parameters) Source

Clip gradients and return the pre-clip norm when available.

Parameters:

Name	Type	Description	Default
parameters	Iterable[Parameter]	Parameters whose non-null gradients should be clipped.	required

run_gradient_clip(context) Source

Clip gradients exposed by the optimization event context.

Parameters:

Name	Type	Description	Default
context	EventContext	Gradient event context containing the model.	required

state_dict() Source

Return clipping configuration for strict checkpoint validation.

load_state_dict(state_dict) Source

Validate restored clipping configuration.

Parameters:

Name	Type	Description	Default
state_dict	dict[str, Any]	State returned by :meth:state_dict.	required

LearningRateScheduler Source

Bases: Callback

Advance a PyTorch learning-rate scheduler during Battery training.

phase selects the monitored metrics for ReduceLROnPlateau. The deprecated stage keyword remains a compatibility alias.

Parameters:

Name	Type	Description	Default
scheduler	LRScheduler	PyTorch scheduler to advance.	required
interval	SchedulerInterval	"step" after optimizer steps or "epoch" after epochs.	'epoch'
phase	SchedulerPhase | None	Metrics phase for ReduceLROnPlateau.	None
metric	str | None	Metric name for ReduceLROnPlateau.	None
stage	SchedulerPhase | None	Deprecated alias for phase.	None

scheduler property

Underlying PyTorch scheduler.

interval property

Scheduler advancement interval.

on_optimizer_step_end(context) Source

Advance step schedulers after actual optimizer steps.

Parameters:

Name	Type	Description	Default
context	EventContext	Optimizer-step event context.	required

on_train_epoch_end(context) Source

Advance ordinary epoch schedulers or train-monitored plateau schedulers.

Parameters:

Name	Type	Description	Default
context	EventContext	Completed training-epoch context and metrics.	required

on_validation_end(context) Source

Advance validation-monitored plateau schedulers.

Parameters:

Name	Type	Description	Default
context	EventContext	Completed validation context and metrics.	required

on_train_end(context) Source

Validate that a requested validation metric was observed.

Parameters:

Name	Type	Description	Default
context	EventContext	Training-end context containing the final epoch.	required

state_dict() Source

Return scheduler and advancement state.

load_state_dict(state_dict) Source

Restore scheduler state after validating its configuration.

Parameters:

Name	Type	Description	Default
state_dict	dict[str, Any]	State returned by :meth:state_dict, including legacy stage state.	required

MixedPrecision Source

Bases: Callback

Apply full or mixed precision across all Battery workflow phases.

Parameters:

Name	Type	Description	Default
precision	Precision	"32-true", "16-mixed", "bf16-mixed", or "amp". Automatic mode selects bf16 on CPU and fp16 otherwise.	'amp'

precision property

Requested precision mode.

effective_precision property

Device-resolved precision mode.

scaler property

Gradient scaler used by fp16 training.

configure(device) Source

Resolve the requested precision for a concrete device.

Parameters:

Name	Type	Description	Default
device	device	Device used by the Battery workflow.	required

autocast() Source

Create an autocast context for the configured precision.

backward(loss) Source

Backpropagate a normalized loss with optional gradient scaling.

Parameters:

Name	Type	Description	Default
loss	Tensor	Scalar loss to backpropagate.	required

optimizer_step(optimizer) Source

Apply an optimizer step and update the gradient scaler.

Parameters:

Name	Type	Description	Default
optimizer	Optimizer	Optimizer whose step should be executed.	required

unscale_(optimizer) Source

Unscale optimizer gradients before operations such as clipping.

Parameters:

Name	Type	Description	Default
optimizer	Optimizer	Optimizer owning the gradients.	required

on_setup(context) Source

Resolve precision using Battery's selected device.

Parameters:

Name	Type	Description	Default
context	EventContext	Setup event context containing the selected device.	required

step_execution_context(context) Source

Provide autocast for train, validation, test, and prediction steps.

Parameters:

Name	Type	Description	Default
context	EventContext	Step context identifying the active phase.	required

run_backward(context) Source

Execute scaled or ordinary backward through the event contract.

Parameters:

Name	Type	Description	Default
context	EventContext	Backward context containing the normalized loss.	required

prepare_gradients(context) Source

Unscale gradients before optional clipping.

Parameters:

Name	Type	Description	Default
context	EventContext	Gradient context containing the optimizer.	required

run_optimizer_step(context) Source

Execute the scaler-aware optimizer step.

Parameters:

Name	Type	Description	Default
context	EventContext	Optimizer context containing the optimizer.	required

state_dict() Source

Return precision and scaler state.

load_state_dict(state_dict) Source

Restore scaler state for a compatible precision configuration.

Parameters:

Name	Type	Description	Default
state_dict	dict[str, Any]	State returned by :meth:state_dict.	required

ModelCheckpoint Source

Bases: Callback

Saves the model when a monitored metric improves.

Parameters:

Name	Type	Description	Default
phase	MonitorPhase | None	"train" or "validation". "val" is deprecated.	None
metric	str | None	The name of the metric to monitor	None
mode	Literal['min', 'max']	One of 'min' or 'max'. In 'min' mode, the model is saved when the monitored metric decreases. In 'max' mode, it is saved when the metric increases	'max'
save_dir	str	Directory to save the model checkpoints (defaults to current directory)	'.'
save_path	str | None	Filename for the saved model. If None, defaults to 'epochs-metric=value.pth'	None
save_top_k	int	Saves specified number of best models (defaults to 1)	1
save_weights_only	bool	Save only model weights instead of full Battery state.	False
stage	MonitorPhase | None	Deprecated keyword alias for phase.	None

Missing directories are created automatically. A .pth suffix is added only when save_path has no explicit suffix. Static templates gain an epoch field when save_top_k is greater than one to avoid overwriting retained weights. The {epoch} filename field uses the one-based epoch number from the event context, matching progress output.

Examples:

checkpoint = ModelCheckpoint(
    phase="validation",
    metric="accuracy",
    mode="max",
    save_path="best_model.pth"
)
battery = Battery(model=model, callbacks=[checkpoint])

best_model_path property

Returns the path of the best saved model.

best_score property

Returns the best score achieved by the monitored metric.

best_k_models property

Returns a dictionary of the top K saved models and their scores.

state_dict() Source

Return checkpoint ranking state for training resumption.

load_state_dict(state_dict) Source

Restore checkpoint ranking state.

Parameters:

Name	Type	Description	Default
state_dict	dict[str, object]	State returned by :meth:state_dict.	required

run_on_train_epoch_end(context) Source

Save model checkpoint after training epoch if metric improved.

Parameters:

Name	Type	Description	Default
context	EventContext	Event context containing training metrics and model.	required

run_on_validation_end(context) Source

Save model checkpoint after validation if metric improved.

Parameters:

Name	Type	Description	Default
context	EventContext	Event context containing validation metrics and model.	required

TerminateOnNonFinite Source

Bases: Callback

Raise when a workflow produces a selected NaN or infinite value.

Parameters:

Name	Type	Description	Default
check_loss	bool	Check training loss before backward and evaluation losses at step completion.	True
check_metrics	bool	Check named batch metrics and final phase metrics.	True

Raises:

Type	Description
TypeError	If either option is not a boolean.
ValueError	If both checks are disabled.
FloatingPointError	When a selected value is NaN or infinite.

state_dict() Source

Return the fixed non-finite check configuration.

load_state_dict(state_dict) Source

Validate checkpoint configuration without changing this callback.

Parameters:

Name	Type	Description	Default
state_dict	dict[str, Any]	Configuration stored by :meth:state_dict.	required

on_before_backward(context) Source

Check training loss before backward or optimizer execution.

Parameters:

Name	Type	Description	Default
context	EventContext	Training optimization context containing loss_tensor.	required

on_step_end(context) Source

Check evaluation loss and named metrics after a phase step.

Parameters:

Name	Type	Description	Default
context	EventContext	Completed train, validation, or test step context.	required

on_phase_end(context) Source

Check final aggregated loss and stateful or collected metrics.

Parameters:

Name	Type	Description	Default
context	EventContext	Completed train, validation, or test phase context.	required