t5x.checkpoint_utils package#

Checkpoint helper functions for managing checkpoints.

Supports marking checkpoints as pinned to exclude them from the checkpointer removal process.

class t5x.checkpoint_utils.CheckpointTypes(value)[source]#

An enumeration.

t5x.checkpoint_utils.detect_checkpoint_type(checkpoint_path, expected=None)[source]#

Returns the checkpoint type by reading the .checkpoint metadata file.

  • checkpoint_path – The path of the .checkpoint file.

  • expected – The expected checkpoint type. If the checkpoint type is not as expected, this function will log a warning but will not raise an error.


The checkpoint type.


Returns whether the checkpoint is pinned, and should NOT be removed.

t5x.checkpoint_utils.pin_checkpoint(ckpt_dir, txt='1')[source]#

Pin a checkpoint so it does not get deleted by the normal pruning process.

Creates a PINNED file in the checkpoint directory to indicate the checkpoint should be excluded from the deletion of old checkpoints.

  • ckpt_dir – The checkpoint step dir that is to be always kept.

  • txt – Text to be written into the checkpoints ALWAYS_KEEP me file.


Full path of the pinned checkpoint file.


Removes the checkpoint dir if it is not pinned.

t5x.checkpoint_utils.remove_dataset_checkpoint(ckpt_dir, train_ds_prefix)[source]#

Removes dataset checkpoints if the checkpoint is not pinned.


Removes the pinned status of the checkpoint so it is open for deletion.