Trainer API¶

Battery remains the public workflow facade. Its checkpoint, fitting, training, standalone validation, testing, prediction, and streaming-prediction methods are documented here even though their implementations are organized into focused private trainer modules. Applications should continue importing only Battery from torch_batteries or torch_batteries.trainer.

Use fit() for combined training and optional validation, train() for training-only work, and validate() for a required single validation pass.

`torch_batteries.trainer` ¶

Event-driven model training, evaluation, and prediction.

Public API¶

Battery — orchestrates device placement, steps, optimization, metrics, callbacks, checkpoints, testing, and prediction.
StepOutput — explicit loss, predictions, targets, and manual metrics returned from a train, validation, or test step.
FitResult — per-epoch training and optional validation histories.
TrainResult — per-epoch training histories with temporary validation compatibility fields.
ValidationResult — aggregate standalone validation loss and metrics.
TestResult — aggregate test loss and metrics.
PredictResult — collected or recursively concatenated prediction output.

`Battery` Source ¶

Bases: CheckpointMixin, TrainingMixin, EvaluationMixin, PredictionMixin

Run event-driven training, evaluation, and prediction for a PyTorch model.

Battery discovers model methods decorated with :func:~torch_batteries.charge and dispatches lifecycle events to the model and configured callbacks. It moves the model and batches to the selected device, aggregates losses and metrics, and can save complete resumable training state.

Parameters:

Name	Type	Description	Default
`model`	`Module`	Model containing the charged step methods. It is moved to `device`.	required
`device`	`str \| device`	Explicit PyTorch device or `"auto"`. Automatic selection prefers CUDA, then MPS, then CPU.	`'auto'`
`optimizer`	`Optimizer \| None`	Optimizer used by :meth:`train`. It is optional for testing and prediction.	`None`
`metrics`	`dict[str, Metric] \| None`	Named callable or stateful metrics. When metrics are configured, train, validation, and test steps must return :class:`StepOutput` with predictions and targets.	`None`
`callbacks`	`list \| None`	Ordered callback objects. Callback order is significant for provider-style optimization events.	`None`
`data_pack`	`DataPack \| None`	Optional event-driven dataset and DataLoader configuration. When attached, workflow loaders may be omitted.	`None`
`metric_error_policy`	`Literal['raise', 'warn']`	`"raise"` to propagate metric lifecycle exceptions. `"warn"` logs the failure and skips that metric for the phase.	`'raise'`

Note

Epoch values exposed through event contexts are one-based. Prediction output stays on its current device unless move_to_cpu=True is requested.

`model` `property` ¶

Get the model.

`device` `property` ¶

Get the device.

`data_pack` `property` ¶

Get the event-driven data configuration attached to this Battery.

`optimizer` `property` `writable` ¶

Get the optimizer.

`metrics` `property` `writable` ¶

Get the metrics dictionary.

`metric_error_policy` `property` `writable` ¶

Get the configured metric exception handling policy.

`stop_training` `property` `writable` ¶

Get the stop_training flag.

`save_checkpoint(path)` Source ¶

Save complete resumable training state atomically.

Parameters:

Name	Type	Description	Default
`path`	`str \| Path`	Destination checkpoint path. Parent directories are created.	required

Raises:

Type	Description
`OSError`	If the destination cannot be created or replaced.
`Exception`	If PyTorch serialization fails.

`load_checkpoint(path)` Source ¶

Restore a full checkpoint or raw model state.

Parameters:

Name	Type	Description	Default
`path`	`str \| Path`	Trusted checkpoint or model-state path.	required

Raises:

Type	Description
`ValueError`	If saved state is incompatible with this Battery.
`TypeError`	If the checkpoint structure is invalid.
`RuntimeError`	If strict PyTorch state restoration fails.

`train(train_loader=None, val_loader=None, epochs=1, verbose=1, *, resume_from=None, resume_epochs_mode='total')` Source ¶

Train with explicit loaders or the attached DataPack.

Validation through this method is deprecated. Use fit for combined training and validation. Calls without validation data do not warn.

Parameters:

Name	Type	Description	Default
`train_loader`	`DataLoader \| None`	Optional sized, non-empty training loader.	`None`
`val_loader`	`DataLoader \| None`	Deprecated. Optional validation loader for direct-loader compatibility. Use `fit` for validated training.	`None`
`epochs`	`int`	Positive epoch count or resume target.	`1`
`verbose`	`int`	`0` for silent, `1` for bars, or `2` for summaries.	`1`
`resume_from`	`str \| Path \| None`	Optional full checkpoint restored before data setup.	`None`
`resume_epochs_mode`	`str`	`"total"` or `"additional"`.	`'total'`

