Struct MemoryOptimizer 

pub struct MemoryOptimizer { /* private fields */ }

Memory optimization manager

Implementations

impl MemoryOptimizer

pub fn new(config: MemoryOptimizationConfig) -> Self

pub fn create_checkpoint( &mut self, layer_index: usize, activations: Vec<Tensor>, ) -> Result<()>

Create a checkpoint for gradient checkpointing

pub fn get_checkpoint_activations( &self, layer_index: usize, ) -> Option<Vec<Tensor>>

Retrieve activations from checkpoint

pub fn offload_to_cpu(&mut self, name: String, tensor: Tensor) -> Result<()>

Offload tensor to CPU to free GPU memory

pub fn retrieve_from_cpu(&mut self, name: &str) -> Option<Tensor>

Retrieve tensor from CPU offloading

pub fn update_memory_usage(&self, delta: isize)

Update memory usage tracking

pub fn get_memory_usage(&self) -> usize

Get current memory usage

pub fn should_cleanup(&self) -> bool

Check if memory cleanup is needed

pub fn cleanup(&mut self) -> Result<usize>

Perform memory cleanup

pub fn store_for_rematerialization( &mut self, key: String, tensors: Vec<Tensor>, ) -> Result<()>

Store tensor for rematerialization

pub fn retrieve_for_rematerialization( &mut self, key: &str, ) -> Option<Vec<Tensor>>

Retrieve tensor for rematerialization

pub fn get_stats(&self) -> MemoryOptimizationStats

Get memory optimization statistics