extensible.extensions.checkpoints#

Classes

CheckpointBase(path)

CheckpointLoader([path, ckpt_spec, permissive])

Loads weights and fixtures at the pre_run stage of a training or evaluation run.

CheckpointSaver([path, saved_fixtures, ...])

Saves checkpoints containing weights and fixtures at in the post_train_step_epoch stage.
class extensible.extensions.checkpoints.CheckpointSaver(path=None, saved_fixtures: Iterable[str] | None = ('epoch_num',), saved_attribs: Iterable[str] | None = None)#

Bases: CheckpointBase, Extension

Saves checkpoints containing weights and fixtures at in the post_train_step_epoch stage.

Parameters:

  • path – The path where all values will be saved.

  • saved_fixtures – The fixtures to save. By default, only 'epoch_num' is saved.

  • saved_attribs – If not specified, all train manager attributes with a state_dict attribute will be saved. This usually includes the model, optimizer and learning rate scheduler.

a checkpoint to load. When training, set to True to continue training if any checkpoints exist (training will fail with the default value False if checkpoints exist).

post_train_step_epoch(train_manager, fixtures, epoch_num)#

Save a checkpoint

class extensible.extensions.checkpoints.CheckpointLoader(path=None, ckpt_spec: int | str | Path = -1, permissive=True)#

Bases: CheckpointBase, Extension

Loads weights and fixtures at the pre_run stage of a training or evaluation run.

By default, the last checkpoint is loaded, if any.

Parameters:

  • path – The path where checkpoints were saved.

  • ckpt_spec – A checkpoint file path or epoch number (or -1 for the latest checkpoint) specifying a checkpoint to load weights and fixtures from.

  • permissive – If ckpt_spec=-1 and there are no checkpoints, do not raise an error.

classmethod load_weights_and_fixtures(ckpt_file, train_manager, fixtures: FixturesDict | None)#

Loads the weights and sets the fixtures from the specified checkpoint file.

pre_run(train_manager, fixtures)#

Loads the weights and fixtures at a particular checkpoint in preparation for standalone evaluations or for finetuning.

Stages:

load_checkpoint:

fixtures:

load_checkpoint: bool