Struct GradientCheckpointing

Source

pub struct GradientCheckpointing<F: Float + Debug> { /* private fields */ }

Expand description

Gradient checkpointing implementation for memory-efficient training

Implementations§

Source §

impl<F: Float + Debug + Clone + 'static + ScalarOperand> GradientCheckpointing<F>

Source

pub fn new(memory_threshold_mb: f64) -> Self

Create a new gradient checkpointing manager

Examples found in repository ?

examples/memory_efficient_example.rs (line 88)

84fn demo_gradient_checkpointing() -> Result<()> {
85    println!("\n📊 Gradient Checkpointing Demo");
86    println!("------------------------------");
87
88    let mut checkpointing = GradientCheckpointing::<f64>::new(100.0); // 100MB threshold
89
90    // Set up checkpoint layers
91    checkpointing.add_checkpoint_layer("conv1".to_string());
92    checkpointing.add_checkpoint_layer("conv3".to_string());
93    checkpointing.add_checkpoint_layer("fc1".to_string());
94
95    println!("Storing activations at checkpoints...");
96
97    // Simulate storing activations during forward pass
98    let conv1_activation = Array3::from_elem((32, 64, 64), 0.5).into_dyn(); // Batch=32, 64x64 feature maps
99    let conv3_activation = Array3::from_elem((32, 128, 32), 0.3).into_dyn(); // Reduced spatial size
100    let fc1_activation = Array2::from_elem((32, 512), 0.2).into_dyn(); // Fully connected
101
102    checkpointing.store_checkpoint("conv1", conv1_activation)?;
103    checkpointing.store_checkpoint("conv3", conv3_activation)?;
104    checkpointing.store_checkpoint("fc1", fc1_activation)?;
105
106    let usage = checkpointing.get_memory_usage();
107    println!("Memory usage after checkpointing:");
108    print_memory_usage(&usage);
109
110    // Simulate retrieving checkpoints during backward pass
111    println!("Retrieving checkpoints for gradient computation...");
112    if let Some(checkpoint) = checkpointing.get_checkpoint("conv1") {
113        println!("Retrieved conv1 checkpoint: shape {:?}", checkpoint.shape());
114    }
115
116    // Clear checkpoints to free memory
117    println!("Clearing checkpoints...");
118    checkpointing.clear_checkpoints();
119
120    let usage = checkpointing.get_memory_usage();
121    println!("Memory usage after clearing:");
122    print_memory_usage(&usage);
123
124    Ok(())
125}

Source

pub fn add_checkpoint_layer(&mut self, layer_name: String)

Add a layer as a checkpoint point

Examples found in repository ?

examples/memory_efficient_example.rs (line 91)

84fn demo_gradient_checkpointing() -> Result<()> {
85    println!("\n📊 Gradient Checkpointing Demo");
86    println!("------------------------------");
87
88    let mut checkpointing = GradientCheckpointing::<f64>::new(100.0); // 100MB threshold
89
90    // Set up checkpoint layers
91    checkpointing.add_checkpoint_layer("conv1".to_string());
92    checkpointing.add_checkpoint_layer("conv3".to_string());
93    checkpointing.add_checkpoint_layer("fc1".to_string());
94
95    println!("Storing activations at checkpoints...");
96
97    // Simulate storing activations during forward pass
98    let conv1_activation = Array3::from_elem((32, 64, 64), 0.5).into_dyn(); // Batch=32, 64x64 feature maps
99    let conv3_activation = Array3::from_elem((32, 128, 32), 0.3).into_dyn(); // Reduced spatial size
100    let fc1_activation = Array2::from_elem((32, 512), 0.2).into_dyn(); // Fully connected
101
102    checkpointing.store_checkpoint("conv1", conv1_activation)?;
103    checkpointing.store_checkpoint("conv3", conv3_activation)?;
104    checkpointing.store_checkpoint("fc1", fc1_activation)?;
105
106    let usage = checkpointing.get_memory_usage();
107    println!("Memory usage after checkpointing:");
108    print_memory_usage(&usage);
109
110    // Simulate retrieving checkpoints during backward pass
111    println!("Retrieving checkpoints for gradient computation...");
112    if let Some(checkpoint) = checkpointing.get_checkpoint("conv1") {
113        println!("Retrieved conv1 checkpoint: shape {:?}", checkpoint.shape());
114    }
115
116    // Clear checkpoints to free memory
117    println!("Clearing checkpoints...");
118    checkpointing.clear_checkpoints();
119
120    let usage = checkpointing.get_memory_usage();
121    println!("Memory usage after clearing:");
122    print_memory_usage(&usage);
123
124    Ok(())
125}

Source

pub fn store_checkpoint( &mut self, layer_name: &str, activation: ArrayD<F>, ) -> Result<()>

Store activation at a checkpoint

