ModelCheckpoint#

class braintools.trainer.ModelCheckpoint(dirpath=None, filename=None, monitor='val_loss', mode='min', save_top_k=3, save_last=True, every_n_epochs=1, every_n_train_steps=None, save_on_train_epoch_end=True, verbose=False)#

Save model checkpoints based on monitored metric.

Parameters:
  • dirpath (str | None) – Directory to save checkpoints. Default: current directory.

  • filename (str | None) – Checkpoint filename template. Can include {epoch}, {step}, and metric names. Default: ‘checkpoint-{epoch:02d}-{val_loss:.4f}’

  • monitor (str) – Metric to monitor for best model selection.

  • mode (str) – One of ‘min’ or ‘max’. In ‘min’ mode, the lowest metric value is best.

  • save_top_k (int) – Number of best models to keep. -1 means keep all.

  • save_last (bool) – Whether to save the last checkpoint regardless of metric.

  • every_n_epochs (int) – Save checkpoint every n epochs.

  • every_n_train_steps (int | None) – Save checkpoint every n training steps.

  • save_on_train_epoch_end (bool) – Whether to run checkpointing at end of training epoch.

  • verbose (bool) – Whether to print checkpoint saving messages.

Examples

>>> checkpoint_callback = ModelCheckpoint(
...     dirpath='checkpoints/',
...     filename='model-{epoch:02d}-{val_loss:.4f}',
...     monitor='val_loss',
...     mode='min',
...     save_top_k=3,
... )
>>> trainer = Trainer(callbacks=[checkpoint_callback])
load_state_dict(state_dict)[source]#

Load state from state dict.

on_fit_end(trainer, module)[source]#

Save last checkpoint if configured.

on_train_batch_end(trainer, module, outputs, batch, batch_idx)[source]#

Check if we should save a checkpoint at step.

on_train_epoch_end(trainer, module)[source]#

Check if we should save a checkpoint at train-epoch end.

When validation runs this epoch we defer to on_validation_epoch_end so the monitored validation metric is available. When there is no validation we save here (subject to save_on_train_epoch_end).

on_validation_epoch_end(trainer, module)[source]#

Checkpoint at validation end, where the monitored metric exists.

This is the correct point to evaluate a validation metric such as monitor='val_loss' — running validation populates it just before this hook fires (unlike on_train_epoch_end, which runs first).

state_dict()[source]#

Return state dict for checkpointing.

Return type:

Dict[str, Any]