CheckpointManager

class braintools.trainer.CheckpointManager(dirpath, max_to_keep=5, filename_template=None, save_best_only=False, monitor=None, mode='min', verbose=True)

Manager for saving and loading checkpoints.

Provides functionality for: - Saving checkpoints with automatic naming - Tracking and keeping only the best N checkpoints - Loading checkpoints by epoch, step, or best metric - Automatic cleanup of old checkpoints

Parameters:

• dirpath (str) – Directory to save checkpoints.

• max_to_keep (int) – Maximum number of checkpoints to keep. Set to -1 to keep all.

• filename_template (str | None) – Template for checkpoint filenames. Available variables: {epoch}, {step}, {metric}, {timestamp}

• save_best_only (bool) – Only save when the monitored metric improves.

• monitor (str | None) – Metric to monitor for best checkpoint selection.

• mode (str) – One of ‘min’ or ‘max’. In ‘min’ mode, lower is better.

• verbose (bool) – Whether to print checkpoint operations.

Examples

>>> manager = CheckpointManager(
...     'checkpoints/',
...     max_to_keep=3,
...     monitor='val_loss',
...     mode='min',
... )
>>> manager.save(model, optimizer, epoch=10, metrics={'val_loss': 0.5})
>>> state = manager.load_best()

property best: str | None

Path to the best checkpoint.

property best_score: float | None

Best score achieved.

property checkpoints: List[Dict[str, Any]]

List of tracked checkpoints.

clear()

Remove all checkpoints.

property latest: str | None

Path to the latest checkpoint.

load(filepath=None, model=None, optimizer=None)

Load a checkpoint.

Parameters:

• filepath (str | None) – Path to checkpoint. If None, loads the latest.

• model (Any | None) – Model to load state into.

• optimizer (Any | None) – Optimizer to load state into.

Returns:

Loaded checkpoint state.

Return type:

Dict[str, Any]

Examples

>>> state = manager.load('checkpoint.ckpt', model=model)

load_best(model=None, optimizer=None)

Load the best checkpoint based on monitored metric.

Parameters:

• model (Any | None) – Model to load state into.

• optimizer (Any | None) – Optimizer to load state into.

Returns:

Loaded checkpoint state.

Return type:

Dict[str, Any]

load_latest(model=None, optimizer=None)

Load the most recent checkpoint.

Parameters:

• model (Any | None) – Model to load state into.

• optimizer (Any | None) – Optimizer to load state into.

Returns:

Loaded checkpoint state.

Return type:

Dict[str, Any]

save(model, optimizer=None, epoch=0, step=None, metrics=None, extra_state=None)

Save a checkpoint.

Parameters:

• model (Any) – Model to save.

• optimizer (Any | None) – Optimizer to save.

• epoch (int) – Current epoch.

• step (int | None) – Current step.

• metrics (Dict[str, float] | None) – Current metrics.

• extra_state (Dict[str, Any] | None) – Additional state to save.

Returns:

Path to saved checkpoint, or None if not saved (save_best_only).

Return type:

str | None

Examples

>>> path = manager.save(model, optimizer, epoch=10)

save_best_as(filepath)

Copy the best checkpoint to a new location.

Parameters:

filepath (str) – Destination path.