Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,304 @@
|
||||
//! Cosine annealing learning rate scheduler
|
||||
//!
|
||||
//! Implements cosine annealing with warm restarts for efficient transformer training.
|
||||
|
||||
use crate::schedulers::LearningRateScheduler;
|
||||
use crate::{Result, TransformerError};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tracing::{debug, trace};
|
||||
|
||||
/// Cosine annealing learning rate scheduler
|
||||
///
|
||||
/// Decreases the learning rate following a cosine function, with optional warm restarts.
|
||||
/// This provides a smooth decay that often leads to better convergence than step-wise decay.
|
||||
///
|
||||
/// # Mathematical Foundation
|
||||
///
|
||||
/// For step t in the current cycle:
|
||||
/// - lr(t) = `min_lr` + (`base_lr` - `min_lr`) * 0.5 * (1 + cos(π * t / `T_max`))
|
||||
///
|
||||
/// Where `T_max` is the maximum number of steps in the current cycle.
|
||||
///
|
||||
/// # Benefits
|
||||
/// - Smooth learning rate decay
|
||||
/// - Can escape local minima with warm restarts
|
||||
/// - Widely used in state-of-the-art transformer training
|
||||
/// - Better final convergence than linear decay
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CosineAnnealingScheduler {
|
||||
/// Maximum learning rate at the beginning of each cycle
|
||||
base_lr: f64,
|
||||
/// Minimum learning rate at the end of each cycle
|
||||
min_lr: f64,
|
||||
/// Number of steps in each cycle
|
||||
t_max: usize,
|
||||
/// Current step count
|
||||
current_step: usize,
|
||||
/// Whether to use warm restarts
|
||||
use_restarts: bool,
|
||||
}
|
||||
|
||||
impl CosineAnnealingScheduler {
|
||||
/// Create a new cosine annealing scheduler
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `base_lr` - Maximum learning rate (must be positive)
|
||||
/// * `min_lr` - Minimum learning rate (must be non-negative and <= `base_lr`)
|
||||
/// * `t_max` - Number of steps in each cycle (must be positive)
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns error if parameters are invalid
|
||||
pub fn new(base_lr: f64, min_lr: f64, t_max: usize) -> Result<Self> {
|
||||
Self::with_restarts(base_lr, min_lr, t_max, true)
|
||||
}
|
||||
|
||||
/// Create a new cosine annealing scheduler with optional restarts
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `base_lr` - Maximum learning rate (must be positive)
|
||||
/// * `min_lr` - Minimum learning rate (must be non-negative and <= `base_lr`)
|
||||
/// * `t_max` - Number of steps in each cycle (must be positive)
|
||||
/// * `use_restarts` - Whether to restart the cycle after `T_max` steps
|
||||
///
|
||||
/// # Errors
|
||||
/// Returns error if parameters are invalid
|
||||
pub fn with_restarts(
|
||||
base_lr: f64,
|
||||
min_lr: f64,
|
||||
t_max: usize,
|
||||
use_restarts: bool,
|
||||
) -> Result<Self> {
|
||||
if base_lr <= 0.0 {
|
||||
return Err(TransformerError::generic(format!(
|
||||
"base_lr {base_lr} must be positive"
|
||||
)));
|
||||
}
|
||||
|
||||
if min_lr < 0.0 {
|
||||
return Err(TransformerError::generic(format!(
|
||||
"min_lr {min_lr} must be non-negative"
|
||||
)));
|
||||
}
|
||||
|
||||
if min_lr > base_lr {
|
||||
return Err(TransformerError::generic(format!(
|
||||
"min_lr {min_lr} must be <= base_lr {base_lr}"
|
||||
)));
|
||||
}
|
||||
|
||||
if t_max == 0 {
|
||||
return Err(TransformerError::generic(format!(
|
||||
"t_max {t_max} must be positive"
|
||||
)));
|
||||
}
|
||||
|
||||
debug!(
|
||||
"Creating cosine annealing scheduler: base_lr={}, min_lr={}, t_max={}, restarts={}",
|
||||
base_lr, min_lr, t_max, use_restarts
|
||||
);
|
||||
|
||||
Ok(Self {
|
||||
base_lr,
|
||||
min_lr,
|
||||
t_max,
|
||||
current_step: 0,
|
||||
use_restarts,
|
||||
})
|
||||
}
|
||||
|
||||
/// Get the minimum learning rate
|
||||
#[must_use]
|
||||
pub fn min_lr(&self) -> f64 {
|
||||
self.min_lr
|
||||
}
|
||||
|
||||
/// Get the cycle length
|
||||
#[must_use]
|
||||
pub fn t_max(&self) -> usize {
|
||||
self.t_max
|
||||
}
|
||||
|
||||
/// Check if using warm restarts
|
||||
#[must_use]
|
||||
pub fn use_restarts(&self) -> bool {
|
||||
self.use_restarts
|
||||
}
|
||||
|
||||
/// Get the current cycle number (0-indexed)
|
||||
///
|
||||
/// Each cycle contains `t_max` + 1 steps (from 0 to `t_max` inclusive).
|
||||
#[must_use]
|
||||
pub fn current_cycle(&self, step: usize) -> usize {
|
||||
if self.use_restarts {
|
||||
step / (self.t_max + 1)
|
||||
} else {
|
||||
0
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the step within the current cycle
|
||||
///
|
||||
/// For a cycle of `t_max` steps:
|
||||
/// - Steps 0 to `t_max` form one complete cycle (`t_max` + 1 total steps)
|
||||
/// - At step `t_max`, LR reaches `min_lr`
|
||||
/// - At step `t_max` + 1, a new cycle begins at `base_lr`
|
||||
#[must_use]
|
||||
pub fn cycle_step(&self, step: usize) -> usize {
|
||||
if self.use_restarts {
|
||||
// Use t_max + 1 as period so step t_max is end of cycle (min_lr)
|
||||
// and step t_max + 1 starts a new cycle (base_lr)
|
||||
step % (self.t_max + 1)
|
||||
} else {
|
||||
step.min(self.t_max)
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the progress within the current cycle (0.0 to 1.0)
|
||||
#[must_use]
|
||||
pub fn cycle_progress(&self, step: usize) -> f64 {
|
||||
let cycle_step = self.cycle_step(step);
|
||||
cycle_step as f64 / self.t_max as f64
|
||||
}
|
||||
}
|
||||
|
||||
impl LearningRateScheduler for CosineAnnealingScheduler {
|
||||
fn get_lr(&self, _epoch: usize, step: usize) -> f64 {
|
||||
let cycle_step = self.cycle_step(step);
|
||||
let progress = cycle_step as f64 / self.t_max as f64;
|
||||
|
||||
// Cosine annealing formula
|
||||
let cosine_factor = 0.5 * (1.0 + (std::f64::consts::PI * progress).cos());
|
||||
let lr = self.min_lr + (self.base_lr - self.min_lr) * cosine_factor;
|
||||
|
||||
trace!(
|
||||
"Cosine step {}: cycle_step={}, progress={:.4}, lr={:.6}",
|
||||
step, cycle_step, progress, lr
|
||||
);
|
||||
|
||||
lr
|
||||
}
|
||||
|
||||
fn step(&mut self) {
|
||||
self.current_step += 1;
|
||||
trace!(
|
||||
"Cosine annealing scheduler stepped to: {}",
|
||||
self.current_step
|
||||
);
|
||||
}
|
||||
|
||||
fn current_step(&self) -> usize {
|
||||
self.current_step
|
||||
}
|
||||
|
||||
fn reset(&mut self) {
|
||||
self.current_step = 0;
|
||||
debug!("Reset cosine annealing scheduler");
|
||||
}
|
||||
|
||||
fn scheduler_type(&self) -> &'static str {
|
||||
if self.use_restarts {
|
||||
"CosineAnnealingWarmRestarts"
|
||||
} else {
|
||||
"CosineAnnealing"
|
||||
}
|
||||
}
|
||||
|
||||
fn base_lr(&self) -> f64 {
|
||||
self.base_lr
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_cosine_annealing_creation() {
|
||||
let scheduler = CosineAnnealingScheduler::new(0.001, 0.00001, 1000).unwrap();
|
||||
assert_eq!(scheduler.base_lr(), 0.001);
|
||||
assert_eq!(scheduler.min_lr(), 0.00001);
|
||||
assert_eq!(scheduler.t_max(), 1000);
|
||||
assert_eq!(scheduler.current_step(), 0);
|
||||
assert!(scheduler.use_restarts());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cosine_annealing_invalid_params() {
|
||||
assert!(CosineAnnealingScheduler::new(-0.001, 0.00001, 1000).is_err());
|
||||
assert!(CosineAnnealingScheduler::new(0.001, -0.00001, 1000).is_err());
|
||||
assert!(CosineAnnealingScheduler::new(0.001, 0.002, 1000).is_err());
|
||||
assert!(CosineAnnealingScheduler::new(0.001, 0.00001, 0).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cosine_annealing_progression() {
|
||||
let scheduler = CosineAnnealingScheduler::new(0.001, 0.00001, 1000).unwrap();
|
||||
|
||||
// At step 0, should be at base_lr
|
||||
let lr_0 = scheduler.get_lr(0, 0);
|
||||
assert!((lr_0 - 0.001).abs() < 1e-10);
|
||||
|
||||
// At step t_max, should be at min_lr
|
||||
let lr_1000 = scheduler.get_lr(0, 1000);
|
||||
assert!((lr_1000 - 0.00001).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cosine_annealing_restarts() {
|
||||
let scheduler = CosineAnnealingScheduler::new(0.001, 0.00001, 100).unwrap();
|
||||
|
||||
// At end of first cycle (step t_max = 100)
|
||||
let lr_100 = scheduler.get_lr(0, 100);
|
||||
assert!((lr_100 - 0.00001).abs() < 1e-6);
|
||||
|
||||
// At start of second cycle (step t_max + 1 = 101)
|
||||
let lr_101 = scheduler.get_lr(0, 101);
|
||||
assert!((lr_101 - 0.001).abs() < 1e-6);
|
||||
|
||||
// Cycle step verification:
|
||||
// - step 100 is at end of cycle 1 (cycle_step = 100)
|
||||
// - step 101 is at start of cycle 2 (cycle_step = 0)
|
||||
assert_eq!(scheduler.cycle_step(100), 100);
|
||||
assert_eq!(scheduler.cycle_step(101), 0);
|
||||
assert_eq!(scheduler.cycle_step(102), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cosine_annealing_without_restarts() {
|
||||
let scheduler =
|
||||
CosineAnnealingScheduler::with_restarts(0.001, 0.00001, 100, false).unwrap();
|
||||
assert!(!scheduler.use_restarts());
|
||||
|
||||
// After t_max, should stay at min_lr
|
||||
let lr_100 = scheduler.get_lr(0, 100);
|
||||
let lr_150 = scheduler.get_lr(0, 150);
|
||||
|
||||
assert!((lr_100 - 0.00001).abs() < 1e-6);
|
||||
assert!((lr_150 - 0.00001).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cycle_calculations() {
|
||||
let scheduler = CosineAnnealingScheduler::new(0.001, 0.00001, 100).unwrap();
|
||||
|
||||
// Each cycle has t_max + 1 = 101 steps (0 to 100 inclusive)
|
||||
// Cycle 0: steps 0-100
|
||||
// Cycle 1: steps 101-201
|
||||
// Cycle 2: steps 202-302
|
||||
assert_eq!(scheduler.current_cycle(0), 0);
|
||||
assert_eq!(scheduler.current_cycle(100), 0); // Still in cycle 0 (last step)
|
||||
assert_eq!(scheduler.current_cycle(101), 1); // Start of cycle 1
|
||||
assert_eq!(scheduler.current_cycle(201), 1); // Last step of cycle 1
|
||||
assert_eq!(scheduler.current_cycle(202), 2); // Start of cycle 2
|
||||
|
||||
assert_eq!(scheduler.cycle_step(0), 0);
|
||||
assert_eq!(scheduler.cycle_step(99), 99);
|
||||
assert_eq!(scheduler.cycle_step(100), 100); // End of cycle 0
|
||||
assert_eq!(scheduler.cycle_step(101), 0); // Start of cycle 1
|
||||
assert_eq!(scheduler.cycle_step(150), 49); // 150 % 101 = 49
|
||||
|
||||
assert_eq!(scheduler.cycle_progress(0), 0.0);
|
||||
assert_eq!(scheduler.cycle_progress(50), 0.5);
|
||||
assert_eq!(scheduler.cycle_progress(100), 1.0); // End of cycle
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user