[TODO] feature: CosineAnnealingLR

[L-M-Sherlock]

PyTorch API:

https://github.com/open-spaced-repetition/fsrs-optimizer/blob/95694b787bb71ac9883db1201af09e334ee5ee0b/src/fsrs_optimizer/fsrs_optimizer.py#L195

Python implementation:

def cosine_annealing_lr(lr, step_count, T_max, eta_min = 0):
    lr = eta_min + (lr - eta_min) * (1 + math.cos(math.pi * step_count / T_max)) / (1 + math.cos(math.pi * (step_count - 1) / T_max))
    return lr

https://github.com/open-spaced-repetition/fsrs-optimizer-tiny/blob/e4256480d958ea7a54a0985a33cf81f545cfa075/seq2one.py#L137-L139

Reference: https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.CosineAnnealingLR.html

[asukaminato0721]

import math


def cosine_annealing_lr(lr, step_count, T_max, eta_min=0.0):
    lr = eta_min + (lr - eta_min) * (1 + math.cos(math.pi * step_count / T_max)) / (
        1 + math.cos(math.pi * (step_count - 1) / T_max)
    )
    return lr


lr = 0.1
step_count = 10.0
t_max = 100.0
eta_min = 0.0
assert math.isclose(
    0.0995287864543317, cosine_annealing_lr(lr, step_count, t_max, eta_min)
)

use std::f64::consts::PI;

fn cosine_annealing_lr(lr: f64, step_count: f64, t_max: f64, eta_min: f64) -> f64 {
    let cosine_arg = PI * step_count / t_max;
    let lr = eta_min + (lr - eta_min) * (1.0 + f64::cos(cosine_arg)) / (1.0 + f64::cos(PI * (step_count - 1.0) / t_max));
    lr
}

fn main() {
    let lr = 0.1;
    let step_count = 10.0;
    let t_max = 100.0;
    let eta_min = 0.0;

    let new_lr = cosine_annealing_lr(lr, step_count, t_max, eta_min);
    assert_eq!(0.09952878645433175, new_lr);
}

ref https://chat.openai.com/share/1adcedee-7c58-4583-88dc-f504b80fce6b

[L-M-Sherlock]

The next step is to implement the interface (or trait?):

https://github.com/burn-rs/burn/blob/d18d1b0bb9cf641e83f356296e7614845dda64f9/burn-core/src/lr_scheduler/base.rs#L4-L17

https://github.com/burn-rs/burn/blob/d18d1b0bb9cf641e83f356296e7614845dda64f9/burn-core/src/lr_scheduler/noam.rs

[asukaminato0721]

I am not so sure about the correctness of implement of this.

use burn::{lr_scheduler::LRScheduler, LearningRate};
#[derive(Clone, Debug)]
pub struct CosineAnnealingLR {
    t_max: f64,
    eta_min: f64,
    init_lr: LearningRate,
    step_count: f64,
}

impl LRScheduler for CosineAnnealingLR {
    type Record = usize;

    fn step(&mut self) -> LearningRate {
        self.step_count += 1.0;
        use std::f64::consts::PI;
        fn cosine_annealing_lr(lr: f64, step_count: f64, t_max: f64, eta_min: f64) -> f64 {
            let cosine_arg = PI * step_count / t_max;
            let lr = eta_min
                + (lr - eta_min) * (1.0 + f64::cos(cosine_arg))
                    / (1.0 + f64::cos(PI * (step_count - 1.0) / t_max));
            lr
        }
        self.init_lr = cosine_annealing_lr(self.init_lr, self.step_count, self.t_max, self.eta_min);
        self.init_lr
    }

    fn to_record(&self) -> Self::Record {
        self.step_count as usize
    }

    fn load_record(mut self, record: Self::Record) -> Self {
        self.step_count = record as f64;
        self
    }
}