Examples found in repository ?

examples/memory_efficient_example.rs (line 102)

84fn demo_gradient_checkpointing() -> Result<()> {
85    println!("\n📊 Gradient Checkpointing Demo");
86    println!("------------------------------");
87
88    let mut checkpointing = GradientCheckpointing::<f64>::new(100.0); // 100MB threshold
89
90    // Set up checkpoint layers
91    checkpointing.add_checkpoint_layer("conv1".to_string());
92    checkpointing.add_checkpoint_layer("conv3".to_string());
93    checkpointing.add_checkpoint_layer("fc1".to_string());
94
95    println!("Storing activations at checkpoints...");
96
97    // Simulate storing activations during forward pass
98    let conv1_activation = Array3::from_elem((32, 64, 64), 0.5).into_dyn(); // Batch=32, 64x64 feature maps
99    let conv3_activation = Array3::from_elem((32, 128, 32), 0.3).into_dyn(); // Reduced spatial size
100    let fc1_activation = Array2::from_elem((32, 512), 0.2).into_dyn(); // Fully connected
101
102    checkpointing.store_checkpoint("conv1", conv1_activation)?;
103    checkpointing.store_checkpoint("conv3", conv3_activation)?;
104    checkpointing.store_checkpoint("fc1", fc1_activation)?;
105
106    let usage = checkpointing.get_memory_usage();
107    println!("Memory usage after checkpointing:");
108    print_memory_usage(&usage);
109
110    // Simulate retrieving checkpoints during backward pass
111    println!("Retrieving checkpoints for gradient computation...");
112    if let Some(checkpoint) = checkpointing.get_checkpoint("conv1") {
113        println!("Retrieved conv1 checkpoint: shape {:?}", checkpoint.shape());
114    }
115
116    // Clear checkpoints to free memory
117    println!("Clearing checkpoints...");
118    checkpointing.clear_checkpoints();
119
120    let usage = checkpointing.get_memory_usage();
121    println!("Memory usage after clearing:");
122    print_memory_usage(&usage);
123
124    Ok(())
125}

Source

pub fn get_checkpoint(&self, layer_name: &str) -> Option<&ArrayD<F>>

Retrieve activation from checkpoint

Examples found in repository ?

examples/memory_efficient_example.rs (line 112)

84fn demo_gradient_checkpointing() -> Result<()> {
85    println!("\n📊 Gradient Checkpointing Demo");
86    println!("------------------------------");
87
88    let mut checkpointing = GradientCheckpointing::<f64>::new(100.0); // 100MB threshold
89
90    // Set up checkpoint layers
91    checkpointing.add_checkpoint_layer("conv1".to_string());
92    checkpointing.add_checkpoint_layer("conv3".to_string());
93    checkpointing.add_checkpoint_layer("fc1".to_string());
94
95    println!("Storing activations at checkpoints...");
96
97    // Simulate storing activations during forward pass
98    let conv1_activation = Array3::from_elem((32, 64, 64), 0.5).into_dyn(); // Batch=32, 64x64 feature maps
99    let conv3_activation = Array3::from_elem((32, 128, 32), 0.3).into_dyn(); // Reduced spatial size
100    let fc1_activation = Array2::from_elem((32, 512), 0.2).into_dyn(); // Fully connected
101
102    checkpointing.store_checkpoint("conv1", conv1_activation)?;
103    checkpointing.store_checkpoint("conv3", conv3_activation)?;
104    checkpointing.store_checkpoint("fc1", fc1_activation)?;
105
106    let usage = checkpointing.get_memory_usage();
107    println!("Memory usage after checkpointing:");
108    print_memory_usage(&usage);
109
110    // Simulate retrieving checkpoints during backward pass
111    println!("Retrieving checkpoints for gradient computation...");
112    if let Some(checkpoint) = checkpointing.get_checkpoint("conv1") {
113        println!("Retrieved conv1 checkpoint: shape {:?}", checkpoint.shape());
114    }
115
116    // Clear checkpoints to free memory
117    println!("Clearing checkpoints...");
118    checkpointing.clear_checkpoints();
119
120    let usage = checkpointing.get_memory_usage();
121    println!("Memory usage after clearing:");
122    print_memory_usage(&usage);
123
124    Ok(())
125}

Source

pub fn clear_checkpoints(&mut self)

Clear checkpoints to free memory

Examples found in repository ?

examples/memory_efficient_example.rs (line 118)

