osobhandClaude Opus 4.6
02d382d5f6
Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
305 lines
9.6 KiB
Rust
305 lines
9.6 KiB
Rust
//! 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
|
|
}
|
|
}
|