cwenzi/neuroflow-cpp

1#include "neuroflow/scheduler.hpp"
2
3#include <algorithm>
4#include <cmath>
5
6namespace neuroflow {
7
8CosineScheduler::CosineScheduler(float lr_max, size_t total_steps,
9                                  float lr_min_ratio, float warmup_ratio)
10    : lr_max_(lr_max),
11      lr_min_(lr_max * lr_min_ratio),
12      warmup_steps_(std::max(size_t(1), static_cast<size_t>(static_cast<float>(total_steps) * warmup_ratio))),
13      total_steps_(total_steps) {}
14
15float CosineScheduler::get_lr(size_t step) const {
16    if (total_steps_ == 0) return lr_max_;
17    if (step >= total_steps_) return lr_min_;
18    if (step < warmup_steps_) {
19        return lr_max_ * static_cast<float>(step + 1) / static_cast<float>(warmup_steps_);
20    }
21    float progress = static_cast<float>(step - warmup_steps_)
22                   / static_cast<float>(total_steps_ - warmup_steps_);
23    return lr_min_ + 0.5f * (lr_max_ - lr_min_)
24                   * (1.0f + std::cos(3.14159265358979f * progress));
25}
26
27CosineScheduler::Phase CosineScheduler::get_phase(size_t step) const {
28    if (step < warmup_steps_) return Phase::WARMUP;
29    if (step < total_steps_) return Phase::COSINE;
30    return Phase::DONE;
31}
32
33} // namespace neuroflow