84fn demo_gradient_checkpointing() -> Result<()> {
85    println!("\n📊 Gradient Checkpointing Demo");
86    println!("------------------------------");
87
88    let mut checkpointing = GradientCheckpointing::<f64>::new(100.0); // 100MB threshold
89
90    // Set up checkpoint layers
91    checkpointing.add_checkpoint_layer("conv1".to_string());
92    checkpointing.add_checkpoint_layer("conv3".to_string());
93    checkpointing.add_checkpoint_layer("fc1".to_string());
94
95    println!("Storing activations at checkpoints...");
96
97    // Simulate storing activations during forward pass
98    let conv1_activation = Array3::from_elem((32, 64, 64), 0.5).into_dyn(); // Batch=32, 64x64 feature maps
99    let conv3_activation = Array3::from_elem((32, 128, 32), 0.3).into_dyn(); // Reduced spatial size
100    let fc1_activation = Array2::from_elem((32, 512), 0.2).into_dyn(); // Fully connected
101
102    checkpointing.store_checkpoint("conv1", conv1_activation)?;
103    checkpointing.store_checkpoint("conv3", conv3_activation)?;
104    checkpointing.store_checkpoint("fc1", fc1_activation)?;
105
106    let usage = checkpointing.get_memory_usage();
107    println!("Memory usage after checkpointing:");
108    print_memory_usage(&usage);
109
110    // Simulate retrieving checkpoints during backward pass
111    println!("Retrieving checkpoints for gradient computation...");
112    if let Some(checkpoint) = checkpointing.get_checkpoint("conv1") {
113        println!("Retrieved conv1 checkpoint: shape {:?}", checkpoint.shape());
114    }
115
116    // Clear checkpoints to free memory
117    println!("Clearing checkpoints...");
118    checkpointing.clear_checkpoints();
119
120    let usage = checkpointing.get_memory_usage();
121    println!("Memory usage after clearing:");
122    print_memory_usage(&usage);
123
124    Ok(())
125}

Source

pub fn get_memory_usage(&self) -> MemoryUsage

Get current memory usage

Examples found in repository ?

examples/memory_efficient_example.rs (line 106)

84fn demo_gradient_checkpointing() -> Result<()> {
85    println!("\n📊 Gradient Checkpointing Demo");
86    println!("------------------------------");
87
88    let mut checkpointing = GradientCheckpointing::<f64>::new(100.0); // 100MB threshold
89
90    // Set up checkpoint layers
91    checkpointing.add_checkpoint_layer("conv1".to_string());
92    checkpointing.add_checkpoint_layer("conv3".to_string());
93    checkpointing.add_checkpoint_layer("fc1".to_string());
94
95    println!("Storing activations at checkpoints...");
96
97    // Simulate storing activations during forward pass
98    let conv1_activation = Array3::from_elem((32, 64, 64), 0.5).into_dyn(); // Batch=32, 64x64 feature maps
99    let conv3_activation = Array3::from_elem((32, 128, 32), 0.3).into_dyn(); // Reduced spatial size
100    let fc1_activation = Array2::from_elem((32, 512), 0.2).into_dyn(); // Fully connected
101
102    checkpointing.store_checkpoint("conv1", conv1_activation)?;
103    checkpointing.store_checkpoint("conv3", conv3_activation)?;
104    checkpointing.store_checkpoint("fc1", fc1_activation)?;
105
106    let usage = checkpointing.get_memory_usage();
107    println!("Memory usage after checkpointing:");
108    print_memory_usage(&usage);
109
110    // Simulate retrieving checkpoints during backward pass
111    println!("Retrieving checkpoints for gradient computation...");
112    if let Some(checkpoint) = checkpointing.get_checkpoint("conv1") {
113        println!("Retrieved conv1 checkpoint: shape {:?}", checkpoint.shape());
114    }
115
116    // Clear checkpoints to free memory
117    println!("Clearing checkpoints...");
118    checkpointing.clear_checkpoints();
119
120    let usage = checkpointing.get_memory_usage();
121    println!("Memory usage after clearing:");
122    print_memory_usage(&usage);
123
124    Ok(())
125}

Source

pub fn recompute_from_checkpoint<L>( &self, layers: &[L], start_layer: &str, _target_layer: &str, _input: &ArrayD<F>, ) -> Result<ArrayD<F>>
where L: Layer<F>,

Recompute forward pass from last checkpoint

Auto Trait Implementations§

§

impl<F> UnwindSafe for GradientCheckpointing<F>
where F: RefUnwindSafe,

Blanket Implementations§

Source §

impl<T> Any for T
where T: 'static + ?Sized,

Source §

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more

Source §

impl<T> Borrow<T> for T
where T: ?Sized,

Source §

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more

Source §

impl<T> BorrowMut<T> for T
where T: ?Sized,

Source §

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more

Source §

impl<T> From<T> for T

Source §

fn from(t: T) -> T

Returns the argument unchanged.

Source §

impl<T, U> Into for T
where U: From<T>,

Source §

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

Source §

impl<T> IntoEither for T

Source §

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 more

Source §

fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
where F: FnOnce(&Self) -> bool,

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

Source §