entrenar/prune/schedule/
types.rs

1//! Type definitions for pruning schedules.
2
3use serde::{Deserialize, Serialize};
4
5/// Pruning schedule defining when sparsity increases during training.
6///
7/// # Variants
8///
9/// - `OneShot`: All pruning happens at a single step
10/// - `Gradual`: Linear interpolation between initial and final sparsity
11/// - `Cubic`: Cubic polynomial schedule for smoother transitions
12///
13/// # Example
14///
15/// ```
16/// use entrenar::prune::PruningSchedule;
17///
18/// // One-shot pruning at step 1000
19/// let oneshot = PruningSchedule::OneShot { step: 1000 };
20/// assert_eq!(oneshot.sparsity_at_step(500), 0.0);
21/// assert_eq!(oneshot.sparsity_at_step(1000), 1.0);
22///
23/// // Gradual pruning from steps 100-1000
24/// let gradual = PruningSchedule::Gradual {
25///     start_step: 100,
26///     end_step: 1000,
27///     initial_sparsity: 0.0,
28///     final_sparsity: 0.5,
29///     frequency: 10,
30/// };
31/// assert_eq!(gradual.sparsity_at_step(50), 0.0);
32/// assert_eq!(gradual.sparsity_at_step(1000), 0.5);
33/// ```
34#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
35#[serde(tag = "type", rename_all = "snake_case")]
36pub enum PruningSchedule {
37    /// Prune once at specified step.
38    OneShot {
39        /// Step at which to apply pruning.
40        step: usize,
41    },
42
43    /// Gradually increase sparsity over steps with linear interpolation.
44    Gradual {
45        /// Step to begin pruning.
46        start_step: usize,
47        /// Step at which final sparsity is reached.
48        end_step: usize,
49        /// Initial sparsity (typically 0.0).
50        initial_sparsity: f32,
51        /// Target final sparsity.
52        final_sparsity: f32,
53        /// Prune every N steps.
54        frequency: usize,
55    },
56
57    /// Cubic sparsity schedule (Zhu & Gupta, 2017).
58    ///
59    /// Formula: s_t = s_f * (1 - (1 - t/T)^3)
60    ///
61    /// This provides faster initial pruning that slows as it approaches
62    /// the target, giving the model more time to adapt.
63    Cubic {
64        /// Step to begin pruning.
65        start_step: usize,
66        /// Step at which final sparsity is reached.
67        end_step: usize,
68        /// Target final sparsity.
69        final_sparsity: f32,
70    },
71}
72
73impl Default for PruningSchedule {
74    fn default() -> Self {
75        PruningSchedule::OneShot { step: 0 }
76    }
77}
entrenar/prune/schedule/types.rs

entrenar/prune/schedule/
types.rs