Returns:

Type	Description
`TrainResult`	Per-epoch loss and named metric histories.

`fit(train_loader=None, val_loader=None, epochs=1, verbose=1, *, resume_from=None, resume_epochs_mode='total')` Source ¶

Train with optional per-epoch validation.

Parameters:

Name	Type	Description	Default
`train_loader`	`DataLoader \| None`	Optional sized, non-empty training loader.	`None`
`val_loader`	`DataLoader \| None`	Optional validation loader for direct-loader mode.	`None`
`epochs`	`int`	Positive epoch count or resume target.	`1`
`verbose`	`int`	`0` for silent, `1` for bars, or `2` for summaries.	`1`
`resume_from`	`str \| Path \| None`	Optional full checkpoint restored before data setup.	`None`
`resume_epochs_mode`	`str`	`"total"` or `"additional"`.	`'total'`

Returns:

Type	Description
`FitResult`	Per-epoch training histories and optional validation histories. Validation
`FitResult`	histories are empty when validation data is unavailable.

`validate(val_loader=None, verbose=1)` Source ¶

Run one standalone validation pass.

Parameters:

Name	Type	Description	Default
`val_loader`	`DataLoader \| None`	Optional sized, non-empty validation loader. When omitted, the attached DataPack must provide validation data.	`None`
`verbose`	`int`	`0` for silent, `1` for a bar, or `2` for a summary.	`1`

Returns:

Type	Description
`ValidationResult`	Aggregate validation loss and optional named validation metrics.

`test(test_loader=None, verbose=1, *, dataset=None)` Source ¶

test(test_loader: DataLoader, verbose: int = 1, *, dataset: None = None) -> TestResult

test(test_loader: None = None, verbose: int = 1, *, dataset: str) -> TestResult

test(test_loader: DataLoader | None = None, verbose: int = 1, *, dataset: None = None) -> TestResult | dict[str, TestResult]

Evaluate an explicit or DataPack-provided test dataset.

Parameters:

Name	Type	Description	Default
`test_loader`	`DataLoader \| None`	Optional sized, non-empty test loader.	`None`
`verbose`	`int`	`0` for silent, `1` for a bar, or `2` for a summary.	`1`
`dataset`	`str \| None`	Optional DataPack test dataset name.	`None`

Returns:

Type	Description
`TestResult \| dict[str, TestResult]`	One result or a mapping of named DataPack results.

`predict(data_loader=None, verbose=1, *, move_to_cpu=False, concatenate=False, dataset=None)` Source ¶

predict(data_loader: DataLoader, verbose: int = 1, *, move_to_cpu: bool = False, concatenate: bool = False, dataset: None = None) -> PredictResult

predict(data_loader: None = None, verbose: int = 1, *, move_to_cpu: bool = False, concatenate: bool = False, dataset: str) -> PredictResult

predict(data_loader: DataLoader | None = None, verbose: int = 1, *, move_to_cpu: bool = False, concatenate: bool = False, dataset: None = None) -> PredictResult | dict[str, PredictResult]

Collect predictions from an explicit or DataPack loader.

Parameters:

Name	Type	Description	Default
`data_loader`	`DataLoader \| None`	Optional sized, non-empty prediction loader.	`None`
`verbose`	`int`	`0` for silent, `1` for a bar, or `2` for a summary.	`1`
`move_to_cpu`	`bool`	Detach tensor outputs and move them to CPU.	`False`
`concatenate`	`bool`	Concatenate compatible batch output structures.	`False`
`dataset`	`str \| None`	Optional DataPack prediction dataset name.	`None`

Returns:

Type	Description
`PredictResult \| dict[str, PredictResult]`	One prediction result or a mapping of named results.

`predict_iter(data_loader=None, verbose=1, *, move_to_cpu=False, dataset=None)` Source ¶

Yield predictions from an explicit or DataPack loader.

Parameters:

Name	Type	Description	Default
`data_loader`	`DataLoader \| None`	Optional sized, non-empty prediction loader.	`None`
`verbose`	`int`	`0` for silent, `1` for a bar, or `2` for a summary.	`1`
`move_to_cpu`	`bool`	Detach tensor outputs and move them to CPU.	`False`
`dataset`	`str \| None`	Optional DataPack prediction dataset name.	`None`

Yields:

Type	Description
`Generator[Any]`	One prediction-step output at a time.

`FitResult` Source ¶

Bases: TypedDict

Result from fitting.

Attributes:

