scvi.train.SaveCheckpoint

class scvi.train.SaveCheckpoint(dirpath=None, filename=None, monitor='validation_loss', load_best_on_end=False, **kwargs)

Bases: ModelCheckpoint

BETA Saves model checkpoints based on a monitored metric.

Inherits from ModelCheckpoint and modifies the default behavior to save the full model state instead of just the state dict. This is necessary for compatibility with BaseModelClass.

The best model save and best model score based on monitor can be accessed post-training with the best_model_path and best_model_score attributes, respectively.

Known issues:

• Does not set train_indices, validation_indices, and test_indices for checkpoints.

• Does not set history for checkpoints. This can be accessed in the final model however.

• Unsupported arguments: save_weights_only and save_last.

Parameters:

• dirpath (str | None (default: None)) – Base directory to save the model checkpoints. If None, defaults to a subdirectory in :attr:scvi.settings.logging_dir formatted with the current date, time, and monitor.

• filename (str | None (default: None)) – Name for the checkpoint directories, which can contain formatting options for auto-filling. If None, defaults to {epoch}-{step}-{monitor}.

• monitor (str (default: 'validation_loss')) – Metric to monitor for checkpointing.

• load_best_on_end (bool (default: False)) – If True, loads the best model state into the model at the end of training.

• **kwargs – Additional keyword arguments passed into the constructor for ModelCheckpoint.

Attributes table

Methods table

on_save_checkpoint(trainer, *args)	Saves the model state on Lightning checkpoint saves.
on_train_end(trainer, pl_module)	Loads the best model state into the model at the end of training.

Attributes

Methods

SaveCheckpoint.on_save_checkpoint(trainer, *args)

Saves the model state on Lightning checkpoint saves.

Return type:

None

SaveCheckpoint.on_train_end(trainer, pl_module)

Loads the best model state into the model at the end of training.

Return type:

None