Skip to main content

Module checkpoint

Module checkpoint 

Source
Expand description

Optimizer checkpointing: save/load optimizer state (momentum buffers, step counts, etc.)

Provides OptimizerCheckpoint, CheckpointManager, and LossTracker for persisting and restoring optimizer state during training.

Structs§

CheckpointManager
Manages writing and reading checkpoints under a directory.
CheckpointMetadata
Metadata attached to a checkpoint (loss values, extra annotations).
LossTracker
Rolling-window tracker for scalar loss values recorded during training.
OptimizerCheckpoint
Serialisable snapshot of an optimizer’s full training state.
ParamState
Per-parameter optimizer state (moment vectors, step counter, shape).

Enums§

CheckpointError
Errors that can occur while managing checkpoints.
CheckpointFormat
Wire format used when writing checkpoints to disk.

Functions§

deserialize_checkpoint
Deserialize an OptimizerCheckpoint from bytes using the chosen format.
serialize_checkpoint
Serialize an OptimizerCheckpoint to bytes using the chosen format.