Struct GradientCheckpointer 

pub struct GradientCheckpointer<A: Float, D: Dimension> { /* private fields */ }

Gradient checkpointing manager

Implementations

impl<A: Float + ScalarOperand + Debug, D: Dimension + Send + Sync> GradientCheckpointer<A, D>

pub fn new(strategy: CheckpointStrategy) -> Self

Create a new gradient checkpointer

pub fn set_max_depth(&mut self, depth: usize)

Set the maximum computation depth

pub fn set_enabled(&mut self, enabled: bool)

Enable or disable checkpointing

pub fn should_checkpoint(&self, depth: usize) -> bool

Check if we should checkpoint at the current depth

pub fn store_checkpoint(&mut self, depth: usize, activation: Array<A, D>)

Store a checkpoint

pub fn get_checkpoint(&self, depth: usize) -> Option<&Array<A, D>>

Retrieve a checkpoint

pub fn remove_checkpoint(&mut self, depth: usize) -> Option<Array<A, D>>

Remove a checkpoint to free memory

pub fn clear_checkpoints(&mut self)

Clear all checkpoints

pub fn memory_usage(&self) -> MemoryUsage

Get memory usage information

pub fn optimize_strategy(&mut self, target_memoryusage: f64)

Optimize checkpointing strategy based on memory usage

pub fn checkpointed_forward<F, Output>( &mut self, depth: usize, input: &Array<A, D>, forward_fn: F, ) -> Result<(Output, Option<Array<A, D>>)>

where F: FnOnce(&Array<A, D>) -> Result<(Output, Array<A, D>)>,

Execute a checkpointed computation

pub fn recompute_from_checkpoint<F>( &self, start_depth: usize, target_depth: usize, recompute_fn: F, ) -> Result<Array<A, D>>

where F: Fn(usize, &Array<A, D>) -> Result<Array<A, D>>,

Recompute activations from checkpoint

Trait Implementations

impl<A: Debug + Float, D: Debug + Dimension> Debug for GradientCheckpointer<A, D>

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.