pub struct CheckpointManager { /* private fields */ }Expand description
Gradient checkpointing manager
Implementations§
Source§impl CheckpointManager
impl CheckpointManager
Sourcepub fn new(config: CheckpointConfig) -> Result<Self, OptimizerError>
pub fn new(config: CheckpointConfig) -> Result<Self, OptimizerError>
Create a new checkpoint manager
Sourcepub fn should_checkpoint(&self) -> bool
pub fn should_checkpoint(&self) -> bool
Check if a checkpoint should be saved at the current step
Sourcepub fn save_checkpoint(
&mut self,
optimizer_state: Option<&[u8]>,
gradients: Option<&HashMap<String, Vec<f32>>>,
model_parameters: Option<&HashMap<String, Vec<f32>>>,
loss: Option<f32>,
learning_rate: f32,
epoch: Option<usize>,
) -> Result<PathBuf, OptimizerError>
pub fn save_checkpoint( &mut self, optimizer_state: Option<&[u8]>, gradients: Option<&HashMap<String, Vec<f32>>>, model_parameters: Option<&HashMap<String, Vec<f32>>>, loss: Option<f32>, learning_rate: f32, epoch: Option<usize>, ) -> Result<PathBuf, OptimizerError>
Save a checkpoint
Sourcepub fn load_checkpoint(
&self,
checkpoint_path: &Path,
) -> Result<Checkpoint, OptimizerError>
pub fn load_checkpoint( &self, checkpoint_path: &Path, ) -> Result<Checkpoint, OptimizerError>
Load a checkpoint
Sourcepub fn latest_checkpoint(&self) -> Option<&PathBuf>
pub fn latest_checkpoint(&self) -> Option<&PathBuf>
Get the latest checkpoint path
Sourcepub fn list_checkpoints(&self) -> Vec<&CheckpointMetadata>
pub fn list_checkpoints(&self) -> Vec<&CheckpointMetadata>
Get all checkpoint metadata
Sourcepub fn current_step(&self) -> usize
pub fn current_step(&self) -> usize
Get current step
Sourcepub fn resume_from_latest(
&mut self,
) -> Result<Option<Checkpoint>, OptimizerError>
pub fn resume_from_latest( &mut self, ) -> Result<Option<Checkpoint>, OptimizerError>
Resume from the latest checkpoint
Sourcepub fn cleanup(&mut self) -> Result<(), OptimizerError>
pub fn cleanup(&mut self) -> Result<(), OptimizerError>
Clean up checkpoints manually
Sourcepub fn statistics(&self) -> CheckpointStatistics
pub fn statistics(&self) -> CheckpointStatistics
Get checkpoint statistics
Auto Trait Implementations§
impl Freeze for CheckpointManager
impl RefUnwindSafe for CheckpointManager
impl Send for CheckpointManager
impl Sync for CheckpointManager
impl Unpin for CheckpointManager
impl UnsafeUnpin for CheckpointManager
impl UnwindSafe for CheckpointManager
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Mutably borrows from an owned value. Read more
impl<ST, DT> CastableFrom<ST, Initialized, Initialized> for DT
impl<ST, DT> CastableFrom<ST, Uninit, Uninit> for DT
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ
fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ
Converts
self into a Left variant of Either<Self, Self>
if into_left is true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
Converts
self into a Left variant of Either<Self, Self>
if into_left(&self) returns true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read more