Name	Type	Description
`train_loss`	`list[float]`	Average training loss for every completed epoch.
`val_loss`	`list[float]`	Average validation loss for every completed epoch, or an empty list when validation data was unavailable.
`train_metrics`	`dict[str, list[float]]`	Named training metric histories.
`val_metrics`	`dict[str, list[float]]`	Named validation metric histories, or an empty mapping when validation data was unavailable.

`PredictResult` Source ¶

Bases: TypedDict

Collected prediction output.

Attributes:

Name	Type	Description
`predictions`	`Any`	Batch list or recursively concatenated prediction structure.

`StepOutput` Source `dataclass` ¶

Explicit output from a training, validation, or test step.

Parameters:

Name	Type	Description	Default
`loss`	`Tensor`	Scalar tensor used for reporting and, during training, backward.	required
`predictions`	`Tensor \| None`	Outputs produced by the same forward pass as `loss`.	`None`
`targets`	`Tensor \| None`	Ground-truth tensors corresponding to `predictions`.	`None`
`metrics`	`dict[str, float \| Tensor]`	Named scalar tensors or numeric values calculated by the step.	`dict()`

`TestResult` Source ¶

Bases: TypedDict

Result from testing.

Attributes:

Name	Type	Description
`test_loss`	`float`	Average test loss across all batches.
`test_metrics`	`dict[str, float]`	Named average test metrics.

`TrainResult` Source ¶

Bases: TypedDict

Result from training.

Attributes:

Name	Type	Description
`train_loss`	`list[float]`	Average training loss for every completed epoch.
`val_loss`	`list[float]`	Deprecated validation loss history retained while `train()` temporarily supports validation. Use `FitResult` instead.
`train_metrics`	`dict[str, list[float]]`	Named training metric histories.
`val_metrics`	`dict[str, list[float]]`	Deprecated validation metric histories retained while `train()` temporarily supports validation. Use `FitResult` instead.

`ValidationResult` Source ¶

Bases: TypedDict

Result from one standalone validation pass.

Attributes:

Name	Type	Description
`val_loss`	`float`	Average validation loss.
`val_metrics`	`dict[str, float]`	Named validation metrics when any are produced.

Trainer API¶

torch_batteries.trainer ¶

Public API¶

Battery Source ¶

model property ¶

device property ¶

data_pack property ¶

optimizer property writable ¶

metrics property writable ¶

metric_error_policy property writable ¶

stop_training property writable ¶

save_checkpoint(path) Source ¶

load_checkpoint(path) Source ¶

train(train_loader=None, val_loader=None, epochs=1, verbose=1, *, resume_from=None, resume_epochs_mode='total') Source ¶

fit(train_loader=None, val_loader=None, epochs=1, verbose=1, *, resume_from=None, resume_epochs_mode='total') Source ¶

validate(val_loader=None, verbose=1) Source ¶

test(test_loader=None, verbose=1, *, dataset=None) Source ¶

predict(data_loader=None, verbose=1, *, move_to_cpu=False, concatenate=False, dataset=None) Source ¶

predict_iter(data_loader=None, verbose=1, *, move_to_cpu=False, dataset=None) Source ¶

FitResult Source ¶

PredictResult Source ¶

StepOutput Source dataclass ¶

TestResult Source ¶

TrainResult Source ¶

ValidationResult Source ¶

`torch_batteries.trainer` ¶

`Battery` Source ¶

`model` `property` ¶

`device` `property` ¶

`data_pack` `property` ¶

`optimizer` `property` `writable` ¶

`metrics` `property` `writable` ¶

`metric_error_policy` `property` `writable` ¶

`stop_training` `property` `writable` ¶

`save_checkpoint(path)` Source ¶

`load_checkpoint(path)` Source ¶

`train(train_loader=None, val_loader=None, epochs=1, verbose=1, *, resume_from=None, resume_epochs_mode='total')` Source ¶

`fit(train_loader=None, val_loader=None, epochs=1, verbose=1, *, resume_from=None, resume_epochs_mode='total')` Source ¶

`validate(val_loader=None, verbose=1)` Source ¶

`test(test_loader=None, verbose=1, *, dataset=None)` Source ¶

`predict(data_loader=None, verbose=1, *, move_to_cpu=False, concatenate=False, dataset=None)` Source ¶

`predict_iter(data_loader=None, verbose=1, *, move_to_cpu=False, dataset=None)` Source ¶

`FitResult` Source ¶

`PredictResult` Source ¶

`StepOutput` Source `dataclass` ¶

`TestResult` Source ¶

`TrainResult` Source ¶

`ValidationResult` Source ¶