entrenar/prune/schedule/gradual/
mod.rs

1//! Gradual pruning schedule methods.
2
3#[cfg(test)]
4mod proptests;
5#[cfg(test)]
6mod tests;
7
8use super::PruningSchedule;
9
10impl PruningSchedule {
11    /// Compute the target sparsity at a given training step for Gradual schedule.
12    pub(super) fn gradual_sparsity_at_step(
13        start_step: usize,
14        end_step: usize,
15        initial_sparsity: f32,
16        final_sparsity: f32,
17        step: usize,
18    ) -> f32 {
19        if step < start_step {
20            initial_sparsity
21        } else if step >= end_step {
22            final_sparsity
23        } else {
24            let progress = (step - start_step) as f32 / (end_step - start_step) as f32;
25            initial_sparsity + progress * (final_sparsity - initial_sparsity)
26        }
27    }
28
29    /// Check if pruning should be applied at this step for Gradual schedule.
30    pub(super) fn gradual_should_prune_at_step(
31        start_step: usize,
32        end_step: usize,
33        frequency: usize,
34        step: usize,
35    ) -> bool {
36        if step < start_step || step > end_step {
37            return false;
38        }
39        if frequency == 0 {
40            return step == start_step;
41        }
42        (step - start_step).is_multiple_of(frequency)
43    }
44
45    /// Get the total number of pruning operations for Gradual schedule.
46    pub(super) fn gradual_num_pruning_steps(
47        start_step: usize,
48        end_step: usize,
49        frequency: usize,
50    ) -> usize {
51        // NonZeroUsize rather than `checked_div`: it keeps `end_step - start_step`
52        // inside the divisor-is-nonzero branch, exactly as the original
53        // `if frequency == 0 { 1 } else { ... }` did. `checked_div` would hoist
54        // the subtraction out and underflow on end_step < start_step, which
55        // `gradual_validate` rejects but this fn does not re-check.
56        match std::num::NonZeroUsize::new(frequency) {
57            None => 1,
58            Some(f) => (end_step - start_step) / f + 1,
59        }
60    }
61
62    /// Validate Gradual schedule.
63    pub(super) fn gradual_validate(
64        start_step: usize,
65        end_step: usize,
66        initial_sparsity: f32,
67        final_sparsity: f32,
68    ) -> Result<(), String> {
69        if end_step <= start_step {
70            return Err(format!(
71                "end_step ({end_step}) must be greater than start_step ({start_step})"
72            ));
73        }
74        if !(0.0..=1.0).contains(&initial_sparsity) {
75            return Err(format!(
76                "initial_sparsity ({initial_sparsity}) must be between 0.0 and 1.0"
77            ));
78        }
79        if !(0.0..=1.0).contains(&final_sparsity) {
80            return Err(format!("final_sparsity ({final_sparsity}) must be between 0.0 and 1.0"));
81        }
82        Ok(())
83    }
84
85    /// Check if Gradual pruning has completed.
86    pub(super) fn gradual_is_complete(end_step: usize, step: usize) -> bool {
87        step > end_step
88    }
89}
entrenar/prune/schedule/gradual/mod.rs

entrenar/prune/schedule/gradual/
mod.rs
